Allow slice rewrite in trace when in IDE

compiler/backend/src/org/jetbrains/kotlin/codegen/serialization/JvmSerializationBindings.java

@@ -59,7 +59,7 @@ public final class JvmSerializationBindings {
         BasicWritableSlice.initSliceDebugNames(JvmSerializationBindings.class);
     }
 
-    private final MutableSlicedMap map = SlicedMapImpl.create();
+    private final MutableSlicedMap map = new SlicedMapImpl(false);
 
     public <K, V> void put(@NotNull SerializationMappingSlice<K, V> slice, @NotNull K key, @NotNull V value) {
         map.put(slice, key, value);

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingTraceContext.java

@@ -79,12 +79,16 @@ public class BindingTraceContext implements BindingTrace {
     };
 
     public BindingTraceContext() {
-        this(BindingTraceFilter.Companion.getACCEPT_ALL());
+        this(false);
     }
 
-    public BindingTraceContext(BindingTraceFilter filter) {
+    public BindingTraceContext(boolean allowSliceRewrite) {
+        this(BindingTraceFilter.Companion.getACCEPT_ALL(), allowSliceRewrite);
+    }
+
+    public BindingTraceContext(BindingTraceFilter filter, boolean allowSliceRewrite) {
         //noinspection ConstantConditions
-        this(TRACK_REWRITES ? new TrackingSlicedMap(TRACK_WITH_STACK_TRACES) : SlicedMapImpl.create(), filter);
+        this(TRACK_REWRITES && !allowSliceRewrite ? new TrackingSlicedMap(TRACK_WITH_STACK_TRACES) : new SlicedMapImpl(allowSliceRewrite), filter);
     }
 
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/CodeAnalyzerInitializer.kt

@@ -16,10 +16,10 @@
 
 package org.jetbrains.kotlin.resolve
 
-import org.jetbrains.kotlin.resolve.lazy.KotlinCodeAnalyzer
-import com.intellij.openapi.project.Project
 import com.intellij.openapi.components.ServiceManager
+import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
+import org.jetbrains.kotlin.resolve.lazy.KotlinCodeAnalyzer
 
 interface CodeAnalyzerInitializer {
     fun initialize(trace: BindingTrace, module: ModuleDescriptor, codeAnalyzer: KotlinCodeAnalyzer)
@@ -36,5 +36,5 @@ class DummyCodeAnalyzerInitializer: CodeAnalyzerInitializer {
         // Do nothing
     }
 
-    override fun createTrace(): BindingTrace = BindingTraceContext()
+    override fun createTrace(): BindingTrace = BindingTraceContext(true)
 }

compiler/frontend/src/org/jetbrains/kotlin/resolve/DelegatingBindingTrace.kt

@@ -26,11 +26,19 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.expressions.typeInfoFactory.createTypeInfo
 import org.jetbrains.kotlin.util.slicedMap.*
 
-open class DelegatingBindingTrace(private val parentContext: BindingContext,
+open class DelegatingBindingTrace(
-                                  private val name: String,
+        private val parentContext: BindingContext,
-                                  withParentDiagnostics: Boolean = true,
+        private val name: String,
-                                  private val filter: BindingTraceFilter = BindingTraceFilter.ACCEPT_ALL) : BindingTrace {
+        withParentDiagnostics: Boolean = true,
-    private val map = if (BindingTraceContext.TRACK_REWRITES) TrackingSlicedMap(BindingTraceContext.TRACK_WITH_STACK_TRACES) else SlicedMapImpl.create()
+        private val filter: BindingTraceFilter = BindingTraceFilter.ACCEPT_ALL,
+        allowSliceRewrite: Boolean = false
+) : BindingTrace {
+    
+    private val map = if (BindingTraceContext.TRACK_REWRITES && !allowSliceRewrite)
+        TrackingSlicedMap(BindingTraceContext.TRACK_WITH_STACK_TRACES)
+    else
+        SlicedMapImpl(allowSliceRewrite)
+
     private val mutableDiagnostics: MutableDiagnosticsWithSuppression?
 
     private inner class MyBindingContext : BindingContext {
@@ -72,10 +80,12 @@ open class DelegatingBindingTrace(private val parentContext: BindingContext,
     constructor(parentContext: BindingContext,
                 debugName: String,
                 resolutionSubjectForMessage: Any?,
-                filter: BindingTraceFilter = BindingTraceFilter.ACCEPT_ALL)
+                filter: BindingTraceFilter = BindingTraceFilter.ACCEPT_ALL,
-        : this(parentContext,
+                allowSliceRewrite: Boolean = false
-               AnalyzingUtils.formDebugNameForBindingTrace(debugName, resolutionSubjectForMessage),
+    ) : this(parentContext,
-               filter = filter)
+             AnalyzingUtils.formDebugNameForBindingTrace(debugName, resolutionSubjectForMessage),
+             filter = filter,
+             allowSliceRewrite = allowSliceRewrite)
 
     override fun getBindingContext(): BindingContext = bindingContext
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/TemporaryBindingTrace.java

@@ -39,7 +39,7 @@ public class TemporaryBindingTrace extends DelegatingBindingTrace {
     protected final BindingTrace trace;
 
     protected TemporaryBindingTrace(@NotNull BindingTrace trace, String debugName, BindingTraceFilter filter) {
-        super(trace.getBindingContext(), debugName, true, filter);
+        super(trace.getBindingContext(), debugName, true, filter, false);
         this.trace = trace;
     }
 

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/CallResolver.java

@@ -587,7 +587,7 @@ public class CallResolver {
         if (CallResolverUtilKt.isInvokeCallOnVariable(call)) return;
 
         DelegatingBindingTrace deltasTraceToCacheResolve = new DelegatingBindingTrace(
-                BindingContext.EMPTY, "delta trace for caching resolve of", context.call, BindingTraceFilter.Companion.getACCEPT_ALL());
+                BindingContext.EMPTY, "delta trace for caching resolve of", context.call, BindingTraceFilter.Companion.getACCEPT_ALL(), false);
         traceToResolveCall.addOwnDataTo(deltasTraceToCacheResolve);
 
         context.resolutionResultsCache.record(call, results, context, tracing, deltasTraceToCacheResolve);

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

@@ -31,13 +31,14 @@ import java.util.Map;
 
 public class SlicedMapImpl implements MutableSlicedMap {
 
-    public static SlicedMapImpl create() {
+    private final boolean alwaysAllowRewrite;
-        return new SlicedMapImpl();
-    }
-
     private final Map<Object, KeyFMap> map = new THashMap<>(0);
     private Multimap<WritableSlice<?, ?>, Object> collectiveSliceKeys = null;
 
+    public SlicedMapImpl(boolean alwaysAllowRewrite) {
+        this.alwaysAllowRewrite = alwaysAllowRewrite;
+    }
+
     @Override
     public <K, V> void put(WritableSlice<K, V> slice, K key, V value) {
         if (!slice.check(key, value)) {
@@ -52,7 +53,7 @@ public class SlicedMapImpl implements MutableSlicedMap {
         Key<V> sliceKey = slice.getKey();
 
         RewritePolicy rewritePolicy = slice.getRewritePolicy();
-        if (rewritePolicy.rewriteProcessingNeeded(key)) {
+        if (!alwaysAllowRewrite && rewritePolicy.rewriteProcessingNeeded(key)) {
             V oldValue = holder.get(sliceKey);
             if (oldValue != null) {
                 //noinspection unchecked

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

@@ -29,6 +29,7 @@ public class TrackingSlicedMap extends SlicedMapImpl {
     private final boolean trackWithStackTraces;
 
     public TrackingSlicedMap(boolean trackWithStackTraces) {
+        super(false);
         this.trackWithStackTraces = trackWithStackTraces;
     }
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/KotlinResolveCache.kt

@@ -159,7 +159,7 @@ private object KotlinResolveDataProvider {
             }
 
             val resolveSession = componentProvider.get<ResolveSession>()
-            val trace = DelegatingBindingTrace(resolveSession.bindingContext, "Trace for resolution of " + analyzableElement)
+            val trace = DelegatingBindingTrace(resolveSession.bindingContext, "Trace for resolution of " + analyzableElement, allowSliceRewrite = true)
 
             val targetPlatform = TargetPlatformDetector.getPlatform(analyzableElement.containingKtFile)
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/project/ResolveElementCache.kt

@@ -344,7 +344,7 @@ class ResolveElementCache(
             }
         }
 
-        val controlFlowTrace = DelegatingBindingTrace(trace.bindingContext, "Element control flow resolve", resolveElement)
+        val controlFlowTrace = DelegatingBindingTrace(trace.bindingContext, "Element control flow resolve", resolveElement, allowSliceRewrite = true)
         ControlFlowInformationProvider(resolveElement, controlFlowTrace, resolveElement.languageVersionSettings).checkDeclaration()
         controlFlowTrace.addOwnDataTo(trace, null, false)
 
@@ -595,7 +595,7 @@ class ResolveElementCache(
     // All additional resolve should be done to separate trace
     private fun createDelegatingTrace(resolveElement: KtElement, filter: BindingTraceFilter): BindingTrace {
         return resolveSession.storageManager.createSafeTrace(
-                DelegatingBindingTrace(resolveSession.bindingContext, "trace to resolve element", resolveElement, filter))
+                DelegatingBindingTrace(resolveSession.bindingContext, "trace to resolve element", resolveElement, filter, allowSliceRewrite = true))
     }
 
     private class BodyResolveContextForLazy(