From 5ad94daaa5817acda863ad1d16e62358be0a58b7 Mon Sep 17 00:00:00 2001
From: Dmitry Gridin <dmitry.gridin@jetbrains.com>
Date: Tue, 23 Jun 2020 16:07:17 +0700
Subject: [PATCH] ImportOptimizer: cleanup code

#KT-31331
---
 .../idea/imports/OptimizedImportsBuilder.kt       | 12 ++++++++++--
 .../kotlin/idea/imports/KotlinImportOptimizer.kt  | 15 +++++++++------
 2 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/idea/ide-common/src/org/jetbrains/kotlin/idea/imports/OptimizedImportsBuilder.kt b/idea/ide-common/src/org/jetbrains/kotlin/idea/imports/OptimizedImportsBuilder.kt
index 19e1f270c84..523fd8caef4 100644
--- a/idea/ide-common/src/org/jetbrains/kotlin/idea/imports/OptimizedImportsBuilder.kt
+++ b/idea/ide-common/src/org/jetbrains/kotlin/idea/imports/OptimizedImportsBuilder.kt
@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.descriptors.impl.TypeAliasConstructorDescriptor
 import org.jetbrains.kotlin.idea.analysis.analyzeAsReplacement
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
+import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.resolve.frontendService
 import org.jetbrains.kotlin.idea.util.ImportInsertHelper
 import org.jetbrains.kotlin.idea.util.getResolutionScope
@@ -82,8 +83,10 @@ class OptimizedImportsBuilder(
     fun buildOptimizedImports(): List<ImportPath>? {
         val facade = file.getResolutionFacade()
         file.importDirectives
-            .filter { it.isAllUnder && data.unresolvedNames.isNotEmpty() || it.importedName in data.unresolvedNames && !it.canResolve(facade) }
-            .mapNotNullTo(importRules) { it.importPath?.let { path -> ImportRule.Add(path) } }
+            .asSequence()
+            .filter { it.mayReferToSomeUnresolvedName() || it.isExistedUnresolvedName(facade) }
+            .mapNotNull { it.importPath }
+            .mapNotNullTo(importRules) { ImportRule.Add(it) }
 
         while (true) {
             ProgressManager.checkCanceled()
@@ -94,6 +97,11 @@ class OptimizedImportsBuilder(
         }
     }
 
+    private fun KtImportDirective.mayReferToSomeUnresolvedName() = isAllUnder && data.unresolvedNames.isNotEmpty()
+
+    private fun KtImportDirective.isExistedUnresolvedName(facade: ResolutionFacade) =
+        importedName in data.unresolvedNames && !canResolve(facade)
+
     private fun getExpressionToAnalyze(element: KtElement): KtExpression? {
         val parent = element.parent
         return when {
diff --git a/idea/src/org/jetbrains/kotlin/idea/imports/KotlinImportOptimizer.kt b/idea/src/org/jetbrains/kotlin/idea/imports/KotlinImportOptimizer.kt
index bffbeb332d9..5c75b7bae7e 100644
--- a/idea/src/org/jetbrains/kotlin/idea/imports/KotlinImportOptimizer.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/imports/KotlinImportOptimizer.kt
@@ -151,12 +151,7 @@ class KotlinImportOptimizer : ImportOptimizer {
 
             val references = element.references.ifEmpty { return }
             val bindingContext = element.analyze(BodyResolveMode.PARTIAL)
-            val isResolved = if (element is KtCallElement)
-                element.getResolvedCall(bindingContext) != null
-            else
-                element.mainReference?.resolveToDescriptors(bindingContext)?.let { descriptors ->
-                    descriptors.isNotEmpty() && descriptors.none { it is ErrorSimpleFunctionDescriptorImpl }
-                } == true
+            val isResolved = hasResolvedDescriptor(element, bindingContext)
 
             for (reference in references) {
                 if (reference !is KtReference) continue
@@ -302,3 +297,11 @@ class KotlinImportOptimizer : ImportOptimizer {
         }
     }
 }
+
+private fun hasResolvedDescriptor(element: KtElement, bindingContext: BindingContext): Boolean =
+    if (element is KtCallElement)
+        element.getResolvedCall(bindingContext) != null
+    else
+        element.mainReference?.resolveToDescriptors(bindingContext)?.let { descriptors ->
+            descriptors.isNotEmpty() && descriptors.none { it is ErrorSimpleFunctionDescriptorImpl }
+        } == true