Add addOwnDataTo() to BindingContext

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

@@ -80,6 +80,11 @@ public interface BindingContext {
         public JetType getType(@NotNull JetExpression expression) {
             return null;
         }
+
+        @Override
+        public void addOwnDataTo(@NotNull BindingTrace trace, boolean commitDiagnostics) {
+            // Do nothing
+        }
     };
 
     WritableSlice<AnnotationDescriptor, JetAnnotationEntry> ANNOTATION_DESCRIPTOR_TO_PSI_ELEMENT = Slices.createSimpleSlice();
@@ -264,4 +269,6 @@ public interface BindingContext {
 
     @Nullable
     JetType getType(@NotNull JetExpression expression);
+
+    void addOwnDataTo(@NotNull BindingTrace trace, boolean commitDiagnostics);
 }

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

@@ -20,19 +20,24 @@ import com.google.common.collect.Lists;
 import com.intellij.openapi.util.Pair;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
+import kotlin.jvm.functions.Function3;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.descriptors.*;
+import org.jetbrains.kotlin.diagnostics.Diagnostic;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.resolve.calls.callUtil.CallUtilPackage;
 import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall;
 import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall;
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo;
+import org.jetbrains.kotlin.resolve.diagnostics.MutableDiagnosticsWithSuppression;
 import org.jetbrains.kotlin.types.JetType;
 import org.jetbrains.kotlin.types.TypeUtils;
 import org.jetbrains.kotlin.types.expressions.JetTypeInfo;
 import org.jetbrains.kotlin.types.expressions.typeInfoFactory.TypeInfoFactoryPackage;
+import org.jetbrains.kotlin.util.slicedMap.MutableSlicedMap;
 import org.jetbrains.kotlin.util.slicedMap.ReadOnlySlice;
+import org.jetbrains.kotlin.util.slicedMap.WritableSlice;
 
 import java.util.Collection;
 
@@ -205,4 +210,29 @@ public class BindingContextUtils {
     ) {
         return bindingContext.get(CONSTRUCTOR_RESOLVED_DELEGATION_CALL, constructorDescriptor);
     }
+
+    static void addOwnDataTo(
+            @NotNull final BindingTrace trace, @Nullable final TraceEntryFilter filter, boolean commitDiagnostics,
+            @NotNull MutableSlicedMap map, MutableDiagnosticsWithSuppression diagnostics
+    ) {
+        map.forEach(new Function3<WritableSlice, Object, Object, Void>() {
+            @Override
+            public Void invoke(WritableSlice slice, Object key, Object value) {
+                if (filter == null || filter.accept(slice, key)) {
+                    trace.record(slice, key, value);
+                }
+
+                return null;
+            }
+        });
+
+        if (!commitDiagnostics) return;
+
+        for (Diagnostic diagnostic : diagnostics.getOwnDiagnostics()) {
+            if (filter == null || filter.accept(null, diagnostic.getPsiElement())) {
+                trace.report(diagnostic);
+            }
+        }
+
+    }
 }

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

@@ -70,6 +70,11 @@ public class BindingTraceContext implements BindingTrace {
         public JetType getType(@NotNull JetExpression expression) {
             return BindingTraceContext.this.getType(expression);
         }
+
+        @Override
+        public void addOwnDataTo(@NotNull BindingTrace trace, boolean commitDiagnostics) {
+            BindingContextUtils.addOwnDataTo(trace, null, commitDiagnostics, map, mutableDiagnostics);
+        }
     };
 
     public BindingTraceContext() {

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

@@ -60,6 +60,10 @@ public class CompositeBindingContext private constructor(
         return CompositeDiagnostics(delegates.map { it.getDiagnostics() })
     }
 
+    override fun addOwnDataTo(trace: BindingTrace, commitDiagnostics: Boolean) {
+        // Do nothing
+    }
+
     private class CompositeDiagnostics(
             private val delegates: List<Diagnostics>
     ) : Diagnostics {

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

@@ -19,7 +19,6 @@ package org.jetbrains.kotlin.resolve;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
-import kotlin.jvm.functions.Function3;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.TestOnly;
@@ -67,7 +66,11 @@ public class DelegatingBindingTrace implements BindingTrace {
         @Override
         public <K, V> Collection<K> getKeys(WritableSlice<K, V> slice) {
             return DelegatingBindingTrace.this.getKeys(slice);
+        }
 
+        @Override
+        public void addOwnDataTo(@NotNull BindingTrace trace, boolean commitDiagnostics) {
+            BindingContextUtils.addOwnDataTo(trace, null, commitDiagnostics, map, mutableDiagnostics);
         }
 
         @NotNull
@@ -162,25 +165,8 @@ public class DelegatingBindingTrace implements BindingTrace {
         clear();
     }
 
-    public void addAllMyDataTo(@NotNull final BindingTrace trace, @Nullable final TraceEntryFilter filter, boolean commitDiagnostics) {
+    public void addAllMyDataTo(@NotNull BindingTrace trace, @Nullable TraceEntryFilter filter, boolean commitDiagnostics) {
-        map.forEach(new Function3<WritableSlice, Object, Object, Void>() {
+        BindingContextUtils.addOwnDataTo(trace, filter, commitDiagnostics, map, mutableDiagnostics);
-            @Override
-            public Void invoke(WritableSlice slice, Object key, Object value) {
-                if (filter == null || filter.accept(slice, key)) {
-                    trace.record(slice, key, value);
-                }
-
-                return null;
-            }
-        });
-
-        if (!commitDiagnostics) return;
-
-        for (Diagnostic diagnostic : mutableDiagnostics.getOwnDiagnostics()) {
-            if (filter == null || filter.accept(null, diagnostic.getPsiElement())) {
-                trace.report(diagnostic);
-            }
-        }
     }
 
     public void clear() {

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

@@ -18,8 +18,6 @@ package org.jetbrains.kotlin.resolve;
 
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
-import org.jetbrains.kotlin.psi.JetExpression;
-import org.jetbrains.kotlin.types.JetType;
 
 public class TemporaryBindingTrace extends DelegatingBindingTrace {
 

compiler/frontend/src/org/jetbrains/kotlin/storage/LockBasedLazyResolveStorageManager.kt

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.util.slicedMap.ReadOnlySlice
 import org.jetbrains.kotlin.util.slicedMap.WritableSlice
+import org.jetbrains.kotlin.resolve.TraceEntryFilter
 import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics
 import com.intellij.util.containers.ContainerUtil
 import org.jetbrains.kotlin.psi.JetExpression
@@ -46,6 +47,10 @@ public class LockBasedLazyResolveStorageManager(private val storageManager: Stor
 
         override fun <K, V> getKeys(slice: WritableSlice<K, V>) = storageManager.compute { context.getKeys<K, V>(slice) }
 
+        override fun addOwnDataTo(trace: BindingTrace, commitDiagnostics: Boolean) {
+            storageManager.compute { context.addOwnDataTo(trace, commitDiagnostics) }
+        }
+
         TestOnly
         override fun <K, V> getSliceContents(slice: ReadOnlySlice<K, V>) = storageManager.compute { context.getSliceContents<K, V>(slice) }
     }

compiler/tests/org/jetbrains/kotlin/test/JetTestUtils.java

@@ -150,6 +150,11 @@ public class JetTestUtils {
                 public JetType getType(@NotNull JetExpression expression) {
                     return DUMMY_TRACE.getType(expression);
                 }
+
+                @Override
+                public void addOwnDataTo(@NotNull BindingTrace trace, boolean commitDiagnostics) {
+                    // do nothing
+                }
             };
         }
 
@@ -228,6 +233,11 @@ public class JetTestUtils {
                 public JetType getType(@NotNull JetExpression expression) {
                     return DUMMY_EXCEPTION_ON_ERROR_TRACE.getType(expression);
                 }
+
+                @Override
+                public void addOwnDataTo(@NotNull BindingTrace trace, boolean commitDiagnostics) {
+                    // do nothing
+                }
             };
         }