From 1a735b08429f44ed648f146c2739add9b82ace45 Mon Sep 17 00:00:00 2001
From: Ilya Chernikov <ilya.chernikov@jetbrains.com>
Date: Wed, 7 Jun 2023 16:45:12 +0200
Subject: [PATCH] K2 scripting: support script-level typealiases

---
 .../kotlin/fir/analysis/checkers/FirHelpers.kt       |  2 +-
 .../kotlin/fir/backend/Fir2IrClassifierStorage.kt    |  6 ++++--
 .../jetbrains/kotlin/fir/backend/Fir2IrConverter.kt  | 12 ++++++++----
 .../testData/codegen/script/topLevelTypealias.kts    |  1 -
 .../tests/script/typealiasInScript.fir.kts           |  9 ---------
 .../diagnostics/tests/script/typealiasInScript.kts   |  1 +
 6 files changed, 14 insertions(+), 17 deletions(-)
 delete mode 100644 compiler/testData/diagnostics/tests/script/typealiasInScript.fir.kts

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirHelpers.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirHelpers.kt
index c19ef3ee8a5..77f01c2a2e6 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirHelpers.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirHelpers.kt
@@ -741,7 +741,7 @@ val CheckerContext.closestNonLocal get() = containingDeclarations.takeWhile { it
 fun CheckerContext.closestNonLocalWith(declaration: FirDeclaration) =
     (containingDeclarations + declaration).takeWhile { it.isNonLocal }.lastOrNull()
 
-val CheckerContext.isTopLevel get() = containingDeclarations.lastOrNull() is FirFile
+val CheckerContext.isTopLevel get() = containingDeclarations.lastOrNull().let { it is FirFile || it is FirScript }
 
 fun FirBasedSymbol<*>.hasAnnotationOrInsideAnnotatedClass(classId: ClassId, session: FirSession): Boolean {
     if (hasAnnotation(classId, session)) return true
diff --git a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrClassifierStorage.kt b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrClassifierStorage.kt
index cb288556c71..c1af1d9bd56 100644
--- a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrClassifierStorage.kt
+++ b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrClassifierStorage.kt
@@ -304,7 +304,7 @@ class Fir2IrClassifierStorage(
 
     fun registerTypeAlias(
         typeAlias: FirTypeAlias,
-        parent: IrFile
+        parent: IrDeclarationParent
     ): IrTypeAlias {
         val signature = signatureComposer.composeSignature(typeAlias)
         return typeAlias.convertWithOffsets { startOffset, endOffset ->
@@ -322,7 +322,9 @@ class Fir2IrClassifierStorage(
                 ).apply {
                     this.parent = parent
                     setTypeParameters(typeAlias)
-                    parent.declarations += this
+                    if (parent is IrFile) {
+                        parent.declarations += this
+                    }
                 }
                 typeAliasCache[typeAlias] = irTypeAlias
                 irTypeAlias
diff --git a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConverter.kt b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConverter.kt
index 71e3678540f..ce825839bb6 100644
--- a/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConverter.kt
+++ b/compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrConverter.kt
@@ -359,15 +359,19 @@ class Fir2IrConverter(
                 processRegularClassMembers(declaration)
             }
             is FirScript -> {
-                assert(parent is IrFile)
+                parent as IrFile
                 declarationStorage.getOrCreateIrScript(declaration).also { irScript ->
                     declarationStorage.enterScope(irScript)
                     irScript.parent = parent
                     for (scriptStatement in declaration.statements) {
                         if (scriptStatement is FirDeclaration) {
-                            if (scriptStatement is FirRegularClass) {
-                                registerClassAndNestedClasses(scriptStatement, irScript)
-                                processClassAndNestedClassHeaders(scriptStatement)
+                            when (scriptStatement) {
+                                is FirRegularClass -> {
+                                    registerClassAndNestedClasses(scriptStatement, irScript)
+                                    processClassAndNestedClassHeaders(scriptStatement)
+                                }
+                                is FirTypeAlias -> classifierStorage.registerTypeAlias(scriptStatement, irScript)
+                                else -> {}
                             }
 
                         }
diff --git a/compiler/testData/codegen/script/topLevelTypealias.kts b/compiler/testData/codegen/script/topLevelTypealias.kts
index eb9f554981b..0bfa4c10a2f 100644
--- a/compiler/testData/codegen/script/topLevelTypealias.kts
+++ b/compiler/testData/codegen/script/topLevelTypealias.kts
@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_K2: JVM_IR
 
 class SimpleClass(val s: String) {
     fun foo() = s
diff --git a/compiler/testData/diagnostics/tests/script/typealiasInScript.fir.kts b/compiler/testData/diagnostics/tests/script/typealiasInScript.fir.kts
deleted file mode 100644
index 06618e9ccf6..00000000000
--- a/compiler/testData/diagnostics/tests/script/typealiasInScript.fir.kts
+++ /dev/null
@@ -1,9 +0,0 @@
-<!TOPLEVEL_TYPEALIASES_ONLY!>typealias TopLevelInScript = String<!>
-
-class C {
-    <!TOPLEVEL_TYPEALIASES_ONLY!>typealias NestedInClass = String<!>
-}
-
-fun foo() {
-    <!TOPLEVEL_TYPEALIASES_ONLY!>typealias Local = String<!>
-}
\ No newline at end of file
diff --git a/compiler/testData/diagnostics/tests/script/typealiasInScript.kts b/compiler/testData/diagnostics/tests/script/typealiasInScript.kts
index 9607616d33b..63320f9d1f3 100644
--- a/compiler/testData/diagnostics/tests/script/typealiasInScript.kts
+++ b/compiler/testData/diagnostics/tests/script/typealiasInScript.kts
@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 typealias TopLevelInScript = String
 
 class C {