Fix tracking slices

compiler/frontend/src/org/jetbrains/jet/lang/resolve/BindingTraceContext.java

@@ -30,11 +30,9 @@ public class BindingTraceContext implements BindingTrace {
     private final List<Diagnostic> diagnostics = Lists.newArrayList();
 
     // This flag is used for debugging of "Rewrite at slice..." exceptions
-    // NOTE: sometimes TrackingSlicedMap throws a ClassCastException (after you have fixed the rewrite).
-    // I gave up debugging it, because it still serves its purpose. Any suggestions on how to fix it are welcome. (abreslav)
     private final static boolean TRACK_REWRITES = false;
-    @SuppressWarnings("ConstantConditions")
-    private final MutableSlicedMap map = TRACK_REWRITES ? new TrackingSlicedMap(SlicedMapImpl.create()) : SlicedMapImpl.create();
+
+    private final MutableSlicedMap map;
 
     private final BindingContext bindingContext = new BindingContext() {
 
@@ -62,6 +60,21 @@ public class BindingTraceContext implements BindingTrace {
         }
     };
 
+    public BindingTraceContext() {
+        //noinspection ConstantConditions
+        this(TRACK_REWRITES ? new TrackingSlicedMap() : SlicedMapImpl.create());
+    }
+
+
+    private BindingTraceContext(MutableSlicedMap map) {
+        this.map = map;
+    }
+
+    @TestOnly
+    public static BindingTraceContext createTraceableBindingTrace() {
+        return new BindingTraceContext(new TrackingSlicedMap());
+    }
+
     @Override
     public void report(@NotNull Diagnostic diagnostic) {
         diagnostics.add(diagnostic);

compiler/frontend/src/org/jetbrains/jet/util/slicedmap/SlicedMapImpl.java

@@ -45,7 +45,7 @@ public class SlicedMapImpl implements MutableSlicedMap {
     private final Map<SlicedMapKey<?, ?>, Object> map;
     private final Multimap<WritableSlice<?, ?>, Object> collectiveSliceKeys = Multimaps.newListMultimap(new HashMap<WritableSlice<?, ?>, Collection<Object>>(), CommonSuppliers.getArrayListSupplier());
 
-    private SlicedMapImpl(Map<SlicedMapKey<?, ?>, Object> map) {
+    protected SlicedMapImpl(Map<SlicedMapKey<?, ?>, Object> map) {
         this.map = map;
     }
 
@@ -100,6 +100,7 @@ public class SlicedMapImpl implements MutableSlicedMap {
         return (V) map.remove(slice.makeKey(key));
     }
 
+    @NotNull
     @Override
     public Iterator<Map.Entry<SlicedMapKey<?, ?>, ?>> iterator() {
         //noinspection unchecked

compiler/frontend/src/org/jetbrains/jet/util/slicedmap/Slices.java

@@ -141,7 +141,9 @@ public class Slices {
                         if (valueNotFound) {
                             for (ReadOnlySlice<K, V> slice : furtherLookupSlices) {
                                 V v = map.get(slice, key);
-                                if (v != null) return v;
+                                if (v != null) {
+                                    return v;
+                                }
                             }
                             return defaultValue;
                         }

compiler/frontend/src/org/jetbrains/jet/util/slicedmap/TrackingSlicedMap.java

@@ -26,13 +26,13 @@ import java.util.Collection;
 import java.util.Iterator;
 import java.util.Map;
 
-public class TrackingSlicedMap implements MutableSlicedMap {
+public class TrackingSlicedMap extends SlicedMapImpl {
 
-    private final MutableSlicedMap delegate;
+    // private final MutableSlicedMap delegate;
     private final Map<ReadOnlySlice<?, ?>, SliceWithStackTrace<?, ?>> sliceTranslationMap = Maps.newHashMap();
 
-    public TrackingSlicedMap(@NotNull MutableSlicedMap delegate) {
-        this.delegate = delegate;
+    public TrackingSlicedMap() {
+        super(Maps.<SlicedMapKey<?, ?>, Object>newLinkedHashMap());
     }
 
     private <K, V> SliceWithStackTrace<K, V> wrapSlice(ReadOnlySlice<K, V> slice) {
@@ -47,45 +47,51 @@ public class TrackingSlicedMap implements MutableSlicedMap {
 
     @Override
     public <K, V> V get(ReadOnlySlice<K, V> slice, K key) {
-        return delegate.get(wrapSlice(slice), key).value;
+        return super.get(wrapSlice(slice), key).value;
     }
 
     @Override
     public <K, V> Collection<K> getKeys(WritableSlice<K, V> slice) {
-        return delegate.getKeys(wrapSlice(slice));
+        return super.getKeys(wrapSlice(slice));
     }
 
     @NotNull
     @Override
     public Iterator<Map.Entry<SlicedMapKey<?, ?>, ?>> iterator() {
         Map<SlicedMapKey<?, ?>, Object> map = Maps.newHashMap();
-        for (Map.Entry<SlicedMapKey<?, ?>, ?> entry : delegate) {
+
+        Iterator<Map.Entry<SlicedMapKey<?, ?>, ?>> iterator = super.iterator();
+
+        //noinspection WhileLoopReplaceableByForEach
+        while (iterator.hasNext()) {
+            Map.Entry<SlicedMapKey<?, ?>, ?> entry = iterator.next();
             map.put(entry.getKey(), ((WithStackTrace<?>) entry.getValue()).value);
         }
+
         //noinspection unchecked
         return (Iterator) map.entrySet().iterator();
     }
 
     @Override
     public <K, V> void put(WritableSlice<K, V> slice, K key, V value) {
-        delegate.put(wrapSlice(slice), key, new WithStackTrace<V>(value));
+        super.put(wrapSlice(slice), key, new WithStackTrace<V>(value));
     }
 
     @Override
     public <K, V> V remove(RemovableSlice<K, V> slice, K key) {
-        return delegate.remove(wrapSlice(slice), key).value;
+        return super.remove(wrapSlice(slice), key).value;
     }
 
     @Override
     public void clear() {
-        delegate.clear();
+        super.clear();
     }
 
     @Override
     @NotNull
     @TestOnly
     public <K, V> ImmutableMap<K, V> getSliceContents(@NotNull ReadOnlySlice<K, V> slice) {
-        return delegate.getSliceContents(slice);
+        return super.getSliceContents(slice);
     }
 
     private static class WithStackTrace<V> {

compiler/tests/org/jetbrains/jet/util/slicemap/TrackingSliceMapTest.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2010-2013 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.jet.util.slicemap;
+
+import junit.framework.TestCase;
+import org.jetbrains.jet.lang.resolve.BindingTraceContext;
+import org.jetbrains.jet.util.slicedmap.ReadOnlySlice;
+import org.jetbrains.jet.util.slicedmap.Slices;
+import org.jetbrains.jet.util.slicedmap.WritableSlice;
+
+public class TrackingSliceMapTest extends TestCase {
+    public void testSimpleSlice() {
+        WritableSlice<String, Integer> SUPER_COMPUTER = Slices.<String, Integer>sliceBuilder().setDebugName("SUPER_COMPUTER").build();
+        BindingTraceContext traceContext = BindingTraceContext.createTraceableBindingTrace();
+
+        traceContext.record(SUPER_COMPUTER, "Answer", 42);
+        Integer answer = traceContext.get(SUPER_COMPUTER, "Answer");
+
+        assertNotNull(answer);
+        assertEquals(42, (int) answer);
+    }
+
+    public void testOppositeSlice() {
+        WritableSlice<Integer, String> COLOR_NAME = Slices.<Integer, String>sliceBuilder().setDebugName("COLOR_NAME").build();
+        WritableSlice<String, Integer> NAME_COLOR = Slices.<String, Integer>sliceBuilder().setOpposite(COLOR_NAME).setDebugName("NAME_COLOR").build();
+
+        BindingTraceContext traceContext = BindingTraceContext.createTraceableBindingTrace();
+
+        traceContext.record(NAME_COLOR, "RED", 0xff0000);
+        String redName = traceContext.get(COLOR_NAME, 0xff0000);
+
+        assertEquals("RED", redName);
+    }
+
+    public void testFurtherSlices() {
+        WritableSlice<String, Integer> NAME_COLOR = Slices.<String, Integer>sliceBuilder().setDebugName("NAME_COLOR").build();
+
+        @SuppressWarnings("unchecked")
+        WritableSlice<String, Object> NAME_OBJECT = Slices.<String, Object>sliceBuilder()
+                .setFurtherLookupSlices(new ReadOnlySlice[] {NAME_COLOR})
+                .setDebugName("NAME_OBJECT").build();
+
+        BindingTraceContext traceContext = BindingTraceContext.createTraceableBindingTrace();
+
+        traceContext.record(NAME_COLOR, "RED", 0xff0000);
+        Object object = traceContext.get(NAME_OBJECT, "RED");
+
+        assertNotNull(object);
+
+        Integer color = (Integer) object;
+        assertEquals(0xff0000, (int) color);
+    }
+}