From 3947ba57daba38b13ead244bbc8a82cbe1117546 Mon Sep 17 00:00:00 2001
From: Dmitrii Gridin <dmitry.gridin@jetbrains.com>
Date: Wed, 9 Aug 2023 14:27:34 +0200
Subject: [PATCH] [LL FIR] support on-air resolve for scripts

We should reanalyze only statements and leave declarations as they are

^KT-61013 Fixed
---
 .../LowLevelFirApiFacadeForResolveOnAir.kt    | 22 +++++++-
 .../testdata/dependentCopy/script.context.txt | 45 ++++++++++++++++
 .../testdata/dependentCopy/script.fir.txt     | 35 +++++++++++++
 .../testdata/dependentCopy/script.kts         | 18 +++++++
 .../dependentCopy/scriptFunction.context.txt  | 51 +++++++++++++++++++
 .../dependentCopy/scriptFunction.fir.txt      |  8 +++
 .../testdata/dependentCopy/scriptFunction.kts | 16 ++++++
 .../scriptInsideLastStatement.context.txt     | 51 +++++++++++++++++++
 .../scriptInsideLastStatement.fir.txt         |  6 +++
 .../scriptInsideLastStatement.kts             | 16 ++++++
 .../scriptLastStatement.context.txt           | 45 ++++++++++++++++
 .../dependentCopy/scriptLastStatement.fir.txt |  6 +++
 .../dependentCopy/scriptLastStatement.kts     | 16 ++++++
 .../scriptLastStatementCall.context.txt       | 45 ++++++++++++++++
 .../scriptLastStatementCall.fir.txt           | 10 ++++
 .../dependentCopy/scriptLastStatementCall.kts | 17 +++++++
 .../dependentCopy/scriptStatement.context.txt | 45 ++++++++++++++++
 .../dependentCopy/scriptStatement.fir.txt     |  6 +++
 .../dependentCopy/scriptStatement.kts         | 16 ++++++
 .../scriptStatementCall.context.txt           | 45 ++++++++++++++++
 .../dependentCopy/scriptStatementCall.fir.txt |  6 +++
 .../dependentCopy/scriptStatementCall.kts     | 17 +++++++
 ...riptDependentCopyContextTestGenerated.java | 42 +++++++++++++++
 .../ScriptDependentCopyFirTestGenerated.java  | 42 +++++++++++++++
 24 files changed, 625 insertions(+), 1 deletion(-)
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/script.context.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/script.fir.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/script.kts
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.context.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.fir.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.kts
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.context.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.fir.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.kts
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.context.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.fir.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.kts
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.context.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.fir.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.kts
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.context.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.fir.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.kts
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.context.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.fir.txt
 create mode 100644 analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.kts

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
index da7d59fab44..7a938b75cb3 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
@@ -13,7 +13,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerC
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerDataContextAllElementsCollector
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.canBePartOfParentDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.FirElementsRecorder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.KtToFirMapping
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.RawFirNonLocalDeclarationBuilder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.RawFirReplacement
@@ -25,6 +24,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.state.LLFirResolvableReso
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.FirElementFinder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.isScriptStatement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.originalDeclaration
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
 import org.jetbrains.kotlin.analysis.utils.printer.parentOfType
@@ -33,6 +33,7 @@ import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.builder.buildFileAnnotationsContainer
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirAnnotation
+import org.jetbrains.kotlin.fir.expressions.FirStatement
 import org.jetbrains.kotlin.fir.realPsi
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.FirResolveContextCollector
@@ -69,6 +70,14 @@ object LowLevelFirApiFacadeForResolveOnAir {
         originalDeclarationParents.zip(fakeDeclarationParents) { original, fake ->
             fake.originalDeclaration = original
         }
+
+        if (targetDeclaration is KtScript && originalDeclaration is KtScript) {
+            originalDeclaration.blockExpression.statements.zip(targetDeclaration.blockExpression.statements) { original, fake ->
+                if (original is KtDeclaration && fake is KtDeclaration) {
+                    fake.originalDeclaration = original
+                }
+            }
+        }
     }
 
     fun <T : KtElement> onAirResolveElement(
@@ -322,6 +331,17 @@ object LowLevelFirApiFacadeForResolveOnAir {
             replacement = replacement,
         )
 
+        /**
+         * We shouldn't touch script declarations because they're independent of a script
+         */
+        if (newFirDeclaration is FirScript && originalFirDeclaration is FirScript) {
+            val newStatements = originalFirDeclaration.statements.zip(newFirDeclaration.statements).map { (original, copied) ->
+                original.takeUnless(FirStatement::isScriptStatement) ?: copied
+            }
+
+            newFirDeclaration.replaceStatements(newStatements)
+        }
+
         session.moduleComponents.firModuleLazyDeclarationResolver.runLazyDesignatedOnAirResolve(
             FirDesignationWithFile(originalDesignationPath.path, newFirDeclaration, originalFirFile),
             collector,
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/script.context.txt b/analysis/low-level-api-fir/testdata/dependentCopy/script.context.txt
new file mode 100644
index 00000000000..6e3739d6c9f
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/script.context.txt
@@ -0,0 +1,45 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirScriptDeclarationsScope
+            Classifiers:
+                FirRegularClassSymbol public? final? class ScriptClass : R|kotlin/Any|
+            Functions
+                FirNamedFunctionSymbol public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                FirNamedFunctionSymbol public final fun scriptFunction(): R|kotlin/Int|
+                FirNamedFunctionSymbol public? final? fun unusedScriptFunction(p: String): <implicit>
+            Properties:
+                FirPropertySymbol public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+    Element 8
+        Scope: FirLocalScope
+            Properties:
+                FirPropertySymbol lval args: R|kotlin/Array<kotlin/String>|
+    Element 9
+        Context receivers:
+            FirScriptSymbol context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+            SCRIPT:
+                lval args: R|kotlin/Array<kotlin/String>|
+
+                public? final? class ScriptClass : R|kotlin/Any|
+                public final fun scriptFunction(): R|kotlin/Int|
+                public? final? fun unusedScriptFunction(p: String): <implicit>
+                R|/scriptFunction|()
+                public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+                Type: kotlin/script/templates/standard/ScriptTemplateWithArgs
+                Label: <script>
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/script.fir.txt b/analysis/low-level-api-fir/testdata/dependentCopy/script.fir.txt
new file mode 100644
index 00000000000..67c1ea885aa
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/script.fir.txt
@@ -0,0 +1,35 @@
+KT element: KtScript
+FIR element: FirScriptImpl
+FIR source kind: KtRealSourceElementKind
+
+FIR element rendered:
+context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+SCRIPT: [ResolvedTo(BODY_RESOLVE)]
+    [ResolvedTo(BODY_RESOLVE)] lval args: R|kotlin/Array<kotlin/String>|
+
+    public? final? [ResolvedTo(RAW_FIR)] class ScriptClass : R|kotlin/Any| {
+        public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=ScriptClass] constructor(): R|ScriptClass| {
+            LAZY_super<R|kotlin/Any|>
+        }
+
+        public? final? [ResolvedTo(RAW_FIR)] val memberProperty: <implicit> = LAZY_EXPRESSION
+            public? [ResolvedTo(RAW_FIR)] [ContainingClassKey=ScriptClass] get(): <implicit>
+
+    }
+
+    public final [ResolvedTo(IMPLICIT_TYPES_BODY_RESOLVE)] fun scriptFunction(): R|kotlin/Int| {
+        ^scriptFunction Int(42)
+    }
+
+    public? final? [ResolvedTo(RAW_FIR)] fun unusedScriptFunction([ResolvedTo(RAW_FIR)] p: String): <implicit> { LAZY_BLOCK }
+
+    R|/scriptFunction|()
+    public final [ResolvedTo(CONTRACTS)] fun foo([ResolvedTo(CONTRACTS)] i: R|kotlin/Int|, [ResolvedTo(CONTRACTS)] action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit| {
+        action#(i#)
+    }
+
+    public final [ResolvedTo(BODY_RESOLVE)] val $$result: R|kotlin/Unit| = R|/foo|(R|/scriptFunction|(), <L> = [ResolvedTo(BODY_RESOLVE)] [MatchingParameterFunctionTypeKey=kotlin/Function1<kotlin/Int, kotlin/Unit>] foo@fun <anonymous>([ResolvedTo(BODY_RESOLVE)] it: R|kotlin/Int|): R|kotlin/Unit| <inline=NoInline>  {
+        R|/scriptFunction|()
+    }
+    )
+        public [ResolvedTo(BODY_RESOLVE)] get(): R|kotlin/Unit|
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/script.kts b/analysis/low-level-api-fir/testdata/dependentCopy/script.kts
new file mode 100644
index 00000000000..e1054d5f740
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/script.kts
@@ -0,0 +1,18 @@
+// LOOK_UP_FOR_ELEMENT_OF_TYPE: KtScript
+
+<caret>class ScriptClass {
+    val memberProperty = 4
+}
+
+fun scriptFunction() = 42
+fun unusedScriptFunction(p: String) = 22
+
+scriptFunction()
+
+fun foo(i: Int, action: (Int) -> Unit) {
+    action(i)
+}
+
+foo(scriptFunction()) {
+    scriptFunction()
+}
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.context.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.context.txt
new file mode 100644
index 00000000000..361464f9138
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.context.txt
@@ -0,0 +1,51 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirScriptDeclarationsScope
+            Classifiers:
+                FirRegularClassSymbol public? final? class ScriptClass : R|kotlin/Any|
+            Functions
+                FirNamedFunctionSymbol public? final? fun foo(i: Int, action: ( (Int) -> Unit )): R|kotlin/Unit|
+                FirNamedFunctionSymbol public? final? fun scriptFunction(): <implicit>
+                FirNamedFunctionSymbol public? final? fun unusedScriptFunction(p: String): <implicit>
+            Properties:
+                FirPropertySymbol public final val $$result: <implicit>
+                    public get(): <implicit>
+    Element 8
+        Scope: FirLocalScope
+            Functions
+                FirNamedFunctionSymbol public final fun unusedScriptFunction(p: R|kotlin/String|): R|kotlin/Int|
+            Properties:
+                FirPropertySymbol lval args: R|kotlin/Array<kotlin/String>|
+    Element 9
+        Context receivers:
+            FirScriptSymbol context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+            SCRIPT:
+                lval args: R|kotlin/Array<kotlin/String>|
+
+                public? final? class ScriptClass : R|kotlin/Any|
+                public? final? fun scriptFunction(): <implicit>
+                public? final? fun unusedScriptFunction(p: String): <implicit>
+                scriptFunction#()
+                public? final? fun foo(i: Int, action: ( (Int) -> Unit )): R|kotlin/Unit|
+                public final val $$result: <implicit>
+                    public get(): <implicit>
+                Type: kotlin/script/templates/standard/ScriptTemplateWithArgs
+                Label: <script>
+    Element 10
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol p: R|kotlin/String|
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.fir.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.fir.txt
new file mode 100644
index 00000000000..c9bcede5958
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.fir.txt
@@ -0,0 +1,8 @@
+KT element: KtNamedFunction
+FIR element: FirSimpleFunctionImpl
+FIR source kind: KtRealSourceElementKind
+
+FIR element rendered:
+public final [ResolvedTo(BODY_RESOLVE)] fun unusedScriptFunction([ResolvedTo(BODY_RESOLVE)] p: R|kotlin/String|): R|kotlin/Int| {
+    ^unusedScriptFunction Int(22)
+}
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.kts b/analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.kts
new file mode 100644
index 00000000000..521d33425fc
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.kts
@@ -0,0 +1,16 @@
+class ScriptClass {
+    val memberProperty = 4
+}
+
+fun scriptFunction() = 42
+fun u<caret>nusedScriptFunction(p: String) = 22
+
+scriptFunction()
+
+fun foo(i: Int, action: (Int) -> Unit) {
+    action(i)
+}
+
+foo(scriptFunction()) {
+    scriptFunction()
+}
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.context.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.context.txt
new file mode 100644
index 00000000000..b6d7c71e27e
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.context.txt
@@ -0,0 +1,51 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirScriptDeclarationsScope
+            Classifiers:
+                FirRegularClassSymbol public? final? class ScriptClass : R|kotlin/Any|
+            Functions
+                FirNamedFunctionSymbol public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                FirNamedFunctionSymbol public final fun scriptFunction(): R|kotlin/Int|
+                FirNamedFunctionSymbol public? final? fun unusedScriptFunction(p: String): <implicit>
+            Properties:
+                FirPropertySymbol public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+    Element 8
+        Scope: FirLocalScope
+            Properties:
+                FirPropertySymbol lval args: R|kotlin/Array<kotlin/String>|
+    Element 9
+        Context receivers:
+            FirScriptSymbol context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+            SCRIPT:
+                lval args: R|kotlin/Array<kotlin/String>|
+
+                public? final? class ScriptClass : R|kotlin/Any|
+                public final fun scriptFunction(): R|kotlin/Int|
+                public? final? fun unusedScriptFunction(p: String): <implicit>
+                R|/scriptFunction|()
+                public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+                Type: kotlin/script/templates/standard/ScriptTemplateWithArgs
+                Label: <script>
+    Element 10
+        Scope: FirLocalScope
+            Properties:
+                FirValueParameterSymbol it: R|kotlin/Int|
+    Element 11
+        Scope: FirLocalScope
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.fir.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.fir.txt
new file mode 100644
index 00000000000..5a188241524
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.fir.txt
@@ -0,0 +1,6 @@
+KT element: KtNameReferenceExpression
+FIR element: FirResolvedNamedReferenceImpl
+FIR source kind: KtRealSourceElementKind
+
+FIR element rendered:
+R|/scriptFunction|
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.kts b/analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.kts
new file mode 100644
index 00000000000..487aefc56bc
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.kts
@@ -0,0 +1,16 @@
+class ScriptClass {
+    val memberProperty = 4
+}
+
+fun scriptFunction() = 42
+fun unusedScriptFunction(p: String) = 22
+
+scriptFunction()
+
+fun foo(i: Int, action: (Int) -> Unit) {
+    action(i)
+}
+
+foo(scriptFunction()) {
+    scri<caret>ptFunction()
+}
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.context.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.context.txt
new file mode 100644
index 00000000000..6e3739d6c9f
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.context.txt
@@ -0,0 +1,45 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirScriptDeclarationsScope
+            Classifiers:
+                FirRegularClassSymbol public? final? class ScriptClass : R|kotlin/Any|
+            Functions
+                FirNamedFunctionSymbol public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                FirNamedFunctionSymbol public final fun scriptFunction(): R|kotlin/Int|
+                FirNamedFunctionSymbol public? final? fun unusedScriptFunction(p: String): <implicit>
+            Properties:
+                FirPropertySymbol public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+    Element 8
+        Scope: FirLocalScope
+            Properties:
+                FirPropertySymbol lval args: R|kotlin/Array<kotlin/String>|
+    Element 9
+        Context receivers:
+            FirScriptSymbol context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+            SCRIPT:
+                lval args: R|kotlin/Array<kotlin/String>|
+
+                public? final? class ScriptClass : R|kotlin/Any|
+                public final fun scriptFunction(): R|kotlin/Int|
+                public? final? fun unusedScriptFunction(p: String): <implicit>
+                R|/scriptFunction|()
+                public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+                Type: kotlin/script/templates/standard/ScriptTemplateWithArgs
+                Label: <script>
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.fir.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.fir.txt
new file mode 100644
index 00000000000..40553415b7d
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.fir.txt
@@ -0,0 +1,6 @@
+KT element: KtNameReferenceExpression
+FIR element: FirResolvedNamedReferenceImpl
+FIR source kind: KtRealSourceElementKind
+
+FIR element rendered:
+R|/foo|
\ No newline at end of file
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.kts b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.kts
new file mode 100644
index 00000000000..526b1ab52b0
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.kts
@@ -0,0 +1,16 @@
+class ScriptClass {
+    val memberProperty = 4
+}
+
+fun scriptFunction() = 42
+fun unusedScriptFunction(p: String) = 22
+
+scriptFunction()
+
+fun foo(i: Int, action: (Int) -> Unit) {
+    action(i)
+}
+
+<caret>foo(scriptFunction()) {
+    scriptFunction()
+}
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.context.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.context.txt
new file mode 100644
index 00000000000..6e3739d6c9f
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.context.txt
@@ -0,0 +1,45 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirScriptDeclarationsScope
+            Classifiers:
+                FirRegularClassSymbol public? final? class ScriptClass : R|kotlin/Any|
+            Functions
+                FirNamedFunctionSymbol public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                FirNamedFunctionSymbol public final fun scriptFunction(): R|kotlin/Int|
+                FirNamedFunctionSymbol public? final? fun unusedScriptFunction(p: String): <implicit>
+            Properties:
+                FirPropertySymbol public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+    Element 8
+        Scope: FirLocalScope
+            Properties:
+                FirPropertySymbol lval args: R|kotlin/Array<kotlin/String>|
+    Element 9
+        Context receivers:
+            FirScriptSymbol context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+            SCRIPT:
+                lval args: R|kotlin/Array<kotlin/String>|
+
+                public? final? class ScriptClass : R|kotlin/Any|
+                public final fun scriptFunction(): R|kotlin/Int|
+                public? final? fun unusedScriptFunction(p: String): <implicit>
+                R|/scriptFunction|()
+                public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+                Type: kotlin/script/templates/standard/ScriptTemplateWithArgs
+                Label: <script>
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.fir.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.fir.txt
new file mode 100644
index 00000000000..018b0a69117
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.fir.txt
@@ -0,0 +1,10 @@
+KT element: KtCallExpression
+FIR element: FirPropertyImpl
+FIR source kind: KtRealSourceElementKind
+
+FIR element rendered:
+public final [ResolvedTo(BODY_RESOLVE)] val $$result: R|kotlin/Unit| = R|/foo|(R|/scriptFunction|(), <L> = [ResolvedTo(BODY_RESOLVE)] [MatchingParameterFunctionTypeKey=kotlin/Function1<kotlin/Int, kotlin/Unit>] foo@fun <anonymous>([ResolvedTo(BODY_RESOLVE)] it: R|kotlin/Int|): R|kotlin/Unit| <inline=NoInline>  {
+    R|/scriptFunction|()
+}
+)
+    public [ResolvedTo(BODY_RESOLVE)] get(): R|kotlin/Unit|
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.kts b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.kts
new file mode 100644
index 00000000000..02a20a12d56
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.kts
@@ -0,0 +1,17 @@
+// LOOK_UP_FOR_ELEMENT_OF_TYPE: KtCallExpression
+class ScriptClass {
+    val memberProperty = 4
+}
+
+fun scriptFunction() = 42
+fun unusedScriptFunction(p: String) = 22
+
+scriptFunction()
+
+fun foo(i: Int, action: (Int) -> Unit) {
+    action(i)
+}
+
+<caret>foo(scriptFunction()) {
+    scriptFunction()
+}
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.context.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.context.txt
new file mode 100644
index 00000000000..6e3739d6c9f
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.context.txt
@@ -0,0 +1,45 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirScriptDeclarationsScope
+            Classifiers:
+                FirRegularClassSymbol public? final? class ScriptClass : R|kotlin/Any|
+            Functions
+                FirNamedFunctionSymbol public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                FirNamedFunctionSymbol public final fun scriptFunction(): R|kotlin/Int|
+                FirNamedFunctionSymbol public? final? fun unusedScriptFunction(p: String): <implicit>
+            Properties:
+                FirPropertySymbol public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+    Element 8
+        Scope: FirLocalScope
+            Properties:
+                FirPropertySymbol lval args: R|kotlin/Array<kotlin/String>|
+    Element 9
+        Context receivers:
+            FirScriptSymbol context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+            SCRIPT:
+                lval args: R|kotlin/Array<kotlin/String>|
+
+                public? final? class ScriptClass : R|kotlin/Any|
+                public final fun scriptFunction(): R|kotlin/Int|
+                public? final? fun unusedScriptFunction(p: String): <implicit>
+                R|/scriptFunction|()
+                public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+                Type: kotlin/script/templates/standard/ScriptTemplateWithArgs
+                Label: <script>
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.fir.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.fir.txt
new file mode 100644
index 00000000000..5a188241524
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.fir.txt
@@ -0,0 +1,6 @@
+KT element: KtNameReferenceExpression
+FIR element: FirResolvedNamedReferenceImpl
+FIR source kind: KtRealSourceElementKind
+
+FIR element rendered:
+R|/scriptFunction|
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.kts b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.kts
new file mode 100644
index 00000000000..ed07a76789a
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.kts
@@ -0,0 +1,16 @@
+class ScriptClass {
+    val memberProperty = 4
+}
+
+fun scriptFunction() = 42
+fun unusedScriptFunction(p: String) = 22
+
+script<caret>Function()
+
+fun foo(i: Int, action: (Int) -> Unit) {
+    action(i)
+}
+
+foo(scriptFunction()) {
+    scriptFunction()
+}
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.context.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.context.txt
new file mode 100644
index 00000000000..6e3739d6c9f
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.context.txt
@@ -0,0 +1,45 @@
+Tower Data Context:
+    Element 0
+        Scope: FirDefaultStarImportingScope
+    Element 1
+        Scope: FirDefaultSimpleImportingScope
+    Element 2
+        Scope: FirExplicitStarImportingScope
+    Element 3
+        Scope: FirDefaultSimpleImportingScope
+    Element 4
+        Scope: FirDefaultSimpleImportingScope
+    Element 5
+        Scope: FirPackageMemberScope
+    Element 6
+        Scope: FirExplicitSimpleImportingScope
+    Element 7
+        Scope: FirScriptDeclarationsScope
+            Classifiers:
+                FirRegularClassSymbol public? final? class ScriptClass : R|kotlin/Any|
+            Functions
+                FirNamedFunctionSymbol public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                FirNamedFunctionSymbol public final fun scriptFunction(): R|kotlin/Int|
+                FirNamedFunctionSymbol public? final? fun unusedScriptFunction(p: String): <implicit>
+            Properties:
+                FirPropertySymbol public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+    Element 8
+        Scope: FirLocalScope
+            Properties:
+                FirPropertySymbol lval args: R|kotlin/Array<kotlin/String>|
+    Element 9
+        Context receivers:
+            FirScriptSymbol context(<script>@R|kotlin/script/templates/standard/ScriptTemplateWithArgs|)
+            SCRIPT:
+                lval args: R|kotlin/Array<kotlin/String>|
+
+                public? final? class ScriptClass : R|kotlin/Any|
+                public final fun scriptFunction(): R|kotlin/Int|
+                public? final? fun unusedScriptFunction(p: String): <implicit>
+                R|/scriptFunction|()
+                public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit|
+                public final val $$result: R|kotlin/Unit|
+                    public get(): R|kotlin/Unit|
+                Type: kotlin/script/templates/standard/ScriptTemplateWithArgs
+                Label: <script>
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.fir.txt b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.fir.txt
new file mode 100644
index 00000000000..ef9e0e17940
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.fir.txt
@@ -0,0 +1,6 @@
+KT element: KtCallExpression
+FIR element: FirFunctionCallImpl
+FIR source kind: KtRealSourceElementKind
+
+FIR element rendered:
+R|/scriptFunction|()
diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.kts b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.kts
new file mode 100644
index 00000000000..05d307f57a1
--- /dev/null
+++ b/analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.kts
@@ -0,0 +1,17 @@
+// LOOK_UP_FOR_ELEMENT_OF_TYPE: KtCallExpression
+class ScriptClass {
+    val memberProperty = 4
+}
+
+fun scriptFunction() = 42
+fun unusedScriptFunction(p: String) = 22
+
+script<caret>Function()
+
+fun foo(i: Int, action: (Int) -> Unit) {
+    action(i)
+}
+
+foo(scriptFunction()) {
+    scriptFunction()
+}
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptDependentCopyContextTestGenerated.java b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptDependentCopyContextTestGenerated.java
index 5863e9a70e6..7a07f02c790 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptDependentCopyContextTestGenerated.java
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptDependentCopyContextTestGenerated.java
@@ -35,4 +35,46 @@ public class ScriptDependentCopyContextTestGenerated extends AbstractScriptDepen
     public void testClassScript() throws Exception {
         runTest("analysis/low-level-api-fir/testdata/dependentCopy/classScript.kts");
     }
+
+    @Test
+    @TestMetadata("script.kts")
+    public void testScript() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/script.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptFunction.kts")
+    public void testScriptFunction() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptInsideLastStatement.kts")
+    public void testScriptInsideLastStatement() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptLastStatement.kts")
+    public void testScriptLastStatement() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptLastStatementCall.kts")
+    public void testScriptLastStatementCall() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptStatement.kts")
+    public void testScriptStatement() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptStatementCall.kts")
+    public void testScriptStatementCall() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.kts");
+    }
 }
diff --git a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptDependentCopyFirTestGenerated.java b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptDependentCopyFirTestGenerated.java
index ddb9f92c88a..d09ee4e329e 100644
--- a/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptDependentCopyFirTestGenerated.java
+++ b/analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/ScriptDependentCopyFirTestGenerated.java
@@ -35,4 +35,46 @@ public class ScriptDependentCopyFirTestGenerated extends AbstractScriptDependent
     public void testClassScript() throws Exception {
         runTest("analysis/low-level-api-fir/testdata/dependentCopy/classScript.kts");
     }
+
+    @Test
+    @TestMetadata("script.kts")
+    public void testScript() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/script.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptFunction.kts")
+    public void testScriptFunction() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptFunction.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptInsideLastStatement.kts")
+    public void testScriptInsideLastStatement() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptInsideLastStatement.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptLastStatement.kts")
+    public void testScriptLastStatement() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatement.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptLastStatementCall.kts")
+    public void testScriptLastStatementCall() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptLastStatementCall.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptStatement.kts")
+    public void testScriptStatement() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptStatement.kts");
+    }
+
+    @Test
+    @TestMetadata("scriptStatementCall.kts")
+    public void testScriptStatementCall() throws Exception {
+        runTest("analysis/low-level-api-fir/testdata/dependentCopy/scriptStatementCall.kts");
+    }
 }