[AA] Fix write access configuration in unit test application environments

- We cannot configure the application's write action accessibility on a
  per-test basis because (1) the application may be shared across
  concurrent tests and (2) the application is usually cached, so the
  configuration will be missed entirely.
- There is actually a much easier solution to allow write access
  selectively: We can enable it in `runWriteAction` blocks, and keep it
  in a thread local to support concurrent test runs. As Analysis API
  tests never call `runWriteAction`, there will be no "analyze cannot be
  called from a write action" error, and if `analyze` is somehow called
  from a write action, it will now be caught.

^KT-63560 fixed

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/base/registerBaseTestServices.kt

@@ -26,7 +26,7 @@ fun TestConfigurationBuilder.registerAnalysisApiBaseTestServices(
 ) {
     useAdditionalService<TestDisposableProvider>(::TestDisposableProviderImpl)
     useAdditionalService<AnalysisApiKtModuleProvider>(::AnalysisApiKtModuleProviderImpl)
-    useAdditionalService<AnalysisApiEnvironmentManager>(::AnalysisApiEnvironmentManagerImpl.bind(testDisposable, configurator))
+    useAdditionalService<AnalysisApiEnvironmentManager>(::AnalysisApiEnvironmentManagerImpl.bind(testDisposable))
     useAdditionalService<ApplicationDisposableProvider> { ExecutionListenerBasedDisposableProvider() }
     useAdditionalService<KotlinStandardLibrariesPathProvider> { StandardLibrariesPathProviderForKotlinProject }
 

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/services/AnalysisApiEnvironmentManager.kt

@@ -40,13 +40,12 @@ abstract class AnalysisApiEnvironmentManager : TestService {
 class AnalysisApiEnvironmentManagerImpl(
     override val testServices: TestServices,
     override val testRootDisposable: Disposable,
-    private val configurator: AnalysisApiTestConfigurator,
 ) : AnalysisApiEnvironmentManager() {
     private val _projectEnvironment: KotlinCoreProjectEnvironment by lazy {
         StandaloneProjectFactory.createProjectEnvironment(
             testRootDisposable,
             testServices.applicationDisposableProvider.getApplicationRootDisposable(),
-            KotlinCoreApplicationEnvironmentMode.UnitTest(configurator.isWriteAccessAllowed),
+            KotlinCoreApplicationEnvironmentMode.UnitTest,
         )
     }
 

analysis/analysis-test-framework/tests/org/jetbrains/kotlin/analysis/test/framework/test/configurators/AnalysisApiTestConfigurator.kt

@@ -36,15 +36,6 @@ abstract class AnalysisApiTestConfigurator {
 
     abstract val serviceRegistrars: List<AnalysisApiTestServiceRegistrar>
 
-    /**
-     * Whether the associated [core application environment][org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreApplicationEnvironment] should
-     * permit write access.
-     *
-     * Write access is disabled by default because the Analysis API should not be used during write actions. Nevertheless, some tests may
-     * need write access, for example to modify PSI.
-     */
-    open val isWriteAccessAllowed get() = false
-
     open fun prepareFilesInModule(files: List<PsiFile>, module: TestModule, testServices: TestServices) {}
 
     open fun doGlobalModuleStateModification(project: Project) {

analysis/low-level-api-fir/testData/declarationModificationService/psiResilience/callExpression/contractCall.kt

@@ -1,6 +1,3 @@
-// IGNORE_FIR
-// Reason: KT-63560
-
 import kotlin.contracts.contract
 import kotlin.contracts.InvocationKind
 

analysis/low-level-api-fir/testData/declarationModificationService/psiResilience/callExpression/printlnFirstStatement.kt

@@ -1,6 +1,3 @@
-// IGNORE_FIR
-// Reason: KT-63560
-
 fun foo(string: String) {
     <expr>println("Hello world, $string!")</expr>
 }

analysis/low-level-api-fir/testData/declarationModificationService/psiResilience/callExpression/printlnSecondStatement.kt

@@ -1,6 +1,3 @@
-// IGNORE_FIR
-// Reason: KT-63560
-
 fun foo(string: String) {
     val a = 7
     <expr>println("Hello world, $string!")</expr>

analysis/low-level-api-fir/testData/declarationModificationService/psiResilience/dotQualifiedExpression/contractCall.kt

@@ -1,6 +1,3 @@
-// IGNORE_FIR
-// Reason: KT-63560
-
 import kotlin.contracts.InvocationKind
 
 inline fun foo(block: () -> Unit) {

analysis/low-level-api-fir/testData/declarationModificationService/psiResilience/dotQualifiedExpression/stringLengthFirstStatement.kt

@@ -1,6 +1,3 @@
-// IGNORE_FIR
-// Reason: KT-63560
-
 fun foo(string: String) {
     <expr>string.length</expr>
 }

analysis/low-level-api-fir/testData/declarationModificationService/psiResilience/dotQualifiedExpression/stringLengthSecondStatement.kt

@@ -1,6 +1,3 @@
-// IGNORE_FIR
-// Reason: KT-63560
-
 fun foo(string: String) {
     println()
     <expr>string.length</expr>

analysis/low-level-api-fir/testData/declarationModificationService/psiResilience/propertyDeclaration/simpleProperty.kt

@@ -1,4 +1 @@
-// IGNORE_FIR
-// Reason: KT-63560
-
 <expr>val property: Int = 5</expr>

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/test/base/AbstractLowLevelApiModifiablePsiSingleFileTest.kt

@@ -40,6 +40,4 @@ object AnalysisApiFirModifiablePsiSourceTestConfigurator : AnalysisApiFirSourceT
         addAll(super.serviceRegistrars)
         add(AnalysisApiModifiablePsiTestServiceRegistrar)
     }
-
-    override val isWriteAccessAllowed: Boolean get() = true
 }

compiler/cli/cli-base/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreApplicationEnvironment.kt

@@ -13,7 +13,9 @@ import com.intellij.lang.MetaLanguage
 import com.intellij.mock.MockApplication
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.util.Computable
 import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.ThrowableComputable
 import com.intellij.openapi.vfs.VirtualFileSystem
 import com.intellij.psi.FileContextProvider
 import com.intellij.psi.augment.PsiAugmentProvider
@@ -25,28 +27,22 @@ import org.jetbrains.kotlin.cli.jvm.compiler.jarfs.FastJarFileSystem
 import org.jetbrains.kotlin.cli.jvm.modules.CoreJrtFileSystem
 
 sealed interface KotlinCoreApplicationEnvironmentMode {
-    val isWriteAccessAllowed: Boolean
+    object Production : KotlinCoreApplicationEnvironmentMode
 
-    object Production : KotlinCoreApplicationEnvironmentMode {
-        override val isWriteAccessAllowed: Boolean get() = true
-    }
-
-    class UnitTest(override val isWriteAccessAllowed: Boolean) : KotlinCoreApplicationEnvironmentMode
+    object UnitTest : KotlinCoreApplicationEnvironmentMode
 
     companion object {
         fun fromUnitTestModeFlag(isUnitTestMode: Boolean): KotlinCoreApplicationEnvironmentMode =
-            if (isUnitTestMode) UnitTest(isWriteAccessAllowed = false) else Production
+            if (isUnitTestMode) UnitTest else Production
     }
 }
 
 class KotlinCoreApplicationEnvironment private constructor(
     parentDisposable: Disposable,
-    private val environmentMode: KotlinCoreApplicationEnvironmentMode,
-) : JavaCoreApplicationEnvironment(parentDisposable, environmentMode is KotlinCoreApplicationEnvironmentMode.UnitTest) {
+    environmentMode: KotlinCoreApplicationEnvironmentMode,
+) : JavaCoreApplicationEnvironment(parentDisposable, environmentMode == KotlinCoreApplicationEnvironmentMode.UnitTest) {
 
     init {
-        application.initializeEnvironmentMode(environmentMode)
-
         registerApplicationService(JavaFileCodeStyleFacadeFactory::class.java, DummyJavaFileCodeStyleFacadeFactory())
         registerFileType(JavaClassFileType.INSTANCE, "sig")
     }
@@ -63,7 +59,7 @@ class KotlinCoreApplicationEnvironment private constructor(
          * is not yet initialized when this function is called from the superclass constructor.
          */
         return if (mock.isUnitTestMode) {
-            MockUnitTestApplication(parentDisposable)
+            KotlinCoreUnitTestApplication(parentDisposable)
         } else {
             mock
         }
@@ -129,25 +125,42 @@ class KotlinCoreApplicationEnvironment private constructor(
     }
 }
 
-private class MockUnitTestApplication(parentDisposable: Disposable) : MockApplication(parentDisposable) {
+/**
+ * A [MockApplication] which allows write actions only in [runWriteAction] blocks.
+ *
+ * The Analysis API is usually not allowed to be used from a write action. To properly support Analysis API tests, the application should
+ * not return `true` for [isWriteAccessAllowed] indiscriminately like [MockApplication]. On the other hand, we have some tests which require
+ * write access, but don't access the Analysis API. Hence, we remember which threads have started a write action with [runWriteAction].
+ */
+private class KotlinCoreUnitTestApplication(parentDisposable: Disposable) : MockApplication(parentDisposable) {
     /**
-     * We can't use [KotlinCoreApplicationEnvironment.environmentMode] in [KotlinCoreApplicationEnvironment.createApplication], because the
-     * corresponding property is not yet initialized when `createApplication` is called from the superclass constructor. So we have to
-     * initialize it later.
+     * We need to remember whether write access is allowed per thread because the application can be shared between multiple concurrent test
+     * runs.
      */
-    lateinit var environmentMode: KotlinCoreApplicationEnvironmentMode
+    private val isWriteAccessAllowedInThread: ThreadLocal<Boolean> = ThreadLocal.withInitial { false }
 
     override fun isUnitTestMode(): Boolean = true
 
-    override fun isWriteAccessAllowed(): Boolean = environmentMode.isWriteAccessAllowed
-}
+    override fun isWriteAccessAllowed(): Boolean = isWriteAccessAllowedInThread.get()
 
-private fun MockApplication.initializeEnvironmentMode(environmentMode: KotlinCoreApplicationEnvironmentMode) {
-    if (this !is MockUnitTestApplication) return
+    override fun runWriteAction(action: Runnable) {
+        withWriteAccessAllowedInThread {
+            action.run()
+        }
+    }
 
-    this.environmentMode = environmentMode
+    override fun <T : Any?> runWriteAction(computation: Computable<T?>): T? =
+        withWriteAccessAllowedInThread { computation.compute() }
 
-    require(isWriteAccessAllowed == environmentMode.isWriteAccessAllowed) {
-        "The mock application's `isWriteAccessAllowed` should correspond to the environment mode configuration."
+    override fun <T : Any?, E : Throwable?> runWriteAction(computation: ThrowableComputable<T?, E?>): T? =
+        withWriteAccessAllowedInThread { computation.compute() }
+
+    private inline fun <A> withWriteAccessAllowedInThread(action: () -> A): A {
+        isWriteAccessAllowedInThread.set(true)
+        try {
+            return action()
+        } finally {
+            isWriteAccessAllowedInThread.set(false)
+        }
     }
 }

compiler/cli/cli-base/src/org/jetbrains/kotlin/cli/jvm/compiler/KotlinCoreEnvironment.kt

@@ -471,7 +471,7 @@ class KotlinCoreEnvironment private constructor(
                 createApplicationEnvironment(
                     parentDisposable,
                     configuration,
-                    KotlinCoreApplicationEnvironmentMode.UnitTest(isWriteAccessAllowed = false),
+                    KotlinCoreApplicationEnvironmentMode.UnitTest,
                 )
             val projectEnv = ProjectEnvironment(parentDisposable, appEnv, configuration)
             return KotlinCoreEnvironment(projectEnv, configuration, extensionConfigs)
@@ -501,7 +501,7 @@ class KotlinCoreEnvironment private constructor(
             val appEnv = createApplicationEnvironment(
                 parentDisposable,
                 configuration,
-                KotlinCoreApplicationEnvironmentMode.UnitTest(isWriteAccessAllowed = false),
+                KotlinCoreApplicationEnvironmentMode.UnitTest,
             )
             return ProjectEnvironment(parentDisposable, appEnv, configuration)
         }
@@ -523,7 +523,7 @@ class KotlinCoreEnvironment private constructor(
         ): KotlinCoreApplicationEnvironment = getOrCreateApplicationEnvironment(
             parentDisposable,
             configuration,
-            KotlinCoreApplicationEnvironmentMode.UnitTest(isWriteAccessAllowed = false),
+            KotlinCoreApplicationEnvironmentMode.UnitTest,
         )
 
         fun getOrCreateApplicationEnvironment(