JVM IR: clear BindingContext after psi2ir

This helps to reduce peak memory in lowerings/codegen by a lot.

A more robust approach would be to have a separate BindingContext for
each file, and clear each of them after running psi2ir on it. This would
also lower peak memory usage in psi2ir.

Provide a fallback workaround compiler argument
-Xir-do-not-clear-binding-context just in case BindingContext is in fact
used somewhere and it's not caught by tests.

compiler/backend/src/org/jetbrains/kotlin/codegen/state/GenerationState.kt

@@ -53,7 +53,7 @@ class GenerationState private constructor(
     val project: Project,
     builderFactory: ClassBuilderFactory,
     val module: ModuleDescriptor,
-    bindingContext: BindingContext,
+    val originalFrontendBindingContext: BindingContext,
     val files: List<KtFile>,
     val configuration: CompilerConfiguration,
     val generateDeclaredClassFilter: GenerateClassFilter,
@@ -178,7 +178,7 @@ class GenerationState private constructor(
     }
 
     val extraJvmDiagnosticsTrace: BindingTrace =
-        DelegatingBindingTrace(bindingContext, "For extra diagnostics in ${this::class.java}", false)
+        DelegatingBindingTrace(originalFrontendBindingContext, "For extra diagnostics in ${this::class.java}", false)
     private val interceptedBuilderFactory: ClassBuilderFactory
     private var used = false
 
@@ -199,13 +199,13 @@ class GenerationState private constructor(
     val moduleName: String = moduleName ?: JvmCodegenUtil.getModuleName(module)
     val classBuilderMode: ClassBuilderMode = builderFactory.classBuilderMode
     val bindingTrace: BindingTrace = DelegatingBindingTrace(
-        bindingContext, "trace in GenerationState",
+        originalFrontendBindingContext, "trace in GenerationState",
         filter = if (wantsDiagnostics) BindingTraceFilter.ACCEPT_ALL else BindingTraceFilter.NO_DIAGNOSTICS
     )
     val bindingContext: BindingContext = bindingTrace.bindingContext
-    val mainFunctionDetector = MainFunctionDetector(bindingContext, languageVersionSettings)
+    val mainFunctionDetector = MainFunctionDetector(originalFrontendBindingContext, languageVersionSettings)
     val typeMapper: KotlinTypeMapper = KotlinTypeMapper(
-        this.bindingContext,
+        bindingContext,
         classBuilderMode,
         this.moduleName,
         languageVersionSettings,
@@ -311,7 +311,7 @@ class GenerationState private constructor(
                         it
                     else
                         BuilderFactoryForDuplicateSignatureDiagnostics(
-                            it, this.bindingContext, diagnostics, this.moduleName, this.languageVersionSettings,
+                            it, bindingContext, diagnostics, this.moduleName, languageVersionSettings,
                             shouldGenerate = { origin -> !shouldOnlyCollectSignatures(origin) },
                         ).apply { duplicateSignatureFactory = this }
                 },
@@ -322,7 +322,7 @@ class GenerationState private constructor(
                 }
             )
             .wrapWith(ClassBuilderInterceptorExtension.getInstances(project)) { classBuilderFactory, extension ->
-                extension.interceptClassBuilderFactory(classBuilderFactory, bindingContext, diagnostics)
+                extension.interceptClassBuilderFactory(classBuilderFactory, originalFrontendBindingContext, diagnostics)
             }
 
         this.factory = ClassFileFactory(this, interceptedBuilderFactory)

compiler/cli/cli-common/src/org/jetbrains/kotlin/cli/common/arguments/K2JVMCompilerArguments.kt

@@ -108,6 +108,12 @@ class K2JVMCompilerArguments : CommonCompilerArguments() {
     )
     var isIrWithStableAbi: Boolean by FreezableVar(false)
 
+    @Argument(
+        value = "-Xir-do-not-clear-binding-context",
+        description = "When using the IR backend, do not clear BindingContext between psi2ir and lowerings"
+    )
+    var doNotClearBindingContext: Boolean by FreezableVar(false)
+
     @Argument(value = "-Xmodule-path", valueDescription = "<path>", description = "Paths where to find Java 9+ modules")
     var javaModulePath: String? by NullableStringFreezableVar(null)
 

compiler/cli/src/org/jetbrains/kotlin/cli/jvm/jvmArguments.kt

@@ -173,6 +173,7 @@ fun CompilerConfiguration.configureAdvancedJvmOptions(arguments: K2JVMCompilerAr
 
     put(JVMConfigurationKeys.IR, arguments.useIR && !arguments.noUseIR)
     put(JVMConfigurationKeys.IS_IR_WITH_STABLE_ABI, arguments.isIrWithStableAbi)
+    put(JVMConfigurationKeys.DO_NOT_CLEAR_BINDING_CONTEXT, arguments.doNotClearBindingContext)
     put(JVMConfigurationKeys.DISABLE_CALL_ASSERTIONS, arguments.noCallAssertions)
     put(JVMConfigurationKeys.DISABLE_RECEIVER_ASSERTIONS, arguments.noReceiverAssertions)
     put(JVMConfigurationKeys.DISABLE_PARAM_ASSERTIONS, arguments.noParamAssertions)

compiler/config.jvm/src/org/jetbrains/kotlin/config/JVMConfigurationKeys.java

@@ -120,6 +120,9 @@ public class JVMConfigurationKeys {
     public static final CompilerConfigurationKey<Boolean> IS_IR_WITH_STABLE_ABI =
             CompilerConfigurationKey.create("Is IR with stable ABI");
 
+    public static final CompilerConfigurationKey<Boolean> DO_NOT_CLEAR_BINDING_CONTEXT =
+            CompilerConfigurationKey.create("When using the IR backend, do not clear BindingContext between psi2ir and lowerings");
+
     public static final CompilerConfigurationKey<Boolean> NO_OPTIMIZED_CALLABLE_REFERENCES =
             CompilerConfigurationKey.create("Do not use optimized callable reference superclasses available from 1.4");
 

compiler/frontend/src/org/jetbrains/kotlin/checkers/utils/CheckerTestUtil.kt

@@ -51,7 +51,8 @@ import java.util.regex.Pattern
 data class DiagnosticsRenderingConfiguration(
     val platform: String?,
     val withNewInference: Boolean,
-    val languageVersionSettings: LanguageVersionSettings?
+    val languageVersionSettings: LanguageVersionSettings?,
+    val skipDebugInfoDiagnostics: Boolean = false,
 )
 
 object CheckerTestUtil {
@@ -130,18 +131,20 @@ object CheckerTestUtil {
             diagnostics.add(ActualDiagnostic(SyntaxErrorDiagnostic(errorElement), configuration.platform, configuration.withNewInference))
         }
 
-        diagnostics.addAll(
-            getDebugInfoDiagnostics(
-                root,
-                bindingContext,
-                markDynamicCalls,
-                dynamicCallDescriptors,
-                configuration,
-                dataFlowValueFactory,
-                moduleDescriptor,
-                diagnosedRanges
+        if (!configuration.skipDebugInfoDiagnostics) {
+            diagnostics.addAll(
+                getDebugInfoDiagnostics(
+                    root,
+                    bindingContext,
+                    markDynamicCalls,
+                    dynamicCallDescriptors,
+                    configuration,
+                    dataFlowValueFactory,
+                    moduleDescriptor,
+                    diagnosedRanges
+                )
             )
-        )
+        }
 
         return diagnostics
     }

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

@@ -37,11 +37,9 @@ public class BindingTraceContext implements BindingTrace {
     /* package */ final static boolean TRACK_WITH_STACK_TRACES = true;
 
     private final MutableSlicedMap map;
-    @Nullable private final MutableDiagnosticsWithSuppression mutableDiagnostics;
-    @NotNull private final BindingTraceFilter filter;
-
-    private final BindingContext bindingContext = new BindingContext() {
+    private final MutableDiagnosticsWithSuppression mutableDiagnostics;
 
+    private final BindingContext bindingContext = new CleanableBindingContext() {
         @NotNull
         @Override
         public Diagnostics getDiagnostics() {
@@ -76,6 +74,11 @@ public class BindingTraceContext implements BindingTrace {
         public void addOwnDataTo(@NotNull BindingTrace trace, boolean commitDiagnostics) {
             BindingContextUtils.addOwnDataTo(trace, null, commitDiagnostics, map, mutableDiagnostics);
         }
+
+        @Override
+        public void clear() {
+            map.clear();
+        }
     };
 
     public BindingTraceContext() {
@@ -87,7 +90,6 @@ public class BindingTraceContext implements BindingTrace {
     }
 
     public BindingTraceContext(BindingTraceFilter filter, boolean allowSliceRewrite) {
-        //noinspection ConstantConditions
         this(TRACK_REWRITES && !allowSliceRewrite ? new TrackingSlicedMap(TRACK_WITH_STACK_TRACES) : new SlicedMapImpl(allowSliceRewrite), filter);
     }
 
@@ -97,7 +99,6 @@ public class BindingTraceContext implements BindingTrace {
         this.mutableDiagnostics = !filter.getIgnoreDiagnostics()
                                   ? new MutableDiagnosticsWithSuppression(bindingContext, Diagnostics.Companion.getEMPTY())
                                   : null;
-        this.filter = filter;
     }
 
     @TestOnly

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

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.resolve
+
+interface CleanableBindingContext : BindingContext {
+    /**
+     * Removes all recorded data except diagnostics.
+     */
+    fun clear()
+}

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmBackendFacade.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.backend.jvm.codegen.DescriptorMetadataSerializer
 import org.jetbrains.kotlin.backend.jvm.lower.MultifileFacadeFileEntry
 import org.jetbrains.kotlin.backend.jvm.serialization.JvmIdSignatureDescriptor
 import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.config.JVMConfigurationKeys
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
 import org.jetbrains.kotlin.descriptors.konan.DeserializedKlibModuleOrigin
 import org.jetbrains.kotlin.descriptors.konan.KlibModuleOrigin
@@ -36,6 +37,7 @@ import org.jetbrains.kotlin.psi2ir.Psi2IrConfiguration
 import org.jetbrains.kotlin.psi2ir.Psi2IrTranslator
 import org.jetbrains.kotlin.psi2ir.PsiSourceManager
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.CleanableBindingContext
 
 object JvmBackendFacade {
     fun doGenerateFiles(files: Collection<KtFile>, state: GenerationState, phaseConfig: PhaseConfig) {
@@ -109,6 +111,12 @@ object JvmBackendFacade {
         // We need to compile all files we reference in Klibs
         irModuleFragment.files.addAll(dependencies.flatMap { it.files })
 
+        if (!state.configuration.getBoolean(JVMConfigurationKeys.DO_NOT_CLEAR_BINDING_CONTEXT)) {
+            val originalBindingContext = state.originalFrontendBindingContext as? CleanableBindingContext
+                ?: error("BindingContext should be cleanable in JVM IR to avoid leaking memory: ${state.originalFrontendBindingContext}")
+            originalBindingContext.clear()
+        }
+
         doGenerateFilesInternal(
             state, irModuleFragment, psi2irContext.symbolTable, psi2irContext.sourceManager, phaseConfig,
             irProviders, extensions, ::DescriptorMetadataSerializer

compiler/testData/cli/jvm/extraHelp.out

@@ -19,6 +19,8 @@ where advanced options include:
                              'enable' since language version 1.3
   -Xdump-declarations-to=<path> Path to JSON file to dump Java to Kotlin declaration mappings
   -Xdisable-standard-script  Disable standard kotlin script support
+  -Xir-do-not-clear-binding-context
+                             When using the IR backend, do not clear BindingContext between psi2ir and lowerings
   -Xemit-jvm-type-annotations Emit JVM type annotations in bytecode
   -Xfriend-paths=<path>      Paths to output directories for friend modules (whose internals should be visible)
   -Xmultifile-parts-inherit  Compile multifile classes as a hierarchy of parts and facade

compiler/tests-common/tests/org/jetbrains/kotlin/checkers/BaseDiagnosticsTest.kt

@@ -34,10 +34,7 @@ import org.jetbrains.kotlin.checkers.diagnostics.factories.SyntaxErrorDiagnostic
 import org.jetbrains.kotlin.checkers.utils.CheckerTestUtil
 import org.jetbrains.kotlin.checkers.utils.DiagnosticsRenderingConfiguration
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
-import org.jetbrains.kotlin.config.JvmTarget
-import org.jetbrains.kotlin.config.LanguageFeature
-import org.jetbrains.kotlin.config.LanguageVersionSettings
-import org.jetbrains.kotlin.config.LanguageVersionSettingsImpl
+import org.jetbrains.kotlin.config.*
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.impl.ModuleDescriptorImpl
 import org.jetbrains.kotlin.diagnostics.*
@@ -269,6 +266,8 @@ abstract class BaseDiagnosticsTest : KotlinMultiFileTestWithJava<TestModule, Tes
                     platform = null,
                     withNewInference,
                     languageVersionSettings,
+                    // When using JVM IR, binding context is empty at the end of compilation, so debug info markers can't be computed.
+                    environment.configuration.getBoolean(JVMConfigurationKeys.IR),
                 ),
                 DataFlowValueFactoryImpl(languageVersionSettings),
                 moduleDescriptor,

idea/jvm-debugger/jvm-debugger-test/test/org/jetbrains/kotlin/idea/debugger/test/DebuggerTestCompilerFacility.kt

@@ -169,6 +169,7 @@ class DebuggerTestCompilerFacility(
         val configuration = CompilerConfiguration()
         configuration.put(JVMConfigurationKeys.JVM_TARGET, jvmTarget)
         configuration.put(JVMConfigurationKeys.IR, useIrBackend)
+        configuration.put(JVMConfigurationKeys.DO_NOT_CLEAR_BINDING_CONTEXT, true)
 
         val state = GenerationUtils.generateFiles(project, files, configuration, ClassBuilderFactories.BINARIES, analysisResult) {
             generateDeclaredClassFilter(GenerationState.GenerateClassFilter.GENERATE_ALL)