001    /*
002     * Copyright 2010-2015 JetBrains s.r.o.
003     *
004     * Licensed under the Apache License, Version 2.0 (the "License");
005     * you may not use this file except in compliance with the License.
006     * You may obtain a copy of the License at
007     *
008     * http://www.apache.org/licenses/LICENSE-2.0
009     *
010     * Unless required by applicable law or agreed to in writing, software
011     * distributed under the License is distributed on an "AS IS" BASIS,
012     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013     * See the License for the specific language governing permissions and
014     * limitations under the License.
015     */
016    
017    package org.jetbrains.kotlin.util.slicedMap;
018    
019    import com.google.common.collect.ImmutableMap;
020    import com.google.common.collect.Maps;
021    import org.jetbrains.annotations.NotNull;
022    import org.jetbrains.annotations.TestOnly;
023    import org.jetbrains.kotlin.utils.Printer;
024    
025    import java.util.Collection;
026    import java.util.Iterator;
027    import java.util.Map;
028    
029    public class TrackingSlicedMap extends SlicedMapImpl {
030        private final Map<ReadOnlySlice<?, ?>, SliceWithStackTrace<?, ?>> sliceTranslationMap = Maps.newHashMap();
031        private final boolean trackWithStackTraces;
032    
033        public TrackingSlicedMap(boolean trackWithStackTraces) {
034            super(Maps.<SlicedMapKey<?, ?>, Object>newLinkedHashMap());
035            this.trackWithStackTraces = trackWithStackTraces;
036        }
037    
038        private <K, V> SliceWithStackTrace<K, V> wrapSlice(ReadOnlySlice<K, V> slice) {
039            SliceWithStackTrace<?, ?> translated = sliceTranslationMap.get(slice);
040            if (translated == null) {
041                translated = new SliceWithStackTrace<K, V>(slice);
042                sliceTranslationMap.put(slice, translated);
043            }
044            //noinspection unchecked
045            return (SliceWithStackTrace) translated;
046        }
047    
048        @Override
049        public <K, V> V get(ReadOnlySlice<K, V> slice, K key) {
050            return super.get(wrapSlice(slice), key).value;
051        }
052    
053        @Override
054        public <K, V> Collection<K> getKeys(WritableSlice<K, V> slice) {
055            return super.getKeys(wrapSlice(slice));
056        }
057    
058        @NotNull
059        @Override
060        public Iterator<Map.Entry<SlicedMapKey<?, ?>, ?>> iterator() {
061            Map<SlicedMapKey<?, ?>, Object> map = Maps.newHashMap();
062    
063            Iterator<Map.Entry<SlicedMapKey<?, ?>, ?>> iterator = super.iterator();
064    
065            //noinspection WhileLoopReplaceableByForEach
066            while (iterator.hasNext()) {
067                Map.Entry<SlicedMapKey<?, ?>, ?> entry = iterator.next();
068                map.put(entry.getKey(), ((TrackableValue<?>) entry.getValue()).value);
069            }
070    
071            //noinspection unchecked
072            return (Iterator) map.entrySet().iterator();
073        }
074    
075        @Override
076        public <K, V> void put(WritableSlice<K, V> slice, K key, V value) {
077            super.put(wrapSlice(slice), key, new TrackableValue<V>(value, trackWithStackTraces));
078        }
079    
080        @Override
081        public <K, V> V remove(RemovableSlice<K, V> slice, K key) {
082            return super.remove(wrapSlice(slice), key).value;
083        }
084    
085        @Override
086        public void clear() {
087            super.clear();
088        }
089    
090        @Override
091        @NotNull
092        @TestOnly
093        public <K, V> ImmutableMap<K, V> getSliceContents(@NotNull ReadOnlySlice<K, V> slice) {
094            return super.getSliceContents(slice);
095        }
096    
097        private static class TrackableValue<V> {
098            private final static StackTraceElement[] EMPTY_STACK_TRACE = new StackTraceElement[0];
099    
100            private final V value;
101            private final StackTraceElement[] stackTrace;
102            private final String threadName;
103    
104            private TrackableValue(V value, boolean storeStack) {
105                this.value = value;
106                this.stackTrace = storeStack ? Thread.currentThread().getStackTrace() : EMPTY_STACK_TRACE;
107                this.threadName = Thread.currentThread().getName();
108            }
109    
110            private Appendable printStackTrace(Appendable appendable) {
111                Printer s = new Printer(appendable);
112                s.println(value);
113                s.println("Thread: " + threadName);
114                s.println("Written at ");
115                StackTraceElement[] trace = stackTrace;
116                for (StackTraceElement aTrace : trace) {
117                    s.println("\tat " + aTrace);
118                }
119                s.println("---------");
120                return appendable;
121            }
122    
123            @Override
124            public String toString() {
125                return printStackTrace(new StringBuilder()).toString();
126            }
127    
128            @Override
129            public boolean equals(Object o) {
130                if (this == o) return true;
131                if (o == null || getClass() != o.getClass()) return false;
132    
133                TrackableValue other = (TrackableValue) o;
134    
135                if (value != null ? !value.equals(other.value) : other.value != null) return false;
136    
137                return true;
138            }
139    
140            @Override
141            public int hashCode() {
142                return value != null ? value.hashCode() : 0;
143            }
144        }
145    
146        public class SliceWithStackTrace<K, V> implements RemovableSlice<K, TrackableValue<V>> {
147    
148            private final ReadOnlySlice<K, V> delegate;
149    
150            private SliceWithStackTrace(@NotNull ReadOnlySlice<K, V> delegate) {
151                this.delegate = delegate;
152            }
153    
154            // Methods of ReadOnlySlice
155    
156            @Override
157            public SlicedMapKey<K, TrackableValue<V>> makeKey(K key) {
158                //noinspection unchecked
159                return (SlicedMapKey) delegate.makeKey(key);
160            }
161    
162            @Override
163            public TrackableValue<V> computeValue(SlicedMap map, K key, TrackableValue<V> value, boolean valueNotFound) {
164                return new TrackableValue<V>(delegate.computeValue(map, key, value == null ? null : value.value, valueNotFound), trackWithStackTraces);
165            }
166    
167            @Override
168            public ReadOnlySlice<K, TrackableValue<V>> makeRawValueVersion() {
169                return wrapSlice(delegate.makeRawValueVersion());
170            }
171    
172            // Methods of WritableSlice
173    
174            private WritableSlice<K, V> getWritableDelegate() {
175                return (WritableSlice<K, V>) delegate;
176            }
177    
178            @Override
179            public boolean isCollective() {
180                return getWritableDelegate().isCollective();
181            }
182    
183            @Override
184            public RewritePolicy getRewritePolicy() {
185                return getWritableDelegate().getRewritePolicy();
186            }
187    
188            @Override
189            public void afterPut(MutableSlicedMap map, K key, TrackableValue<V> value) {
190                getWritableDelegate().afterPut(map, key, value.value);
191            }
192    
193            @Override
194            public boolean check(K key, TrackableValue<V> value) {
195                return getWritableDelegate().check(key, value.value);
196            }
197        }
198    }