Force resolve check for dispatch thread instead of isWriteAccessAllowed

isWriteAccessAllowed = dispatch thread + write lock

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/ResolveInWriteActionManager.kt

@@ -20,6 +20,18 @@ fun <T> allowResolveInWriteAction(runnable: () -> T): T {
     return ResolveInWriteActionManager.runWithResolveAllowedInWriteAction(runnable)
 }
 
+/**
+ * Temporary allow resolve in write action.
+ *
+ * All resolve should be banned from dispatch thread. This method should be carefully used for the transition
+ * period at the places where developer expects everything is already cached.
+ *
+ * (A better check is needed instead that would assert no resolve result are obtained outside of caches.)
+ */
+fun <T> allowResolveExpectedToBeCached(runnable: () -> T): T {
+    return ResolveInWriteActionManager.runWithResolveAllowedInWriteAction(runnable)
+}
+
 /**
  * Force resolve check in tests.
  */
@@ -35,21 +47,23 @@ class ResolveInWriteActionException(message: String? = null) : IllegalThreadStat
 internal object ResolveInWriteActionManager {
     private val LOG = Logger.getInstance(ResolveInWriteActionManager::class.java)
 
-    private val isResolveInWriteActionAllowed: ThreadLocal<Boolean> = ThreadLocal.withInitial { false }
-    private val isForceCheckInTests: ThreadLocal<Boolean> = ThreadLocal.withInitial { false }
-    private val errorHandler: ThreadLocal<(() -> Unit)?> = ThreadLocal.withInitial { null }
+    // Should be accessed only from dispatch thread
+    private var isResolveAllowed: Boolean = false
+    private var isForceCheckInTests: Boolean = false
+    private var errorHandler: (() -> Unit)? = null
 
     fun assertNoResolveUnderWriteAction() {
-        if (isResolveInWriteActionAllowed.get()) return
-
         val application = ApplicationManager.getApplication() ?: return
+        if (!application.isDispatchThread) {
+            return
+        }
 
-        if (!application.isWriteAccessAllowed) return
+        if (isResolveAllowed) return
 
         if (application.isUnitTestMode) {
-            if (!isForceCheckInTests.get()) return
+            if (!isForceCheckInTests) return
 
-            val handler = errorHandler.get()
+            val handler = errorHandler
             if (handler != null) {
                 handler()
                 return
@@ -66,26 +80,27 @@ internal object ResolveInWriteActionManager {
 
     internal inline fun <T> runWithResolveAllowedInWriteAction(runnable: () -> T): T {
         val wasSet =
-            if (ApplicationManager.getApplication()?.isWriteAccessAllowed == true && isResolveInWriteActionAllowed.get() == false) {
-                isResolveInWriteActionAllowed.set(true)
+            if (ApplicationManager.getApplication()?.isDispatchThread == true && !isResolveAllowed) {
+                isResolveAllowed = true
                 true
             } else {
                 false
             }
-
         try {
             return runnable()
         } finally {
             if (wasSet) {
-                isResolveInWriteActionAllowed.set(false)
+                isResolveAllowed = false
             }
         }
     }
 
     internal fun <T> runWithForceResolveAllowedCheckInTests(errorHandler: (() -> Unit)?, runnable: () -> T): T {
-        val wasSet = if (isForceCheckInTests.get() == false) {
-            isForceCheckInTests.set(true)
-            this.errorHandler.set(errorHandler)
+        ApplicationManager.getApplication()?.assertIsDispatchThread()
+
+        val wasSet = if (!isForceCheckInTests) {
+            isForceCheckInTests = true
+            this.errorHandler = errorHandler
             true
         } else {
             false
@@ -95,8 +110,8 @@ internal object ResolveInWriteActionManager {
             return runnable()
         } finally {
             if (wasSet) {
-                isForceCheckInTests.set(false)
-                this.errorHandler.set(null)
+                isForceCheckInTests = false
+                this.errorHandler = null
             }
         }
     }

idea/src/org/jetbrains/kotlin/idea/quickfix/KotlinQuickFixAction.kt

@@ -10,6 +10,7 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
+import org.jetbrains.kotlin.idea.caches.resolve.allowResolveExpectedToBeCached
 import org.jetbrains.kotlin.psi.KtFile
 
 abstract class KotlinQuickFixAction<out T : PsiElement>(element: T) : QuickFixActionBase<T>(element) {
@@ -17,7 +18,11 @@ abstract class KotlinQuickFixAction<out T : PsiElement>(element: T) : QuickFixAc
 
     override fun isAvailableImpl(project: Project, editor: Editor?, file: PsiFile): Boolean {
         val ktFile = file as? KtFile ?: return false
-        return isAvailable(project, editor, ktFile)
+        return allowResolveExpectedToBeCached {
+            // Quick fixes availability is checked in UI thread but only after background highlighting pass is finished
+            // so we are expecting that every relevant results are already cached.
+            isAvailable(project, editor, ktFile)
+        }
     }
 
     final override fun invoke(project: Project, editor: Editor?, file: PsiFile) {

idea/testData/quickfix/allowResolveInWriteAction.txt

@@ -1,19 +1,22 @@
 # Actions that are allowed to resolve in write action. Normally this set shouldn't be extended and eventually should be dropped.
 # Please consider rewriting a quick-fix and remove resolve from it before adding a new entry to this list.
 
+org.jetbrains.kotlin.idea.core.overrideImplement.ImplementAsConstructorParameter
 org.jetbrains.kotlin.idea.core.overrideImplement.ImplementMembersHandler
+org.jetbrains.kotlin.idea.inspections.ConflictingExtensionPropertyInspection$DeleteRedundantExtensionAction
 org.jetbrains.kotlin.idea.inspections.IntentionBasedInspection$HighPriorityIntentionBasedQuickFix
 org.jetbrains.kotlin.idea.inspections.KotlinUnusedImportInspection$OptimizeImportsQuickFix
+org.jetbrains.kotlin.idea.inspections.migration.ObsoleteCoroutineUsageFix
 org.jetbrains.kotlin.idea.inspections.SafeDeleteFix
 org.jetbrains.kotlin.idea.inspections.SuspiciousCollectionReassignmentInspection$ChangeTypeToMutableFix
 org.jetbrains.kotlin.idea.inspections.SuspiciousCollectionReassignmentInspection$JoinWithInitializerFix
 org.jetbrains.kotlin.idea.inspections.UnusedReceiverParameterInspection$RemoveReceiverFix
-org.jetbrains.kotlin.idea.inspections.migration.ObsoleteCoroutineUsageFix
 org.jetbrains.kotlin.idea.intentions.AddPropertyAccessorsIntention
 org.jetbrains.kotlin.idea.intentions.AddPropertyGetterIntention
 org.jetbrains.kotlin.idea.intentions.AddPropertySetterIntention
 org.jetbrains.kotlin.idea.intentions.ConvertPropertyInitializerToGetterIntention
 org.jetbrains.kotlin.idea.intentions.CreateKotlinSubClassIntention
+org.jetbrains.kotlin.idea.intentions.MoveMemberToCompanionObjectIntention
 org.jetbrains.kotlin.idea.intentions.MovePropertyToConstructorIntention
 org.jetbrains.kotlin.idea.intentions.SpecifyTypeExplicitlyIntention
 org.jetbrains.kotlin.idea.quickfix.AddAnnotationTargetFix
@@ -26,9 +29,14 @@ org.jetbrains.kotlin.idea.quickfix.AddStartProjectionsForInnerClass
 org.jetbrains.kotlin.idea.quickfix.AddWhenRemainingBranchesFix
 org.jetbrains.kotlin.idea.quickfix.AssignToPropertyFix
 org.jetbrains.kotlin.idea.quickfix.ChangeAccessorTypeFix
+org.jetbrains.kotlin.idea.quickfix.ChangeCallableReturnTypeFix$ForCalled
 org.jetbrains.kotlin.idea.quickfix.ChangeCallableReturnTypeFix$ForEnclosing
+org.jetbrains.kotlin.idea.quickfix.ChangeCallableReturnTypeFix$ForOverridden
+org.jetbrains.kotlin.idea.quickfix.ChangeFunctionLiteralReturnTypeFix
 org.jetbrains.kotlin.idea.quickfix.ChangeParameterTypeFix
 org.jetbrains.kotlin.idea.quickfix.ChangeToMutableCollectionFix
+org.jetbrains.kotlin.idea.quickfix.ChangeVariableTypeFix
+org.jetbrains.kotlin.idea.quickfix.ChangeVariableTypeFix$ForOverridden
 org.jetbrains.kotlin.idea.quickfix.ChangeVisibilityFix$ChangeToInternalFix
 org.jetbrains.kotlin.idea.quickfix.ChangeVisibilityFix$ChangeToPrivateFix
 org.jetbrains.kotlin.idea.quickfix.ChangeVisibilityFix$ChangeToProtectedFix
@@ -36,6 +44,8 @@ org.jetbrains.kotlin.idea.quickfix.ChangeVisibilityFix$ChangeToPublicFix
 org.jetbrains.kotlin.idea.quickfix.ConvertExtensionPropertyInitializerToGetterFix
 org.jetbrains.kotlin.idea.quickfix.ConvertMemberToExtensionFix
 org.jetbrains.kotlin.idea.quickfix.ConvertToAnonymousObjectFix
+org.jetbrains.kotlin.idea.quickfix.createFromUsage.createVariable.CreateLocalVariableActionFactory$createAction$2
+org.jetbrains.kotlin.idea.quickfix.DeprecatedJavaAnnotationFix
 org.jetbrains.kotlin.idea.quickfix.ImportFix
 org.jetbrains.kotlin.idea.quickfix.InitializePropertyQuickFixFactory$AddInitializerFix
 org.jetbrains.kotlin.idea.quickfix.InitializePropertyQuickFixFactory$InitializeWithConstructorParameter
@@ -45,6 +55,7 @@ org.jetbrains.kotlin.idea.quickfix.InsertDelegationCallQuickfix
 org.jetbrains.kotlin.idea.quickfix.KotlinSuppressIntentionAction
 org.jetbrains.kotlin.idea.quickfix.LetImplementInterfaceFix
 org.jetbrains.kotlin.idea.quickfix.LowPriorityQuickFixWithDelegateFactory
+org.jetbrains.kotlin.idea.quickfix.migration.MigrateExternalExtensionFix
 org.jetbrains.kotlin.idea.quickfix.QuickFixWithDelegateFactory
 org.jetbrains.kotlin.idea.quickfix.RemoveNameFromFunctionExpressionFix
 org.jetbrains.kotlin.idea.quickfix.RemovePartsFromPropertyFix
@@ -55,14 +66,12 @@ org.jetbrains.kotlin.idea.quickfix.RenameModToRemFix
 org.jetbrains.kotlin.idea.quickfix.RenameParameterToMatchOverriddenMethodFix
 org.jetbrains.kotlin.idea.quickfix.RenameUnresolvedReferenceFix
 org.jetbrains.kotlin.idea.quickfix.ReplaceInfixOrOperatorCallFix
+org.jetbrains.kotlin.idea.quickfix.replaceWith.DeprecatedSymbolUsageFix
+org.jetbrains.kotlin.idea.quickfix.replaceWith.DeprecatedSymbolUsageInWholeProjectFix
 org.jetbrains.kotlin.idea.quickfix.RestrictedRetentionForExpressionAnnotationFactory$AddSourceRetentionFix
 org.jetbrains.kotlin.idea.quickfix.SimplifyComparisonFix
 org.jetbrains.kotlin.idea.quickfix.SmartCastImpossibleInIfThenFactory$createQuickFix$1
 org.jetbrains.kotlin.idea.quickfix.SpecifyOverrideExplicitlyFix
 org.jetbrains.kotlin.idea.quickfix.SpecifyTypeExplicitlyFix
 org.jetbrains.kotlin.idea.quickfix.SuperClassNotInitialized$AddParenthesisFix
-org.jetbrains.kotlin.idea.quickfix.WrapWithSafeLetCallFix
-org.jetbrains.kotlin.idea.quickfix.createFromUsage.createVariable.CreateLocalVariableActionFactory$createAction$2
-org.jetbrains.kotlin.idea.quickfix.migration.MigrateExternalExtensionFix
-org.jetbrains.kotlin.idea.quickfix.replaceWith.DeprecatedSymbolUsageFix
-org.jetbrains.kotlin.idea.quickfix.replaceWith.DeprecatedSymbolUsageInWholeProjectFix
+org.jetbrains.kotlin.idea.quickfix.WrapWithSafeLetCallFix