[EE-IR] Support mutations by evaluator fragments

- box captured variables using `Ref`s, the same mechanism that's used
  for shared variables in closures.

compiler/ir/backend.jvm/entrypoint/src/org/jetbrains/kotlin/backend/jvm/JvmIrCodegenFactory.kt

@@ -8,10 +8,7 @@ package org.jetbrains.kotlin.backend.jvm
 import org.jetbrains.kotlin.analyzer.hasJdkCapability
 import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
 import org.jetbrains.kotlin.backend.common.extensions.IrPluginContextImpl
-import org.jetbrains.kotlin.backend.common.phaser.NamedCompilerPhase
-import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
-import org.jetbrains.kotlin.backend.common.phaser.invokeToplevel
-import org.jetbrains.kotlin.backend.common.phaser.then
+import org.jetbrains.kotlin.backend.common.phaser.*
 import org.jetbrains.kotlin.backend.jvm.ir.getKtFile
 import org.jetbrains.kotlin.backend.jvm.serialization.JvmIdSignatureDescriptor
 import org.jetbrains.kotlin.codegen.CodegenFactory
@@ -50,7 +47,7 @@ open class JvmIrCodegenFactory(
     private val externalMangler: JvmDescriptorMangler? = null,
     private val externalSymbolTable: SymbolTable? = null,
     private val jvmGeneratorExtensions: JvmGeneratorExtensionsImpl = JvmGeneratorExtensionsImpl(configuration),
-    private val prefixPhases: NamedCompilerPhase<JvmBackendContext, IrModuleFragment>? = null,
+    private val prefixPhases: CompilerPhase<JvmBackendContext, IrModuleFragment, IrModuleFragment>? = null,
     private val evaluatorFragmentInfoForPsi2Ir: EvaluatorFragmentInfo? = null,
 ) : CodegenFactory {
     data class JvmIrBackendInput(

compiler/ir/backend.jvm/lower/src/org/jetbrains/kotlin/backend/jvm/lower/FragmentSharedVariablesLowering.kt

@@ -0,0 +1,119 @@
+/*
+ * Copyright 2010-2021 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.backend.jvm.lower
+
+import org.jetbrains.kotlin.backend.common.FileLoweringPass
+import org.jetbrains.kotlin.backend.common.IrElementTransformerVoidWithContext
+import org.jetbrains.kotlin.backend.common.ir.copyTo
+import org.jetbrains.kotlin.backend.common.phaser.makeIrModulePhase
+import org.jetbrains.kotlin.backend.jvm.JvmBackendContext
+import org.jetbrains.kotlin.ir.IrStatement
+import org.jetbrains.kotlin.ir.declarations.IrDeclarationOrigin
+import org.jetbrains.kotlin.ir.declarations.IrFile
+import org.jetbrains.kotlin.ir.declarations.IrFunction
+import org.jetbrains.kotlin.ir.expressions.IrExpression
+import org.jetbrains.kotlin.ir.expressions.IrGetValue
+import org.jetbrains.kotlin.ir.expressions.IrSetValue
+import org.jetbrains.kotlin.ir.symbols.IrValueParameterSymbol
+import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
+import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid
+
+// Used from the IntelliJ IDEA Kotlin Debugger Plug-In
+@Suppress("unused")
+val fragmentSharedVariablesLowering = makeIrModulePhase(
+    ::FragmentSharedVariablesLowering,
+    name = "FragmentSharedVariablesLowering",
+    description = "Promotes captured variables that are modified by the fragment to shared variables"
+)
+
+// This lowering is a preprocessor for IR in order to support the compilation
+// scheme used by the "Evaluate Expression..." mechanism of the IntelliJ plug-in
+// for Kotlin debugging.
+//
+// Fragments are compiled as the body of an enclosing function that close the free
+// variables of the fragment as parameters. The values of these are then extracted
+// from the stack at the current breakpoint, and the fragment code is invoked with
+// these values to evaluate the expression.
+//
+// If the parameter is a shared variable, e.g. `IntRef` (the same mechanism used to
+// implement captures of lambdas) the value extracted from the stack is
+// automatically boxed in a `Ref` before being passed to the fragment.
+//
+// Upon return, all `Ref`s are written back into the stack, thus allowing fragments
+// to modify the state of the program being debugged.
+//
+// This lowering promotes these parameters to `Ref`s, as deemed appropriate by
+// psi2ir's Fragment generation.
+//
+// The reason for this "phasing" is that the JVM specific infrastructure (e.g.
+// symbols for `Ref`s) have not been loaded when psi2ir runs, as psi2ir is designed
+// to be backend agnostic. So, we "tag" the appropriate parameters with a new
+// JvmIrDeclarationOrigin that we can then detect in this lowering.
+//
+// See `FragmentDeclarationGenerator.kt:declareParameter` for the front half
+// of this logic.
+class FragmentSharedVariablesLowering(
+    val context: JvmBackendContext
+) : IrElementTransformerVoidWithContext(), FileLoweringPass {
+
+    companion object {
+        // Echo of GENERATED_FUNCTION_NAME in the JVM Debugger plug-in.
+        // TODO: Find a good common dependency of JVM Debugger and IR Compiler and deduplicate this
+        const val GENERATED_FUNCTION_NAME = "generated_for_debugger_fun"
+    }
+
+    override fun lower(irFile: IrFile) {
+        irFile.transformChildrenVoid(this)
+    }
+
+    override fun visitFunctionNew(declaration: IrFunction): IrStatement {
+        if (declaration.name.asString() != GENERATED_FUNCTION_NAME) {
+            return super.visitFunctionNew(declaration)
+        }
+
+        val promotedParameters = promoteParametersForCapturesToRefs(declaration)
+        replaceUseOfPromotedParametersWithRefs(declaration, promotedParameters)
+        return declaration
+    }
+
+    private fun promoteParametersForCapturesToRefs(declaration: IrFunction): Map<IrValueParameterSymbol, IrValueParameterSymbol> {
+        val promotedParameters = mutableMapOf<IrValueParameterSymbol, IrValueParameterSymbol>()
+        declaration.valueParameters = declaration.valueParameters.map {
+            if (it.origin == IrDeclarationOrigin.SHARED_VARIABLE_IN_EVALUATOR_FRAGMENT) {
+                val newParameter =
+                    it.copyTo(
+                        declaration,
+                        type = context.sharedVariablesManager.getIrType(it.type),
+                        origin = IrDeclarationOrigin.DEFINED
+                    )
+                promotedParameters[it.symbol] = newParameter.symbol
+                newParameter
+            } else {
+                it
+            }
+        }
+        return promotedParameters
+    }
+
+    private fun replaceUseOfPromotedParametersWithRefs(
+        declaration: IrFunction,
+        promotedParameters: Map<IrValueParameterSymbol, IrValueParameterSymbol>
+    ) {
+        declaration.body!!.transformChildrenVoid(object : IrElementTransformerVoid() {
+            override fun visitGetValue(expression: IrGetValue): IrExpression {
+                expression.transformChildrenVoid(this)
+                val newDeclaration = promotedParameters[expression.symbol] ?: return expression
+                return context.sharedVariablesManager.getSharedValue(newDeclaration, expression)
+            }
+
+            override fun visitSetValue(expression: IrSetValue): IrExpression {
+                expression.transformChildrenVoid(this)
+                val newDeclaration = promotedParameters[expression.symbol] ?: return expression
+                return context.sharedVariablesManager.setSharedValue(newDeclaration, expression)
+            }
+        })
+    }
+}

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

@@ -19,7 +19,7 @@ import org.jetbrains.kotlin.ir.expressions.IrExpression
 import org.jetbrains.kotlin.ir.expressions.IrGetValue
 import org.jetbrains.kotlin.ir.expressions.IrSetValue
 import org.jetbrains.kotlin.ir.expressions.impl.*
-import org.jetbrains.kotlin.ir.symbols.IrVariableSymbol
+import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.symbols.impl.IrVariableSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.name.FqName
@@ -39,7 +39,7 @@ class JvmSharedVariablesManager(
         name = Name.identifier("Ref")
     }
 
-    private class RefProvider(val refClass: IrClass, elementType: IrType) {
+    class RefProvider(val refClass: IrClass, elementType: IrType) {
         val refConstructor = refClass.addConstructor {
             origin = IrDeclarationOrigin.IR_BUILTINS_STUB
         }
@@ -75,7 +75,7 @@ class JvmSharedVariablesManager(
         RefProvider(refClass, refClass.typeParameters[0].defaultType)
     }
 
-    private fun getProvider(valueType: IrType): RefProvider =
+    fun getProvider(valueType: IrType): RefProvider =
         if (valueType.isPrimitiveType())
             primitiveRefProviders.getValue(valueType.classifierOrFail)
         else
@@ -129,7 +129,7 @@ class JvmSharedVariablesManager(
             putValueArgument(0, value)
         }
 
-    override fun getSharedValue(sharedVariableSymbol: IrVariableSymbol, originalGet: IrGetValue): IrExpression =
+    override fun getSharedValue(sharedVariableSymbol: IrValueSymbol, originalGet: IrGetValue): IrExpression =
         with(originalGet) {
             val unboxedType = InlineClassAbi.unboxType(symbol.owner.type)
             val provider = getProvider(unboxedType ?: symbol.owner.type)
@@ -138,7 +138,7 @@ class JvmSharedVariablesManager(
             unboxedType?.let { unsafeCoerce(unboxedRead, it, symbol.owner.type) } ?: unboxedRead
         }
 
-    override fun setSharedValue(sharedVariableSymbol: IrVariableSymbol, originalSet: IrSetValue): IrExpression =
+    override fun setSharedValue(sharedVariableSymbol: IrValueSymbol, originalSet: IrSetValue): IrExpression =
         with(originalSet) {
             val unboxedType = InlineClassAbi.unboxType(symbol.owner.type)
             val unboxedValue = unboxedType?.let { unsafeCoerce(value, symbol.owner.type, it) } ?: value
@@ -146,6 +146,13 @@ class JvmSharedVariablesManager(
             val receiver = IrGetValueImpl(startOffset, endOffset, sharedVariableSymbol)
             IrSetFieldImpl(startOffset, endOffset, provider.elementField.symbol, receiver, unboxedValue, type, origin)
         }
+
+    @Suppress("MemberVisibilityCanBePrivate") // Used by FragmentSharedVariablesLowering
+    fun getIrType(originalType: IrType): IrType {
+        val provider = getProvider(InlineClassAbi.unboxType(originalType) ?: originalType)
+        val typeArguments = provider.refClass.typeParameters.map { originalType }
+        return provider.refClass.typeWith(typeArguments)
+    }
 }
 
 private inline fun IrFactory.addClass(