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.Maps;
020    import kotlin.jvm.functions.Function3;
021    import org.jetbrains.annotations.NotNull;
022    import org.jetbrains.kotlin.utils.Printer;
023    
024    import java.util.Collection;
025    import java.util.Map;
026    
027    public class TrackingSlicedMap extends SlicedMapImpl {
028        private final Map<ReadOnlySlice<?, ?>, SliceWithStackTrace<?, ?>> sliceTranslationMap = Maps.newHashMap();
029        private final boolean trackWithStackTraces;
030    
031        public TrackingSlicedMap(boolean trackWithStackTraces) {
032            this.trackWithStackTraces = trackWithStackTraces;
033        }
034    
035        private <K, V> SliceWithStackTrace<K, V> wrapSlice(ReadOnlySlice<K, V> slice) {
036            SliceWithStackTrace<?, ?> translated = sliceTranslationMap.get(slice);
037            if (translated == null) {
038                translated = new SliceWithStackTrace<K, V>(slice);
039                sliceTranslationMap.put(slice, translated);
040            }
041            //noinspection unchecked
042            return (SliceWithStackTrace) translated;
043        }
044    
045        @Override
046        public <K, V> V get(ReadOnlySlice<K, V> slice, K key) {
047            return super.get(wrapSlice(slice), key).value;
048        }
049    
050        @Override
051        public <K, V> Collection<K> getKeys(WritableSlice<K, V> slice) {
052            return super.getKeys(wrapSlice(slice));
053        }
054    
055        @Override
056        public void forEach(@NotNull final Function3<WritableSlice, Object, Object, Void> f) {
057            super.forEach(new Function3<WritableSlice, Object, Object, Void>() {
058                @Override
059                public Void invoke(WritableSlice slice, Object key, Object value) {
060                    f.invoke(((SliceWithStackTrace) slice).getWritableDelegate(), key, ((TrackableValue<?>) value).value);
061                    return null;
062                }
063            });
064        }
065    
066        @Override
067        public <K, V> void put(WritableSlice<K, V> slice, K key, V value) {
068            super.put(wrapSlice(slice), key, new TrackableValue<V>(value, trackWithStackTraces));
069        }
070    
071        private static class TrackableValue<V> {
072            private final static StackTraceElement[] EMPTY_STACK_TRACE = new StackTraceElement[0];
073    
074            private final V value;
075            private final StackTraceElement[] stackTrace;
076            private final String threadName;
077    
078            private TrackableValue(V value, boolean storeStack) {
079                this.value = value;
080                this.stackTrace = storeStack ? Thread.currentThread().getStackTrace() : EMPTY_STACK_TRACE;
081                this.threadName = Thread.currentThread().getName();
082            }
083    
084            private Appendable printStackTrace(Appendable appendable) {
085                Printer s = new Printer(appendable);
086                s.println(value);
087                s.println("Thread: " + threadName);
088                s.println("Written at ");
089                StackTraceElement[] trace = stackTrace;
090                for (StackTraceElement aTrace : trace) {
091                    s.println("\tat " + aTrace);
092                }
093                s.println("---------");
094                return appendable;
095            }
096    
097            @Override
098            public String toString() {
099                return printStackTrace(new StringBuilder()).toString();
100            }
101    
102            @Override
103            public boolean equals(Object o) {
104                if (this == o) return true;
105                if (o == null || getClass() != o.getClass()) return false;
106    
107                TrackableValue other = (TrackableValue) o;
108    
109                if (value != null ? !value.equals(other.value) : other.value != null) return false;
110    
111                return true;
112            }
113    
114            @Override
115            public int hashCode() {
116                return value != null ? value.hashCode() : 0;
117            }
118        }
119    
120        private class SliceWithStackTrace<K, V>
121                extends AbstractWritableSlice<K, TrackableValue<V>>
122                implements WritableSlice<K, TrackableValue<V>> {
123    
124            private final ReadOnlySlice<K, V> delegate;
125    
126            private SliceWithStackTrace(@NotNull ReadOnlySlice<K, V> delegate) {
127                super(delegate.toString());
128                this.delegate = delegate;
129            }
130    
131            // Methods of ReadOnlySlice
132    
133            @Override
134            public TrackableValue<V> computeValue(SlicedMap map, K key, TrackableValue<V> value, boolean valueNotFound) {
135                return new TrackableValue<V>(delegate.computeValue(map, key, value == null ? null : value.value, valueNotFound), trackWithStackTraces);
136            }
137    
138            @Override
139            public ReadOnlySlice<K, TrackableValue<V>> makeRawValueVersion() {
140                return wrapSlice(delegate.makeRawValueVersion());
141            }
142    
143            // Methods of WritableSlice
144    
145            private WritableSlice<K, V> getWritableDelegate() {
146                return (WritableSlice<K, V>) delegate;
147            }
148    
149            @Override
150            public boolean isCollective() {
151                return getWritableDelegate().isCollective();
152            }
153    
154            @Override
155            public RewritePolicy getRewritePolicy() {
156                return getWritableDelegate().getRewritePolicy();
157            }
158    
159            @Override
160            public void afterPut(MutableSlicedMap map, K key, TrackableValue<V> value) {
161                getWritableDelegate().afterPut(map, key, value.value);
162            }
163    
164            @Override
165            public boolean check(K key, TrackableValue<V> value) {
166                return getWritableDelegate().check(key, value.value);
167            }
168        }
169    }