From 54d6fc846be0535243549d9f47ce18a0a4a3f02e Mon Sep 17 00:00:00 2001
From: Ilya Kirillov <ilya.kirillov@jetbrains.com>
Date: Mon, 8 Apr 2019 19:31:34 +0300
Subject: [PATCH] New J2K: Add file range support for nullability analysis

---
 .../kotlin/nj2k/nullabilityAnalysis/context.kt         |  4 ++++
 .../nj2k/nullabilityAnalysis/nullabilityAnalysis.kt    | 10 ++++++----
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/context.kt b/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/context.kt
index 213ee7c748a..1e319e55fe8 100644
--- a/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/context.kt
+++ b/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/context.kt
@@ -5,9 +5,13 @@
 
 package org.jetbrains.kotlin.nj2k.nullabilityAnalysis
 
+import com.intellij.openapi.editor.RangeMarker
+import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
+import org.jetbrains.kotlin.psi.psiUtil.elementsInRange
 import org.jetbrains.kotlin.utils.addToStdlib.safeAs
 
 internal class ContextCreator(private val getNullability: (KtTypeElement) -> Nullability) {
diff --git a/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/nullabilityAnalysis.kt b/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/nullabilityAnalysis.kt
index 5a0ef40103b..91ad895c99a 100644
--- a/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/nullabilityAnalysis.kt
+++ b/nj2k/nj2k-services/src/org/jetbrains/kotlin/nj2k/nullabilityAnalysis/nullabilityAnalysis.kt
@@ -82,10 +82,11 @@ class NullabilityAnalysisFacade(
             runWriteAction {
                 analysisScope.prepareTypeElements(prepareTypeElement)
                 val context = ContextCreator(getTypeElementNullability).createContext(analysisScope)
-                if (debugPrint)
+                if (debugPrint) {
                     with(Printer(context)) {
-                        analysisScope.single().addTypeVariablesNames()
+                        analysisScope.forEach { it.addTypeVariablesNames() }
                     }
+                }
 
                 val constraints = ConstraintsCollector(context, debugPrint).collectConstraints(analysisScope)
                 Solver(context, debugPrint).solveConstraints(constraints)
@@ -97,8 +98,9 @@ class NullabilityAnalysisFacade(
 }
 
 private fun AnalysisScope.prepareTypeElements(prepareTypeElement: (KtTypeElement) -> Unit) {
-    val typeElements = flatMap { it.collectDescendantsOfType<KtTypeElement>() }
-    typeElements.forEach { typeElement ->
+    val typeElements = flatMap { it.collectDescendantsOfType<KtTypeReference>() }
+    typeElements.forEach { typeReference ->
+        val typeElement = typeReference.typeElement ?: return@forEach
         if (typeElement.parentOfType<KtSuperTypeCallEntry>() == null) {
             prepareTypeElement(typeElement)
         }