FIR IDE: forbid executing write action inside analyse call

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSessionProvider.kt

@@ -5,12 +5,12 @@
 
 package org.jetbrains.kotlin.idea.frontend.api
 
+import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.service
 import com.intellij.openapi.progress.ProgressIndicator
 import com.intellij.openapi.progress.Task
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ReadActionConfinementValidityTokenFactory
-import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityToken
 import org.jetbrains.kotlin.idea.frontend.api.tokens.ValidityTokenFactory
 import org.jetbrains.kotlin.idea.util.application.runReadAction
 import org.jetbrains.kotlin.psi.KtElement
@@ -27,7 +27,11 @@ annotation class KtAnalysisSessionProviderInternals
  * Should not be used directly, consider using [analyse]/[analyseWithReadAction]/[analyseInModalWindow] instead
  */
 @InvalidWayOfUsingAnalysisSession
-abstract class KtAnalysisSessionProvider {
+abstract class KtAnalysisSessionProvider : Disposable {
+    @Suppress("LeakingThis")
+    @OptIn(KtInternalApiMarker::class)
+    val noWriteActionInAnalyseCallChecker = NoWriteActionInAnalyseCallChecker(this)
+
     @InvalidWayOfUsingAnalysisSession
     abstract fun getAnalysisSession(contextElement: KtElement, factory: ValidityTokenFactory): KtAnalysisSession
 
@@ -42,17 +46,20 @@ abstract class KtAnalysisSessionProvider {
     inline fun <R> analyse(contextElement: KtElement, tokenFactory: ValidityTokenFactory, action: KtAnalysisSession.() -> R): R =
         analyse(getAnalysisSession(contextElement, tokenFactory), tokenFactory, action)
 
-    @OptIn(KtAnalysisSessionProviderInternals::class)
+    @OptIn(KtAnalysisSessionProviderInternals::class, KtInternalApiMarker::class)
     @InvalidWayOfUsingAnalysisSession
     inline fun <R> analyse(analysisSession: KtAnalysisSession, factory: ValidityTokenFactory, action: KtAnalysisSession.() -> R): R {
+        noWriteActionInAnalyseCallChecker.beforeEnteringAnalysisContext()
         factory.beforeEnteringAnalysisContext()
         return try {
             analysisSession.action()
         } finally {
             factory.afterLeavingAnalysisContext()
+            noWriteActionInAnalyseCallChecker.afterLeavingAnalysisContext()
         }
     }
 
+    override fun dispose() {}
 }
 
 /**

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/NoWriteActionInAnalyseCallChecker.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.frontend.api
+
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationListener
+import com.intellij.openapi.application.ApplicationManager
+import org.jetbrains.kotlin.miniStdLib.multithreadings.javaThreadLocal
+
+@KtInternalApiMarker
+class NoWriteActionInAnalyseCallChecker(parentDisposable: Disposable) {
+    init {
+        val listener = object : ApplicationListener {
+            override fun writeActionFinished(action: Any) {
+                if (currentAnalysisContextEnteringCount > 0) {
+                    throw WriteActionStartInsideAnalysisContextException()
+                }
+            }
+        }
+        ApplicationManager.getApplication().addApplicationListener(listener, parentDisposable)
+    }
+
+    fun beforeEnteringAnalysisContext() {
+        currentAnalysisContextEnteringCount++
+    }
+
+    fun afterLeavingAnalysisContext() {
+        currentAnalysisContextEnteringCount--
+    }
+
+    private var currentAnalysisContextEnteringCount by javaThreadLocal(0)
+}
+
+class WriteActionStartInsideAnalysisContextException : IllegalStateException(
+    "write action should be never executed inside analysis context (e,g. analyse call)"
+)

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/annotations.kt

@@ -0,0 +1,9 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.frontend.api
+
+@RequiresOptIn
+annotation class KtInternalApiMarker

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSessionContractsTest.kt

@@ -0,0 +1,66 @@
+/*
+ * Copyright 2010-2021 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.frontend.api
+
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.testFramework.LightProjectDescriptor
+import junit.framework.Assert
+import org.jetbrains.kotlin.idea.frontend.api.tokens.AlwaysAccessibleValidityTokenFactory
+import org.jetbrains.kotlin.idea.frontend.api.tokens.HackToForceAllowRunningAnalyzeOnEDT
+import org.jetbrains.kotlin.idea.frontend.api.tokens.hackyAllowRunningOnEdt
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
+import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor
+import org.jetbrains.kotlin.idea.util.application.runWriteAction
+import org.jetbrains.kotlin.psi.KtFile
+
+class KtAnalysisSessionContractsTest : KotlinLightCodeInsightFixtureTestCase() {
+    override val captureExceptions: Boolean get() = false
+
+    fun testThatWriteActionCannotBeCalledInsideAnalyseCall() {
+        val fakeFile = createFakeKtFile()
+
+        var wasThrown = false
+        ApplicationManager.getApplication().invokeAndWait {
+            @OptIn(HackToForceAllowRunningAnalyzeOnEDT::class)
+            hackyAllowRunningOnEdt {
+                analyse(fakeFile) {
+                    try {
+                        runWriteAction { }
+                    } catch (_: WriteActionStartInsideAnalysisContextException) {
+                        wasThrown = true
+                    }
+                }
+            }
+        }
+        Assert.assertTrue("WriteActionStartInsideAnalysisContextException was not thrown", wasThrown)
+    }
+
+    fun testThatWriteActionCannotBeCalledInsideAnalyseCallWithCustomToken() {
+        val fakeFile = createFakeKtFile()
+
+        var wasThrown = false
+        ApplicationManager.getApplication().invokeAndWait {
+            @OptIn(HackToForceAllowRunningAnalyzeOnEDT::class)
+            hackyAllowRunningOnEdt {
+                analyseWithCustomToken(fakeFile, AlwaysAccessibleValidityTokenFactory) {
+                    try {
+                        runWriteAction { }
+                    } catch (_: WriteActionStartInsideAnalysisContextException) {
+                        wasThrown = true
+                    }
+                }
+            }
+        }
+        Assert.assertTrue("WriteActionStartInsideAnalysisContextException was not thrown", wasThrown)
+    }
+
+
+    private fun createFakeKtFile(): KtFile =
+        myFixture.configureByText("file.kt", "fun a() {}") as KtFile
+
+    override fun getProjectDescriptor(): LightProjectDescriptor =
+        KotlinLightProjectDescriptor.INSTANCE
+}