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

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()
+    }
 }

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()
     }
 }

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)