From b6bd3e15c52ee2ceb50652b2a53d15b3092da5b3 Mon Sep 17 00:00:00 2001
From: Nikolay Krasko <nikolay.krasko@jetbrains.com>
Date: Fri, 21 Jul 2017 15:54:18 +0300
Subject: [PATCH] Fix outdated highlighting for Kotlin injected fragments
 (KT-18842)

Do not use cached value for injection in non-dispatch threads and throw
PCE if value isn't ready. Otherwise outdated value may be stored in
highlighting.

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

diff --git a/idea/src/org/jetbrains/kotlin/idea/ProgressUtil.kt b/idea/src/org/jetbrains/kotlin/idea/ProgressUtil.kt
index 9e2806c52be..a6a37a68aba 100644
--- a/idea/src/org/jetbrains/kotlin/idea/ProgressUtil.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/ProgressUtil.kt
@@ -22,27 +22,19 @@ import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.progress.util.ProgressIndicatorUtils
 import com.intellij.openapi.project.Project
 
-fun <T : Any> runInReadActionWithWriteActionPriority(f: () -> T): T? {
+fun <T : Any> runInReadActionWithWriteActionPriorityWithPCE(f: () -> T): T {
     var r: T? = null
     if (!with(ApplicationManager.getApplication()) { isDispatchThread && isUnitTestMode }) {
-        ProgressIndicatorUtils.runInReadActionWithWriteActionPriority {
+        val complete = ProgressIndicatorUtils.runInReadActionWithWriteActionPriority {
             r = f()
         }
+
+        // There is a write action in progress or pending, so no point in counting the result
+        if (!complete) throw ProcessCanceledException()
     }
     else {
         r = f()
     }
-    return r
-}
-
-fun <T : Any> runInReadActionWithWriteActionPriorityWithPCE(f: () -> T): T {
-    var r: T? = null
-    val complete = ProgressIndicatorUtils.runInReadActionWithWriteActionPriority {
-        r = f()
-    }
-
-    // There is a write action in progress or pending, so no point in counting the result
-    if (!complete) throw ProcessCanceledException()
 
     return r!!
 }
diff --git a/idea/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt b/idea/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt
index 24d6bf6731c..2205dd1e194 100644
--- a/idea/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/injection/KotlinLanguageInjector.kt
@@ -38,7 +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.runInReadActionWithWriteActionPriorityWithPCE
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
@@ -74,19 +74,18 @@ class KotlinLanguageInjector(
 
         val needImmediateAnswer = with(ApplicationManager.getApplication()) { isDispatchThread && !isUnitTestMode }
         val kotlinCachedInjection = ktHost.cachedInjectionWithModification
-        val cachedBaseInjection = kotlinCachedInjection?.baseInjection ?: ABSENT_KOTLIN_INJECTION
         val modificationCount = PsiManager.getInstance(project).modificationTracker.modificationCount
 
         val baseInjection = when {
-            needImmediateAnswer -> cachedBaseInjection
+            needImmediateAnswer -> kotlinCachedInjection?.baseInjection ?: ABSENT_KOTLIN_INJECTION
             kotlinCachedInjection != null && (modificationCount == kotlinCachedInjection.modificationCount) ->
                 kotlinCachedInjection.baseInjection
             else -> {
-                runInReadActionWithWriteActionPriority {
+                runInReadActionWithWriteActionPriorityWithPCE {
                     val computedInjection = computeBaseInjection(ktHost, support, registrar) ?: ABSENT_KOTLIN_INJECTION
                     ktHost.cachedInjectionWithModification = KotlinCachedInjection(modificationCount, computedInjection)
                     computedInjection
-                } ?: cachedBaseInjection
+                }
             }
         }