K2 REPL: Implement IR lowering for REPL snippets

#KT-59854 in-progress
JetBrains · Jan 10, 2025 · b539c8d · b539c8d
1 parent f6b42bb
commit b539c8d
Show file tree

Hide file tree

Showing 8 changed files with 480 additions and 18 deletions.
diff --git a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/utils/OriginUtils.kt b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/utils/OriginUtils.kt
@@ -45,6 +45,7 @@ fun FirClass.irOrigin(c: Fir2IrComponents): IrDeclarationOrigin = when {
     c.firProvider.getFirClassifierContainerFileIfAny(symbol) != null -> IrDeclarationOrigin.DEFINED
     else -> when (val origin = origin) {
         is FirDeclarationOrigin.Plugin -> GeneratedByPlugin(origin.key)
+        is FirDeclarationOrigin.FromOtherReplSnippet -> IrDeclarationOrigin.REPL_FROM_OTHER_SNIPPET
         else -> IrDeclarationOrigin.IR_EXTERNAL_DECLARATION_STUB
     }
 }

diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ClosureAnnotator.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ClosureAnnotator.kt
@@ -324,10 +324,13 @@ class ClosureAnnotator(irElement: IrElement, declaration: IrDeclaration) {
         private fun processScriptCapturing(receiverExpression: IrExpression?, declaration: IrDeclaration, data: ClosureBuilder?) {
             if (receiverExpression == null) {
                 val parent = declaration.parent
-                if (parent is IrScript) {
-                    data?.seeVariable(parent.thisReceiver!!.symbol)
-                } else if (parent is IrClass && parent.origin == IrDeclarationOrigin.SCRIPT_CLASS) {
-                    data?.seeVariable(parent.thisReceiver!!.symbol)
+                when (parent) {
+                    is IrScript -> {
+                        data?.seeVariable(parent.thisReceiver!!.symbol)
+                    }
+                    is IrClass if (parent.origin == IrDeclarationOrigin.SCRIPT_CLASS || parent.origin == IrDeclarationOrigin.REPL_FROM_OTHER_SNIPPET) -> {
+                        data?.seeVariable(parent.thisReceiver!!.symbol)
+                    }
                 }
             }
         }

diff --git a/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/IrDeclarationOrigin.kt b/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/declarations/IrDeclarationOrigin.kt
@@ -49,6 +49,8 @@ interface IrDeclarationOrigin {
         val SCRIPT_IMPLICIT_RECEIVER by IrDeclarationOriginImpl
         val SCRIPT_PROVIDED_PROPERTY by IrDeclarationOriginImpl
         val SCRIPT_RESULT_PROPERTY by IrDeclarationOriginImpl
+        val REPL_SNIPPET_CLASS by IrDeclarationOriginImpl
+        val REPL_FROM_OTHER_SNIPPET by IrDeclarationOriginImpl
         val GENERATED_DATA_CLASS_MEMBER by IrDeclarationOriginImpl
         val GENERATED_SINGLE_FIELD_VALUE_CLASS_MEMBER by IrDeclarationOriginImpl
         val GENERATED_MULTI_FIELD_VALUE_CLASS_MEMBER by IrDeclarationOriginImpl

diff --git a/...er/src/org/jetbrains/kotlin/scripting/compiler/plugin/extensions/ReplLoweringExtension.kt b/...er/src/org/jetbrains/kotlin/scripting/compiler/plugin/extensions/ReplLoweringExtension.kt
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2010-2024 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.scripting.compiler.plugin.extensions
+
+import org.jetbrains.kotlin.backend.common.extensions.IrGenerationExtension
+import org.jetbrains.kotlin.backend.common.extensions.IrPluginContext
+import org.jetbrains.kotlin.ir.declarations.IrModuleFragment
+import org.jetbrains.kotlin.platform.jvm.isJvm
+import org.jetbrains.kotlin.scripting.compiler.plugin.irLowerings.ReplSnippetsToClassesLowering
+
+class ReplLoweringExtension : IrGenerationExtension {
+    override fun generate(
+        moduleFragment: IrModuleFragment,
+        pluginContext: IrPluginContext,
+    ) {
+        val platform = pluginContext.platform
+        when {
+            platform.isJvm() -> {
+                ReplSnippetsToClassesLowering(pluginContext).lower(moduleFragment)
+            }
+        }
+    }
+}