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 7a3df552703..ca31f2023af 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 @@ -72,15 +72,26 @@ object LowLevelFirApiFacadeForResolveOnAir { fake.originalDeclaration = original } + /** + * We assume that [targetDeclaration] should have the same number of declarations + * + * @see restoreOriginalDeclarationsInScript + */ 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 - } + val originalStatements = originalDeclaration.declarationSequence() + val newStatements = targetDeclaration.declarationSequence() + originalStatements.zip(newStatements) { original, fake -> + fake.originalDeclaration = original } } } + private fun KtScript.declarationSequence(): Sequence { + val sequence = blockExpression.statements.asSequence().filter { it !is KtScriptInitializer && it is KtDeclaration } + @Suppress("UNCHECKED_CAST") + return sequence as Sequence + } + fun onAirResolveElement( firResolveSession: LLFirResolveSession, place: T, @@ -338,11 +349,7 @@ object LowLevelFirApiFacadeForResolveOnAir { * 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) + restoreOriginalDeclarationsInScript(originalScript = originalFirDeclaration, newScript = newFirDeclaration) } session.moduleComponents.firModuleLazyDeclarationResolver.runLazyDesignatedOnAirResolve( @@ -354,6 +361,59 @@ object LowLevelFirApiFacadeForResolveOnAir { return newFirDeclaration } + /** + * We assume that [newScript] has the same declarations as [originalScript] + */ + private fun restoreOriginalDeclarationsInScript(originalScript: FirScript, newScript: FirScript) { + val updatedStatements = ArrayList(newScript.statements.size) + val originalDeclarations = originalScript.statements.iterator() + for (recreatedStatement in newScript.statements) { + updatedStatements += if (recreatedStatement.isScriptStatement) { + recreatedStatement + } else { + originalDeclarations.nextDeclaration() ?: scriptDeclarationInconsistencyError(originalScript, newScript) + } + } + + val nextDeclaration = originalDeclarations.nextDeclaration() + if (nextDeclaration != null) { + scriptDeclarationInconsistencyError(originalScript, newScript) + } + + newScript.replaceStatements(updatedStatements) + } + + private fun scriptDeclarationInconsistencyError(originalScript: FirScript, newScript: FirScript): Nothing { + val originalDeclarations = originalScript.statements.filterNot(FirStatement::isScriptStatement) + val newDeclarations = newScript.statements.filterNot(FirStatement::isScriptStatement) + errorWithAttachment("New script has ${if (newDeclarations.size > originalDeclarations.size) "more" else "less"} declarations") { + withFirEntry("originalScript", originalScript) + withFirEntry("newScript", newScript) + withEntryGroup("originalDeclarations") { + for ((index, declaration) in originalDeclarations.withIndex()) { + withFirEntry(index.toString(), declaration) + } + } + + withEntryGroup("newDeclarations") { + for ((index, declaration) in newDeclarations.withIndex()) { + withFirEntry(index.toString(), declaration) + } + } + } + } + + private fun Iterator.nextDeclaration(): FirStatement? { + while (hasNext()) { + val statement = next() + if (statement.isScriptStatement) continue + + return statement + } + + return null + } + private fun computeDesignation(originalDeclaration: KtElement, originalFirFile: FirFile): FirElementFinder.FirDeclarationDesignation? { if (originalDeclaration is KtCodeFragment) { val firCodeFragment = originalFirFile.codeFragment diff --git a/analysis/low-level-api-fir/testdata/dependentCopy/moreStatementsInCopy.context.txt b/analysis/low-level-api-fir/testdata/dependentCopy/moreStatementsInCopy.context.txt index f9b1591c42f..302d1281cb3 100644 --- a/analysis/low-level-api-fir/testdata/dependentCopy/moreStatementsInCopy.context.txt +++ b/analysis/low-level-api-fir/testdata/dependentCopy/moreStatementsInCopy.context.txt @@ -16,11 +16,12 @@ Tower Data Context: Element 7 Scope: FirScriptDeclarationsScope Functions - FirNamedFunctionSymbol public? final? fun foo(i: Int, action: ( (Int) -> Unit )): R|kotlin/Unit| - FirNamedFunctionSymbol public? final? fun foo(i: Int, action: ( (Int) -> Unit )): R|kotlin/Unit| - FirNamedFunctionSymbol public? final? fun scriptFunction(): - FirNamedFunctionSymbol public? final? fun unusedScriptFunction(p: String): - FirNamedFunctionSymbol public? final? fun unusedScriptFunction(p: String): + 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: R|kotlin/String|): R|kotlin/Int| + Properties: + FirPropertySymbol public final val $$result: R|kotlin/Int| + public get(): R|kotlin/Int| Element 8 Scope: FirLocalScope Properties: @@ -31,10 +32,13 @@ Tower Data Context: SCRIPT: lval args: R|kotlin/Array| - public? final? fun scriptFunction(): - public? final? fun unusedScriptFunction(p: String): - public? final? fun unusedScriptFunction(p: String): - public? final? fun foo(i: Int, action: ( (Int) -> Unit )): R|kotlin/Unit| - public? final? fun foo(i: Int, action: ( (Int) -> Unit )): R|kotlin/Unit| + public final fun scriptFunction(): R|kotlin/Int| + R|/scriptFunction|() + public final fun unusedScriptFunction(p: R|kotlin/String|): R|kotlin/Int| + R|/scriptFunction|() + public final fun foo(i: R|kotlin/Int|, action: R|(kotlin/Int) -> kotlin/Unit|): R|kotlin/Unit| + R|/foo|(R|/scriptFunction|(), = foo@fun (it: R|kotlin/Int|): R|kotlin/Unit| ) + public final val $$result: R|kotlin/Int| + public get(): R|kotlin/Int| Type: kotlin/script/templates/standard/ScriptTemplateWithArgs Label: