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 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): + 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| + Element 9 + Context receivers: + FirScriptSymbol context(