Implement REPL support in IR scripting

2021-01-10 21:17:14 +01:00
parent c066b7843c
commit 375441832e
13 changed files with 242 additions and 69 deletions
@@ -48,7 +48,7 @@ import org.jetbrains.kotlin.serialization.deserialization.descriptors.Deserializ
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.typeUtil.replaceAnnotations

-class JvmGeneratorExtensionsImpl(private val generateFacades: Boolean = true) : GeneratorExtensions(), JvmGeneratorExtensions {
+open class JvmGeneratorExtensionsImpl(private val generateFacades: Boolean = true) : GeneratorExtensions(), JvmGeneratorExtensions {
    override val classNameOverride: MutableMap<IrClass, JvmClassName> = mutableMapOf()

    override val samConversion: SamConversion
@@ -37,17 +37,23 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi2ir.Psi2IrConfiguration
 import org.jetbrains.kotlin.psi2ir.Psi2IrTranslator
 import org.jetbrains.kotlin.psi2ir.generators.DeclarationStubGeneratorImpl
+import org.jetbrains.kotlin.psi2ir.generators.GeneratorExtensions
 import org.jetbrains.kotlin.psi2ir.generators.generateTypicalIrProviderList
 import org.jetbrains.kotlin.resolve.CleanableBindingContext

-class JvmIrCodegenFactory(private val phaseConfig: PhaseConfig) : CodegenFactory {
+open class JvmIrCodegenFactory(
+    private val phaseConfig: PhaseConfig,
+    private val externalMangler: JvmManglerDesc? = null,
+    private val externalSymbolTable: SymbolTable? = null,
+    private val jvmGeneratorExtensions: JvmGeneratorExtensionsImpl = JvmGeneratorExtensionsImpl()
+) : CodegenFactory {
    data class JvmIrBackendInput(
        val state: GenerationState,
        val irModuleFragment: IrModuleFragment,
        val symbolTable: SymbolTable,
        val phaseConfig: PhaseConfig,
        val irProviders: List<IrProvider>,
-        val extensions: JvmGeneratorExtensions,
+        val extensions: JvmGeneratorExtensionsImpl,
        val backendExtension: JvmBackendExtension,
        val notifyCodegenStart: () -> Unit,
    )
@@ -59,18 +65,22 @@ class JvmIrCodegenFactory(private val phaseConfig: PhaseConfig) : CodegenFactory

    @JvmOverloads
    fun convertToIr(state: GenerationState, files: Collection<KtFile>, ignoreErrors: Boolean = false): JvmIrBackendInput {
-        val extensions = JvmGeneratorExtensionsImpl()
-        val mangler = JvmManglerDesc(MainFunctionDetector(state.bindingContext, state.languageVersionSettings))
+        val (mangler, symbolTable) =
+            if (externalSymbolTable != null) externalMangler!! to externalSymbolTable
+            else {
+                val mangler = JvmManglerDesc(MainFunctionDetector(state.bindingContext, state.languageVersionSettings))
+                val symbolTable = SymbolTable(JvmIdSignatureDescriptor(mangler), IrFactoryImpl, JvmNameProvider)
+                mangler to symbolTable
+            }
        val psi2ir = Psi2IrTranslator(state.languageVersionSettings, Psi2IrConfiguration(ignoreErrors))
-        val symbolTable = SymbolTable(JvmIdSignatureDescriptor(mangler), IrFactoryImpl, JvmNameProvider)
        val messageLogger = state.configuration[IrMessageLogger.IR_MESSAGE_LOGGER] ?: IrMessageLogger.None
-        val psi2irContext = psi2ir.createGeneratorContext(state.module, state.bindingContext, symbolTable, extensions)
+        val psi2irContext = psi2ir.createGeneratorContext(state.module, state.bindingContext, symbolTable, jvmGeneratorExtensions)
        val pluginExtensions = IrGenerationExtension.getInstances(state.project)
        val functionFactory = IrFunctionFactory(psi2irContext.irBuiltIns, symbolTable)
        psi2irContext.irBuiltIns.functionFactory = functionFactory

        val stubGenerator =
-            DeclarationStubGeneratorImpl(psi2irContext.moduleDescriptor, symbolTable, state.languageVersionSettings, extensions)
+            DeclarationStubGeneratorImpl(psi2irContext.moduleDescriptor, symbolTable, state.languageVersionSettings, jvmGeneratorExtensions)
        val frontEndContext = object : TranslationPluginContext {
            override val moduleDescriptor: ModuleDescriptor
                get() = psi2irContext.moduleDescriptor
@@ -153,7 +163,7 @@ class JvmIrCodegenFactory(private val phaseConfig: PhaseConfig) : CodegenFactory
            symbolTable,
            phaseConfig,
            irProviders,
-            extensions,
+            jvmGeneratorExtensions,
            JvmBackendExtension.Default,
        ) {}
    }
@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.declarations.impl.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
+import org.jetbrains.kotlin.ir.interpreter.toIrConst
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.symbols.impl.IrAnonymousInitializerSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
@@ -29,6 +30,7 @@ import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.util.OperatorNameConventions

 internal val scriptsToClassesPhase = makeCustomPhase<JvmBackendContext, IrModuleFragment>(
    name = "ScriptsToClasses",
@@ -77,12 +79,13 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext) {
            irScriptClass.superTypes += irScript.baseClass
            irScriptClass.parent = irFile
            irScriptClass.metadata = irScript.metadata
+            irScript.targetClass = irScriptClass.symbol
        }
    }

    private fun finalizeScriptClass(irScriptClass: IrClass, irScript: IrScript, symbolRemapper: ScriptsToClassesSymbolRemapper) {
        val typeRemapper = SimpleTypeRemapper(symbolRemapper)
-        val scriptTransformer = ScriptToClassTransformer(irScript, irScriptClass, symbolRemapper, typeRemapper)
+        val scriptTransformer = ScriptToClassTransformer(irScript, irScriptClass, symbolRemapper, typeRemapper, context)
        irScriptClass.thisReceiver = irScript.thisReceiver.run {
            transform(scriptTransformer, null)
        }
@@ -112,6 +115,9 @@ private class ScriptsToClassesLowering(val context: JvmBackendContext) {
                }
            }

+            irScript.earlierScriptsParameter?.let { earlierScriptdParameter ->
+                addConstructorParameter(earlierScriptdParameter, false)
+            }
            irScript.explicitCallParameters.forEach { addConstructorParameter(it, false) }
            irScript.implicitReceiversParameters.forEach { addConstructorParameter(it, false) }
            irScript.providedProperties.forEach { addConstructorParameter(it.first, false) }
@@ -300,7 +306,8 @@ private class ScriptToClassTransformer(
    val irScript: IrScript,
    val irScriptClass: IrClass,
    val symbolRemapper: SymbolRemapper,
-    val typeRemapper: TypeRemapper
+    val typeRemapper: TypeRemapper,
+    val context: JvmBackendContext
 ) : IrElementTransformerVoid() {

    private fun IrType.remapType() = typeRemapper.remapType(this)
@@ -425,6 +432,74 @@ private class ScriptToClassTransformer(
        }
        visitExpression(expression)
    }
+
+    private fun getAccessCallForEarlierScript(expression: IrDeclarationReference, maybeScriptType: IrType): IrExpression? {
+        if (irScript.earlierScripts?.isEmpty() != false) return null
+        val scriptSymbol = maybeScriptType.classifierOrNull ?: return null
+        val scriptSymbolOwner = scriptSymbol.owner
+        val earlierScriptIndex = when {
+            scriptSymbolOwner is IrScript ->
+                irScript.earlierScripts!!.indexOfFirst { it == scriptSymbol }
+            (scriptSymbolOwner as? IrClass)?.origin == IrDeclarationOrigin.SCRIPT_CLASS -> {
+                irScript.earlierScripts!!.indexOfFirst { it.owner.targetClass == scriptSymbol }
+            }
+            else -> return null
+        }
+        if (earlierScriptIndex >= 0) {
+            val objArray = context.irBuiltIns.arrayClass
+            val objArrayGet = objArray.functions.single { it.owner.name == OperatorNameConventions.GET }
+            val builder = context.createIrBuilder(expression.symbol)
+            val getPrevScriptObjectExpression = builder.irCall(objArrayGet).apply {
+                dispatchReceiver = builder.irGet(objArray.defaultType, irScript.earlierScriptsParameter!!.symbol)
+                putValueArgument(0, earlierScriptIndex.toIrConst(objArrayGet.owner.valueParameters.first().type))
+            }
+            val prevScriptClassType =
+                when {
+                    scriptSymbolOwner is IrScript -> scriptSymbolOwner.targetClass?.owner
+                    (scriptSymbolOwner as? IrClass)?.origin == IrDeclarationOrigin.SCRIPT_CLASS -> scriptSymbolOwner
+                    else -> null
+                }
+            return if (prevScriptClassType == null) getPrevScriptObjectExpression
+            else builder.irImplicitCast(getPrevScriptObjectExpression, prevScriptClassType.defaultType)
+        }
+        return null
+    }
+
+    override fun visitGetField(expression: IrGetField): IrExpression {
+        if (irScript.earlierScripts != null) {
+            val receiver = expression.receiver
+            if (receiver is IrGetValue && receiver.symbol.owner.name == Name.special("<this>")) {
+                val newReceiver = getAccessCallForEarlierScript(expression, receiver.type)
+                if (newReceiver != null) {
+                    val newGetField =
+                        IrGetFieldImpl(expression.startOffset, expression.endOffset, expression.symbol, expression.type, newReceiver)
+                    return super.visitGetField(newGetField)
+                }
+            }
+        }
+        return super.visitGetField(expression)
+    }
+
+    override fun visitCall(expression: IrCall): IrExpression {
+        if (irScript.earlierScripts != null) {
+            val target = expression.symbol.owner
+            val receiver: IrValueParameter? = target.dispatchReceiverParameter
+            if (target.origin == IrDeclarationOrigin.DEFAULT_PROPERTY_ACCESSOR && receiver?.name == Name.special("<this>")) {
+                val newReceiver = getAccessCallForEarlierScript(expression, receiver.type)
+                if (newReceiver != null) {
+                    val builder = context.createIrBuilder(expression.symbol)
+                    val newCall = builder.irCall(expression.symbol, target.returnType, origin = expression.origin).also {
+                        it.copyTypeAndValueArgumentsFrom(expression)
+                        it.dispatchReceiver = newReceiver
+                    }
+                    return super.visitCall(newCall)
+                }
+            }
+        }
+        return super.visitCall(expression)
+    }
+
+
 }

 private class ScriptsToClassesSymbolRemapper(
@@ -41,5 +41,5 @@ open class GeneratorExtensions : StubGeneratorExtensions() {
    open val shouldPreventDeprecatedIntegerValueTypeLiteralConversion: Boolean
        get() = false

-    open fun getPreviousScripts(): List<IrScriptSymbol> = emptyList()
+    open fun getPreviousScripts(): List<IrScriptSymbol>? = null
 }
@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.psi2ir.generators
 import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
 import org.jetbrains.kotlin.descriptors.ParameterDescriptor
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
+import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.assertCast
 import org.jetbrains.kotlin.ir.declarations.DescriptorMetadataSource
 import org.jetbrains.kotlin.ir.declarations.IrDeclaration
@@ -18,10 +19,13 @@ import org.jetbrains.kotlin.ir.expressions.IrStatementOrigin
 import org.jetbrains.kotlin.ir.expressions.impl.IrCompositeImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrSetFieldImpl
+import org.jetbrains.kotlin.ir.symbols.impl.IrValueParameterSymbolImpl
+import org.jetbrains.kotlin.ir.types.toArrayOrPrimitiveArrayType
 import org.jetbrains.kotlin.ir.util.indexOrMinusOne
 import org.jetbrains.kotlin.ir.util.isCrossinline
 import org.jetbrains.kotlin.ir.util.isNoinline
 import org.jetbrains.kotlin.ir.util.varargElementType
+import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtDestructuringDeclaration
 import org.jetbrains.kotlin.psi.KtScript
 import org.jetbrains.kotlin.psi.KtScriptInitializer
@@ -39,27 +43,12 @@ class ScriptGenerator(declarationGenerator: DeclarationGenerator) : DeclarationG
    fun generateScriptDeclaration(ktScript: KtScript): IrDeclaration? {
        val descriptor = getOrFail(BindingContext.DECLARATION_TO_DESCRIPTOR, ktScript) as ScriptDescriptor

-        val previousScripts = context.extensions.getPreviousScripts()
-
        return context.symbolTable.declareScript(descriptor).buildWithScope { irScript ->

            irScript.metadata = DescriptorMetadataSource.Class(descriptor)

-            // A workaround for the JS/REPL backend:
-            // JS backend doesn't save previously executed snippets anywhere, they should be taken for now from the context.symbolTable.listExistedScripts()
-            // on the other hand imported scripts are also stored there so to avoid clashes the imported scripts should be filtered out
-            // NOTE: that JVM IR is not properly tested with REPL, so it might have the same problem
-            // TODO: design and implement other schema for handling previous snippets
            val importedScripts = descriptor.implicitReceivers.filterIsInstanceTo(HashSet<ScriptDescriptor>())

-            // TODO: since script could reference instances of previous one their receivers have to be enlisted in its scope
-            // Remove this code once script is no longer represented by Class
-            previousScripts.forEach {
-                if (it.owner != irScript && it.descriptor !in importedScripts) {
-                    context.symbolTable.introduceValueParameter(it.owner.thisReceiver)
-                }
-            }
-
            val startOffset = ktScript.pureStartOffset
            val endOffset = ktScript.pureEndOffset

@@ -93,6 +82,24 @@ class ScriptGenerator(declarationGenerator: DeclarationGenerator) : DeclarationG
            // TODO: implement implicit receiver parameters handlin properly
            var parametersIndex = 0

+            irScript.earlierScripts = context.extensions.getPreviousScripts()?.filter {
+                // TODO: probably unnecessary filtering
+                it.owner != irScript && it.descriptor !in importedScripts
+            }
+            irScript.earlierScripts?.forEach {
+                context.symbolTable.introduceValueParameter(it.owner.thisReceiver)
+            }
+            irScript.earlierScriptsParameter = irScript.earlierScripts?.takeIf { it.isNotEmpty() }?.let {
+                val baseType = context.irBuiltIns.anyType
+                val arrayType = baseType.toArrayOrPrimitiveArrayType(context.irBuiltIns)
+                context.irFactory.createValueParameter(
+                    UNDEFINED_OFFSET, UNDEFINED_OFFSET, IrDeclarationOrigin.SCRIPT_EARLIER_SCRIPTS, IrValueParameterSymbolImpl(),
+                    Name.special("<earlierScripts>"), parametersIndex++,
+                    arrayType, null,
+                    false, false, false, false
+                ).also { it.parent = irScript }
+            }
+
            irScript.explicitCallParameters = descriptor.explicitConstructorParameters.map { valueParameterDescriptor ->
                parametersIndex++
                valueParameterDescriptor.toIrValueParameter(startOffset, endOffset, IrDeclarationOrigin.SCRIPT_CALL_PARAMETER).also { it.parent = irScript }
@@ -111,14 +118,17 @@ class ScriptGenerator(declarationGenerator: DeclarationGenerator) : DeclarationG
                        startOffset, endOffset, IrDeclarationOrigin.SCRIPT_PROVIDED_PROPERTY, parameter, type
                    )
                    val irProperty =
-                        PropertyGenerator(declarationGenerator).generateSyntheticProperty(ktScript, providedProperty, valueParameter, generateSyntheticAccessors = true)
+                        PropertyGenerator(declarationGenerator).generateSyntheticProperty(
+                            ktScript,
+                            providedProperty,
+                            valueParameter,
+                            generateSyntheticAccessors = true
+                        )
                    irProperty.origin = IrDeclarationOrigin.SCRIPT_PROVIDED_PROPERTY
                    irScript.statements += irProperty
                    valueParameter to irProperty.symbol
                }

-            irScript.earlierScripts = previousScripts
-
            for (d in ktScript.declarations) {
                when (d) {
                    is KtScriptInitializer -> {
@@ -129,7 +139,11 @@ class ScriptGenerator(declarationGenerator: DeclarationGenerator) : DeclarationG
                        if (d == ktScript.declarations.last() && descriptor.resultValue != null) {
                            descriptor.resultValue!!.let { resultDescriptor ->
                                PropertyGenerator(declarationGenerator)
-                                    .generateSyntheticPropertyWithInitializer(ktScript, resultDescriptor, generateSyntheticAccessors = true) {
+                                    .generateSyntheticPropertyWithInitializer(
+                                        ktScript,
+                                        resultDescriptor,
+                                        generateSyntheticAccessors = true
+                                    ) {
                                        // TODO: check if this is a correct place to do it
                                        it.visibility = DescriptorVisibilities.PUBLIC
                                        irExpressionBody
@@ -42,6 +42,7 @@ interface IrDeclarationOrigin {

    object SCRIPT_CLASS : IrDeclarationOriginImpl("SCRIPT_CLASS")
    object SCRIPT_STATEMENT : IrDeclarationOriginImpl("SCRIPT_STATEMENT")
+    object SCRIPT_EARLIER_SCRIPTS : IrDeclarationOriginImpl("SCRIPT_EARLIER_SCRIPTS")
    object SCRIPT_CALL_PARAMETER : IrDeclarationOriginImpl("SCRIPT_CALL_PARAMETER")
    object SCRIPT_IMPLICIT_RECEIVER : IrDeclarationOriginImpl("SCRIPT_IMPLICIT_RECEIVER")
    object SCRIPT_PROVIDED_PROPERTY : IrDeclarationOriginImpl("SCRIPT_PROVIDED_PROPERTY")
@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.ir.declarations

 import org.jetbrains.kotlin.ir.expressions.IrStatementContainer
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrPropertySymbol
 import org.jetbrains.kotlin.ir.symbols.IrScriptSymbol
 import org.jetbrains.kotlin.ir.types.IrType
@@ -32,5 +33,9 @@ abstract class IrScript :

    abstract var resultProperty: IrPropertySymbol?

+    abstract var earlierScriptsParameter: IrValueParameter?
+
    abstract var earlierScripts: List<IrScriptSymbol>?
+
+    abstract var targetClass: IrClassSymbol?
 }
@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.ir.ObsoleteDescriptorBasedAPI
 import org.jetbrains.kotlin.ir.UNDEFINED_OFFSET
 import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.IrConstructorCall
+import org.jetbrains.kotlin.ir.symbols.IrClassSymbol
 import org.jetbrains.kotlin.ir.symbols.IrPropertySymbol
 import org.jetbrains.kotlin.ir.symbols.IrScriptSymbol
 import org.jetbrains.kotlin.ir.types.IrType
@@ -52,7 +53,9 @@ class IrScriptImpl(
    override lateinit var implicitReceiversParameters: List<IrValueParameter>
    override lateinit var providedProperties: List<Pair<IrValueParameter, IrPropertySymbol>>
    override var resultProperty: IrPropertySymbol? = null
+    override var earlierScriptsParameter: IrValueParameter? = null
    override var earlierScripts: List<IrScriptSymbol>? = null
+    override var targetClass: IrClassSymbol? = null

    @ObsoleteDescriptorBasedAPI
    override val descriptor: ScriptDescriptor
@@ -118,6 +118,8 @@ open class DeepCopyIrTreeWithSymbols(
        ).also { scriptCopy ->
            scriptCopy.thisReceiver = declaration.thisReceiver.transform()
            declaration.statements.mapTo(scriptCopy.statements) { it.transform() }
+            scriptCopy.earlierScripts = declaration.earlierScripts
+            scriptCopy.earlierScriptsParameter = declaration.earlierScriptsParameter
            scriptCopy.explicitCallParameters = declaration.explicitCallParameters.map { it.transform() }
            scriptCopy.implicitReceiversParameters = declaration.implicitReceiversParameters.map { it.transform() }
            scriptCopy.providedProperties = declaration.providedProperties.map { it.first.transform() to it.second }
@@ -8,7 +8,6 @@ package kotlin.script.experimental.jvmhost.repl
 import kotlinx.coroutines.runBlocking
 import org.jetbrains.kotlin.backend.common.push
 import org.jetbrains.kotlin.cli.common.repl.*
-import org.jetbrains.kotlin.cli.common.repl.ReplCompilerWithoutCheck
 import org.jetbrains.kotlin.scripting.compiler.plugin.impl.KJvmReplCompilerBase
 import org.jetbrains.kotlin.scripting.compiler.plugin.repl.JvmReplCompilerStageHistory
 import org.jetbrains.kotlin.scripting.compiler.plugin.repl.JvmReplCompilerState
@@ -78,7 +77,7 @@ class JvmReplCompiler(
                )
            }
            else -> {
-                val message = res.reports.joinToString("\n") { it.message }
+                val message = res.reports.joinToString("\n")
                if (res.isIncomplete()) {
                    ReplCompileResult.Incomplete(message)
                } else {
@@ -75,7 +75,7 @@ class LazyScriptDescriptor(
        resolveSession.trace.record(BindingContext.SCRIPT, scriptInfo.script, this)
    }

-    override fun getResultValue(): ReplResultPropertyDescriptor? {
+    private val _resultValue: () -> ReplResultPropertyDescriptor? = resolveSession.storageManager.createNullableLazyValue {
        val expression = scriptInfo.script
            .getChildOfType<KtBlockExpression>()
            ?.getChildrenOfType<KtScriptInitializer>()?.lastOrNull()
@@ -85,7 +85,7 @@ class LazyScriptDescriptor(
            resolveSession.trace.bindingContext.getType(it)
        }

-        return if (type != null && !type.isUnit() && !type.isNothing()) {
+        if (type != null && !type.isUnit() && !type.isNothing()) {
            resultFieldName()?.let {
                ReplResultPropertyDescriptor(
                    it,
@@ -98,6 +98,8 @@ class LazyScriptDescriptor(
        } else null
    }

+    override fun getResultValue(): ReplResultPropertyDescriptor? = _resultValue()
+
    fun resultFieldName(): Name? {
        // TODO: implement robust REPL/script selection
        val replSnippetId =
@@ -7,6 +7,10 @@ package org.jetbrains.kotlin.scripting.compiler.plugin.impl


 import com.intellij.openapi.Disposable
+import org.jetbrains.kotlin.backend.common.phaser.PhaseConfig
+import org.jetbrains.kotlin.backend.jvm.*
+import org.jetbrains.kotlin.backend.jvm.serialization.JvmIdSignatureDescriptor
+import org.jetbrains.kotlin.cli.common.CLIConfigurationKeys
 import org.jetbrains.kotlin.cli.common.environment.setIdeaIoUseFallback
 import org.jetbrains.kotlin.cli.common.messages.AnalyzerWithCompilerReport
 import org.jetbrains.kotlin.cli.common.messages.MessageCollector
@@ -16,13 +20,16 @@ import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.codegen.ClassBuilderFactories
 import org.jetbrains.kotlin.codegen.KotlinCodegenFacade
 import org.jetbrains.kotlin.codegen.state.GenerationState
+import org.jetbrains.kotlin.config.JVMConfigurationKeys
+import org.jetbrains.kotlin.config.languageVersionSettings
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
+import org.jetbrains.kotlin.idea.MainFunctionDetector
+import org.jetbrains.kotlin.ir.backend.jvm.serialization.JvmManglerDesc
+import org.jetbrains.kotlin.ir.declarations.impl.IrFactoryImpl
+import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
-import org.jetbrains.kotlin.scripting.compiler.plugin.repl.JvmReplCompilerStageHistory
-import org.jetbrains.kotlin.scripting.compiler.plugin.repl.JvmReplCompilerState
-import org.jetbrains.kotlin.scripting.compiler.plugin.repl.ReplCodeAnalyzerBase
-import org.jetbrains.kotlin.scripting.compiler.plugin.repl.ReplImplicitsExtensionsResolutionFilter
+import org.jetbrains.kotlin.scripting.compiler.plugin.repl.*
 import org.jetbrains.kotlin.scripting.definitions.ScriptDependenciesProvider
 import org.jetbrains.kotlin.scripting.resolve.skipExtensionsResolutionForImplicits
 import org.jetbrains.kotlin.scripting.resolve.skipExtensionsResolutionForImplicitsExceptInnermost
@@ -112,45 +119,35 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
                    )
                }

-                val no = scriptPriority.getAndIncrement()
+                val snippetNo = scriptPriority.getAndIncrement()

                val analysisResult =
-                    compilationState.analyzerEngine.analyzeReplLineWithImportedScripts(snippetKtFile, sourceFiles.drop(1), snippet, no)
+                    compilationState.analyzerEngine.analyzeReplLineWithImportedScripts(
+                        snippetKtFile,
+                        sourceFiles.drop(1),
+                        snippet,
+                        snippetNo
+                    )
                AnalyzerWithCompilerReport.reportDiagnostics(analysisResult.diagnostics, errorHolder)

                val scriptDescriptor = when (analysisResult) {
-                    is ReplCodeAnalyzerBase.ReplLineAnalysisResult.WithErrors -> return failure(
-                        messageCollector
-                    )
+                    is ReplCodeAnalyzerBase.ReplLineAnalysisResult.WithErrors -> return failure(messageCollector)
                    is ReplCodeAnalyzerBase.ReplLineAnalysisResult.Successful -> {
                        (analysisResult.scriptDescriptor as? ScriptDescriptor)
-                            ?: return failure(
-                                snippet,
-                                messageCollector,
-                                "Unexpected script descriptor type ${analysisResult.scriptDescriptor::class}"
-                            )
+                            ?: throw AssertionError("Unexpected script descriptor type ${analysisResult.scriptDescriptor::class}")
                    }
-                    else -> return failure(
-                        snippet,
-                        messageCollector,
-                        "Unexpected result ${analysisResult::class.java}"
-                    )
+                    else -> throw AssertionError("Unexpected result ${analysisResult::class.java}")
                }

-                val generationState = GenerationState.Builder(
-                    snippetKtFile.project,
-                    ClassBuilderFactories.BINARIES,
-                    compilationState.analyzerEngine.module,
-                    compilationState.analyzerEngine.trace.bindingContext,
-                    sourceFiles,
-                    compilationState.environment.configuration
-                ).build().apply {
-                    scriptSpecific.earlierScriptsForReplInterpreter = history.map { it.item }
-                    beforeCompile()
-                }
-                KotlinCodegenFacade.generatePackage(generationState, snippetKtFile.script!!.containingKtFile.packageFqName, sourceFiles)
+                val isIr = context.environment.configuration.getBoolean(JVMConfigurationKeys.IR)

-                history.push(LineId(no, 0, snippet.hashCode()), scriptDescriptor)
+                val generationState = if (isIr) {
+                    generateWithBackendIr(compilationState, snippetKtFile, sourceFiles)
+                } else {
+                    generateWithOldBackend(compilationState, snippetKtFile, sourceFiles)
+                }
+
+                history.push(LineId(snippetNo, 0, snippet.hashCode()), scriptDescriptor)

                val dependenciesProvider = ScriptDependenciesProvider.getInstance(context.environment.project)
                val compiledScript =
@@ -175,6 +172,55 @@ open class KJvmReplCompilerBase<AnalyzerT : ReplCodeAnalyzerBase> protected cons
            }
        }.last()

+    private fun generateWithOldBackend(
+        compilationState: ReplCompilationState<AnalyzerT>,
+        snippetKtFile: KtFile,
+        sourceFiles: List<KtFile>
+    ): GenerationState {
+        val generationState = GenerationState.Builder(
+            snippetKtFile.project,
+            ClassBuilderFactories.BINARIES,
+            compilationState.analyzerEngine.module,
+            compilationState.analyzerEngine.trace.bindingContext,
+            sourceFiles,
+            compilationState.environment.configuration
+        ).build().also { generationState ->
+            generationState.scriptSpecific.earlierScriptsForReplInterpreter = history.map { it.item }
+            generationState.beforeCompile()
+        }
+        KotlinCodegenFacade.generatePackage(generationState, snippetKtFile.script!!.containingKtFile.packageFqName, sourceFiles)
+
+        return generationState
+    }
+
+    private fun generateWithBackendIr(
+        compilationState: ReplCompilationState<AnalyzerT>,
+        snippetKtFile: KtFile,
+        sourceFiles: List<KtFile>
+    ): GenerationState {
+        val generatorExtensions = object : JvmGeneratorExtensionsImpl() {
+            override fun getPreviousScripts() = history.map { compilationState.symbolTable.referenceScript(it.item) }
+        }
+        val codegenFactory = JvmIrCodegenFactory(
+            compilationState.environment.configuration.get(CLIConfigurationKeys.PHASE_CONFIG) ?: PhaseConfig(jvmPhases),
+            compilationState.mangler, compilationState.symbolTable, generatorExtensions
+        )
+        val generationState = GenerationState.Builder(
+            snippetKtFile.project,
+            ClassBuilderFactories.BINARIES,
+            compilationState.analyzerEngine.module,
+            compilationState.analyzerEngine.trace.bindingContext,
+            sourceFiles,
+            compilationState.environment.configuration
+        )
+            .codegenFactory(codegenFactory)
+            .build()
+
+        generationState.codegenFactory.generateModule(generationState, generationState.files)
+
+        return generationState
+    }
+
    override suspend fun invoke(
        script: SourceCode,
        scriptCompilationConfiguration: ScriptCompilationConfiguration
@@ -294,4 +340,15 @@ class ReplCompilationState<AnalyzerT : ReplCodeAnalyzerBase>(
        // ReplCodeAnalyzer1(context.environment)
        analyzerInit(context, implicitsResolutionFilter)
    }
+
+    private val manglerAndSymbolTable by lazy {
+        val mangler = JvmManglerDesc(
+            MainFunctionDetector(analyzerEngine.trace.bindingContext, environment.configuration.languageVersionSettings)
+        )
+        val symbolTable = SymbolTable(JvmIdSignatureDescriptor(mangler), IrFactoryImpl, JvmNameProvider)
+        mangler to symbolTable
+    }
+
+    override val mangler: JvmManglerDesc get() = manglerAndSymbolTable.first
+    override val symbolTable: SymbolTable get() = manglerAndSymbolTable.second
 }
@@ -7,13 +7,16 @@ package org.jetbrains.kotlin.scripting.compiler.plugin.repl

 import com.intellij.openapi.Disposable
 import com.intellij.openapi.util.Disposer
-import org.jetbrains.kotlin.cli.common.repl.*
+import org.jetbrains.kotlin.cli.common.repl.BasicReplStageHistory
+import org.jetbrains.kotlin.cli.common.repl.IReplStageState
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
 import org.jetbrains.kotlin.descriptors.ScriptDescriptor
+import org.jetbrains.kotlin.ir.backend.jvm.serialization.JvmManglerDesc
+import org.jetbrains.kotlin.ir.util.SymbolTable
 import org.jetbrains.kotlin.resolve.calls.tower.ImplicitsExtensionsResolutionFilter
 import java.util.concurrent.locks.ReentrantReadWriteLock
 import kotlin.concurrent.write
-import kotlin.script.experimental.api.*
+import kotlin.script.experimental.api.ScriptCompilationConfiguration

 class JvmReplCompilerStageHistory<CompilationT : JvmReplCompilerState.Compilation>(private val state: JvmReplCompilerState<CompilationT>) :
    BasicReplStageHistory<ScriptDescriptor>(state.lock)
@@ -62,5 +65,7 @@ class JvmReplCompilerState<CompilationT : JvmReplCompilerState.Compilation>(
        val environment: KotlinCoreEnvironment
        val analyzerEngine: ReplCodeAnalyzerBase
        val implicitsResolutionFilter: ImplicitsExtensionsResolutionFilter
+        val mangler: JvmManglerDesc
+        val symbolTable: SymbolTable
    }
 }