From 98e9238763298d54351a3f78986a55f797e2d4a1 Mon Sep 17 00:00:00 2001
From: Nikolay Krasko <nikolay.krasko@jetbrains.com>
Date: Fri, 15 Sep 2017 16:05:06 +0300
Subject: [PATCH] Force canceling injection search when it's executed without
 progress (KT-19901)

More accurate fix for KT-19901.
There're cases when injectors are counted without progress indicator but
under read action. In that cases typing can't be started until exit from
the computing that produces lags in typing.

Tested with long spek-like test file.

 #KT-19901 Fixed
---
 .../org/jetbrains/kotlin/idea/ProgressUtil.kt  | 12 ++++++++++++
 .../idea/injection/KotlinLanguageInjector.kt   | 18 +++++++++++++++---
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/idea/src/org/jetbrains/kotlin/idea/ProgressUtil.kt b/idea/src/org/jetbrains/kotlin/idea/ProgressUtil.kt
index a6a37a68aba..f956f593b05 100644
--- a/idea/src/org/jetbrains/kotlin/idea/ProgressUtil.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/ProgressUtil.kt
@@ -39,6 +39,18 @@ fun <T : Any> runInReadActionWithWriteActionPriorityWithPCE(f: () -> T): T {
     return r!!
 }
 
+fun <T : Any> runInReadActionWithWriteActionPriority(f: () -> T): T? {
+    var r: T? = null
+    val complete = ProgressIndicatorUtils.runInReadActionWithWriteActionPriority {
+        r = f()
+    }
+
+    if (!complete) return null
+    return r
+}
+
+
+
 fun <T : Any> Project.runSynchronouslyWithProgress(progressTitle: String, canBeCanceled: Boolean, action: () -> T): T? {
     var result: T? = null
     ProgressManager.getInstance().runProcessWithProgressSynchronously({ result = action() }, progressTitle, canBeCanceled, this)
diff --git a/idea/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt b/idea/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt
index ba6881b19f5..18ee5f0cf4e 100644
--- a/idea/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt
@@ -38,6 +38,7 @@ import org.intellij.plugins.intelliLang.util.AnnotationUtilEx
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.references.KtReference
+import org.jetbrains.kotlin.idea.runInReadActionWithWriteActionPriority
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
@@ -85,9 +86,20 @@ class KotlinLanguageInjector(
                 // Cache is up-to-date
                 kotlinCachedInjection.baseInjection
             else -> {
-                val computedInjection = computeBaseInjection(ktHost, support, registrar) ?: ABSENT_KOTLIN_INJECTION
-                ktHost.cachedInjectionWithModification = KotlinCachedInjection(modificationCount, computedInjection)
-                computedInjection
+                fun computeAndCache(): BaseInjection {
+                    val computedInjection = computeBaseInjection(ktHost, support, registrar) ?: ABSENT_KOTLIN_INJECTION
+                    ktHost.cachedInjectionWithModification = KotlinCachedInjection(modificationCount, computedInjection)
+                    return computedInjection
+                }
+
+                if (ApplicationManager.getApplication().isReadAccessAllowed && ProgressManager.getInstance().progressIndicator == null) {
+                    // The action cannot be canceled by caller and by internal checkCanceled() calls.
+                    // Force creating new indicator that is canceled on write action start, otherwise there might be lags in typing.
+                    runInReadActionWithWriteActionPriority(::computeAndCache) ?: kotlinCachedInjection?.baseInjection ?: ABSENT_KOTLIN_INJECTION
+                }
+                else {
+                    computeAndCache()
+                }
             }
         }