From 8f2aaf62c2398591d60e69a8c5e619af08a546a1 Mon Sep 17 00:00:00 2001
From: Natalia Selezneva <natalia.selezneva@jetbrains.com>
Date: Tue, 13 Aug 2019 16:17:56 +0300
Subject: [PATCH] Implement modification stamp for ScriptInitializer

This should influence ResolveElementCache to correctly update caches
because it rely on modifications stamps of the declaration
Before this change file.modificationStamp was used for script initializers in file, so they all became outdated after some change in file
---
 .../kotlin/psi/KtAnonymousInitializer.kt      | 12 +++++++
 .../src/org/jetbrains/kotlin/psi/KtElement.kt |  1 +
 .../kotlin/idea/ResolveElementCacheTest.kt    | 34 +++++++++++++++++++
 3 files changed, 47 insertions(+)

diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/KtAnonymousInitializer.kt b/compiler/psi/src/org/jetbrains/kotlin/psi/KtAnonymousInitializer.kt
index 34c33c9fce3..c67c0c50028 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/KtAnonymousInitializer.kt
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/KtAnonymousInitializer.kt
@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.stubs.KotlinPlaceHolderStub
 import org.jetbrains.kotlin.psi.stubs.elements.KtStubElementTypes
 import org.jetbrains.kotlin.utils.sure
+import java.util.concurrent.atomic.AtomicLong
 
 interface KtAnonymousInitializer : KtDeclaration, KtStatementExpression {
     val containingDeclaration: KtDeclaration
@@ -57,4 +58,15 @@ class KtScriptInitializer(node: ASTNode) : KtDeclarationImpl(node), KtAnonymousI
         get() = getParentOfType<KtScript>(true).sure { "Should only be present in script" }
 
     override fun <R, D> accept(visitor: KtVisitor<R, D>, data: D) = visitor.visitScriptInitializer(this, data)
+
+    private val modificationStamp = AtomicLong()
+
+    override fun subtreeChanged() {
+        super.subtreeChanged()
+        modificationStamp.getAndIncrement()
+    }
+
+    fun getModificationStamp(): Long {
+        return modificationStamp.get()
+    }
 }
\ No newline at end of file
diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/KtElement.kt b/compiler/psi/src/org/jetbrains/kotlin/psi/KtElement.kt
index 2b8f846d8b1..b95358fcf5e 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/KtElement.kt
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/KtElement.kt
@@ -33,6 +33,7 @@ fun KtElement.getModificationStamp(): Long {
         is KtFile -> this.modificationStamp
         is KtDeclarationStub<*> -> this.modificationStamp
         is KtSuperTypeList -> this.modificationStamp
+        is KtScriptInitializer -> this.getModificationStamp()
         else -> (parent as KtElement).getModificationStamp()
     }
 }
diff --git a/idea/tests/org/jetbrains/kotlin/idea/ResolveElementCacheTest.kt b/idea/tests/org/jetbrains/kotlin/idea/ResolveElementCacheTest.kt
index ef4e4c74f5a..04c602c8631 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/ResolveElementCacheTest.kt
+++ b/idea/tests/org/jetbrains/kotlin/idea/ResolveElementCacheTest.kt
@@ -533,6 +533,40 @@ class C(param1: String = "", param2: Int = 0) {
         assert(!parameterDescriptor.type.containsError())
     }
 
+    fun testTypingInScriptInitializer() {
+        val file = myFixture.configureByText(
+            "Test.kts", """
+            run { 1 +<caret> 1 }
+            run { 2 + 2 }
+            """
+        ) as KtFile
+
+        val script = file.script ?: error("File should be a script")
+
+        val statement1 = (script.blockExpression.statements.first() as? KtScriptInitializer)?.body
+            ?: error("Cannot find first expression in script")
+        val statement2 = (script.blockExpression.statements.last() as? KtScriptInitializer)?.body
+            ?: error("Cannot find last expression in script")
+        val bindingContext1Before = statement1.analyze()
+        val bindingContext2Before = statement2.analyze()
+
+        val caret = myFixture.elementByOffset.getParentOfType<KtExpression>(true) ?: error("Cannot find element at caret")
+
+        myFixture.project.executeWriteCommand("") {
+            caret.parent.addAfter(KtPsiFactory(project).createWhiteSpace(), caret)
+        }
+
+        val bindingContext1After = statement1.analyze()
+        val bindingContext2After = statement2.analyze()
+
+        assert(bindingContext1Before !== bindingContext1After) {
+            "Analysis for first statement must change because statement was changed"
+        }
+        assert(bindingContext2Before === bindingContext2After) {
+            "Analysis for second statement must not change, only first statement was changed"
+        }
+    }
+
     private fun checkResolveMultiple(mode: BodyResolveMode, vararg expressions: KtExpression): BindingContext {
         val resolutionFacade = expressions.first().getResolutionFacade()
         val bindingContext = resolutionFacade.analyze(expressions.asList(), mode)