1/*
2 * Copyright (C) 2017 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17package com.android.systemui.util.leak;
18
19import java.lang.ref.Reference;
20import java.lang.ref.ReferenceQueue;
21import java.lang.ref.WeakReference;
22import java.util.Collection;
23import java.util.HashMap;
24import java.util.Map;
25import java.util.Set;
26
27/**
28 * Like WeakHashMap, but uses identity instead of equality when comparing keys.
29 */
30public class WeakIdentityHashMap<K,V> {
31
32    private final HashMap<WeakReference<K>,V> mMap = new HashMap<>();
33    private final ReferenceQueue<Object> mRefQueue = new ReferenceQueue<>();
34
35    private void cleanUp() {
36        Reference<?> ref;
37        while ((ref = mRefQueue.poll()) != null) {
38            mMap.remove(ref);
39        }
40    }
41
42    public void put(K key, V value) {
43        cleanUp();
44        mMap.put(new CmpWeakReference<>(key, mRefQueue), value);
45    }
46
47    public V get(K key) {
48        cleanUp();
49        return mMap.get(new CmpWeakReference<>(key));
50    }
51
52    public Collection<V> values() {
53        cleanUp();
54        return mMap.values();
55    }
56
57    public Set<Map.Entry<WeakReference<K>, V>> entrySet() {
58        return mMap.entrySet();
59    }
60
61    public int size() {
62        cleanUp();
63        return mMap.size();
64    }
65
66    public boolean isEmpty() {
67        cleanUp();
68        return mMap.isEmpty();
69    }
70
71    private static class CmpWeakReference<K> extends WeakReference<K> {
72        private final int mHashCode;
73
74        public CmpWeakReference(K key) {
75            super(key);
76            mHashCode = System.identityHashCode(key);
77        }
78
79        public CmpWeakReference(K key, ReferenceQueue<Object> refQueue) {
80            super(key, refQueue);
81            mHashCode = System.identityHashCode(key);
82        }
83
84        @Override
85        public boolean equals(Object o) {
86            if (o == this) {
87                return true;
88            }
89            K k = get();
90            if (k != null && o instanceof CmpWeakReference) {
91                return ((CmpWeakReference) o).get() == k;
92            }
93            return false;
94        }
95
96        @Override
97        public int hashCode() {
98            return mHashCode;
99        }
100    }
101}
102