From 0bc3fc9203e8a32da4f55657bab12a95481bcc63 Mon Sep 17 00:00:00 2001
From: Andrei Klunnyi <andrei.klunnyi@jetbrains.com>
Date: Mon, 10 Jul 2023 15:36:19 +0200
Subject: [PATCH] KT-59801 [FIR] scripting: explicit import breaks reference
 shortening

FIR-tree for scripts is built either for KtScript or KtFile depending on
the file structure. The presence of import statements results in KtFile.
There are lazy trees (some nodes are stubs) that are not yet supported
for scripts. We build such ones for KtFile case.

This commit prevents lazy tree construction for scripts at references
shortening.

^KT-59801 fixed
---
 .../api/fir/components/KtFirReferenceShortener.kt   | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirReferenceShortener.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirReferenceShortener.kt
index 2e1e56502ff..31b4e053626 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirReferenceShortener.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/components/KtFirReferenceShortener.kt
@@ -94,14 +94,15 @@ internal class KtFirReferenceShortener(
             kDocQualifiersToShorten = emptyList(),
         )
 
-        val towerContext = when (declarationToVisit) {
-            is KtFile -> {
-                LowLevelFirApiFacadeForResolveOnAir.getOnAirTowerDataContextProviderForTheWholeFile(firResolveSession, declarationToVisit)
-            }
-            else -> {
+        val towerContext =
+            if (declarationToVisit !is KtFile || declarationToVisit.isScript()) {
                 LowLevelFirApiFacadeForResolveOnAir.getOnAirGetTowerContextProvider(firResolveSession, declarationToVisit)
+            } else {
+                LowLevelFirApiFacadeForResolveOnAir.getOnAirTowerDataContextProviderForTheWholeFile(
+                    firResolveSession,
+                    declarationToVisit
+                )
             }
-        }
         //TODO: collect all usages of available symbols in the file and prevent importing symbols that could introduce name clashes, which
         // may alter the meaning of existing code.
         val collector = ElementsToShortenCollector(