Introduce KeyWithSlice and use it as Key for UserDataHolder in SlicedMap instead of cashing mapping from key to slice(see SlicedMapImpl#keyIdToSlice)

compiler/frontend/src/org/jetbrains/kotlin/util/slicedMap/AbstractWritableSlice.kt

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2010-2015 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.kotlin.util.slicedMap
+
+abstract class AbstractWritableSlice<K, V>(debugName: String) : KeyWithSlice<K, V, WritableSlice<K, V>>(debugName), WritableSlice<K, V> {
+    override val slice: WritableSlice<K, V>
+        get() = this
+
+    override fun getKey(): AbstractWritableSlice<K, V> = this
+}

compiler/frontend/src/org/jetbrains/kotlin/util/slicedMap/BasicWritableSlice.java

@@ -16,13 +16,12 @@
 
 package org.jetbrains.kotlin.util.slicedMap;
 
-import com.intellij.openapi.util.Key;
 import org.jetbrains.annotations.NotNull;
 
 import java.lang.reflect.Field;
 import java.lang.reflect.Modifier;
 
-public class BasicWritableSlice<K, V> extends Key<V> implements WritableSlice<K, V> {
+public class BasicWritableSlice<K, V> extends AbstractWritableSlice<K, V> {
 
     public static Void initSliceDebugNames(Class<?> declarationOwner) {
         for (Field field : declarationOwner.getFields()) {
@@ -55,12 +54,6 @@ public class BasicWritableSlice<K, V> extends Key<V> implements WritableSlice<K,
         this.isCollective = isCollective;
     }
 
-    @Override
-    @NotNull
-    public Key<V> getKey() {
-        return this;
-    }
-
     // True to put, false to skip
     @Override
     public boolean check(K key, V value) {

compiler/frontend/src/org/jetbrains/kotlin/util/slicedMap/DelegatingSlice.java

@@ -16,7 +16,6 @@
 
 package org.jetbrains.kotlin.util.slicedMap;
 
-import com.intellij.openapi.util.Key;
 import org.jetbrains.annotations.NotNull;
 
 /**
@@ -51,7 +50,7 @@ public class DelegatingSlice<K, V> implements WritableSlice<K, V> {
 
     @Override
     @NotNull
-    public Key<V> getKey() {
+    public KeyWithSlice<K, V, WritableSlice<K, V>> getKey() {
         return delegate.getKey();
     }
 

compiler/frontend/src/org/jetbrains/kotlin/util/slicedMap/KeyWithSlice.kt

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2010-2015 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.kotlin.util.slicedMap
+
+import com.intellij.openapi.util.Key
+
+abstract class KeyWithSlice<K, V, out Slice : ReadOnlySlice<K, V>>(debugName: String) : Key<V>(debugName) {
+    abstract val slice: Slice
+}

compiler/frontend/src/org/jetbrains/kotlin/util/slicedMap/ReadOnlySlice.java

@@ -16,12 +16,11 @@
 
 package org.jetbrains.kotlin.util.slicedMap;
 
-import com.intellij.openapi.util.Key;
 import org.jetbrains.annotations.NotNull;
 
 public interface ReadOnlySlice<K, V> {
     @NotNull
-    Key<V> getKey();
+    KeyWithSlice<K, V, ? extends ReadOnlySlice<K, V>> getKey();
 
     V computeValue(SlicedMap map, K key, V value, boolean valueNotFound);
 

compiler/frontend/src/org/jetbrains/kotlin/util/slicedMap/SlicedMapImpl.java

@@ -22,7 +22,6 @@ import com.google.common.collect.Multimap;
 import com.intellij.openapi.util.Key;
 import com.intellij.openapi.util.UserDataHolder;
 import gnu.trove.THashMap;
-import gnu.trove.TIntObjectHashMap;
 import kotlin.Function3;
 import org.jetbrains.annotations.NotNull;
 
@@ -31,8 +30,6 @@ import java.util.Map;
 
 public class SlicedMapImpl implements MutableSlicedMap {
 
-    private static final TIntObjectHashMap<WritableSlice> keyIdToSlice = new TIntObjectHashMap<WritableSlice>();
-
     public static SlicedMapImpl create() {
         return new SlicedMapImpl();
     }
@@ -69,11 +66,6 @@ public class SlicedMapImpl implements MutableSlicedMap {
             collectiveSliceKeys.put(slice, key);
         }
 
-        int sliceKeyId = sliceKey.hashCode();
-        if (!keyIdToSlice.contains(sliceKeyId)) {
-            keyIdToSlice.put(sliceKeyId, slice);
-        }
-
         holder.putUserData(sliceKey, value);
         slice.afterPut(this, key, value);
     }
@@ -126,10 +118,9 @@ public class SlicedMapImpl implements MutableSlicedMap {
             if (holder == null) continue;
 
             for (Key<?> sliceKey : holder.getKeys()) {
-                WritableSlice slice = keyIdToSlice.get(sliceKey.hashCode());
                 Object value = holder.getUserData(sliceKey);
 
-                f.invoke(slice, key, value);
+                f.invoke(((AbstractWritableSlice) sliceKey).getSlice(), key, value);
             }
         }
     }

compiler/frontend/src/org/jetbrains/kotlin/util/slicedMap/TrackingSlicedMap.java

@@ -18,7 +18,6 @@ package org.jetbrains.kotlin.util.slicedMap;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
-import com.intellij.openapi.util.Key;
 import kotlin.Function3;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.TestOnly;
@@ -137,7 +136,7 @@ public class TrackingSlicedMap extends SlicedMapImpl {
         }
     }
 
-    private class SliceWithStackTrace<K, V> extends Key<TrackableValue<V>> implements RemovableSlice<K, TrackableValue<V>> {
+    private class SliceWithStackTrace<K, V> extends AbstractWritableSlice<K, TrackableValue<V>> implements RemovableSlice<K, TrackableValue<V>> {
 
         private final ReadOnlySlice<K, V> delegate;
 
@@ -148,12 +147,6 @@ public class TrackingSlicedMap extends SlicedMapImpl {
 
         // Methods of ReadOnlySlice
 
-        @Override
-        @NotNull
-        public Key<TrackableValue<V>> getKey() {
-            return this;
-        }
-
         @Override
         public TrackableValue<V> computeValue(SlicedMap map, K key, TrackableValue<V> value, boolean valueNotFound) {
             return new TrackableValue<V>(delegate.computeValue(map, key, value == null ? null : value.value, valueNotFound), trackWithStackTraces);

compiler/frontend/src/org/jetbrains/kotlin/util/slicedMap/WritableSlice.java

@@ -16,7 +16,13 @@
 
 package org.jetbrains.kotlin.util.slicedMap;
 
+import org.jetbrains.annotations.NotNull;
+
 public interface WritableSlice<K, V> extends ReadOnlySlice<K, V> {
+    @NotNull
+    @Override
+    KeyWithSlice<K, V, WritableSlice<K, V>> getKey();
+
     // True to put, false to skip
     boolean check(K key, V value);