From 86348ed981597bd975e9820bc9c6e1498e64206a Mon Sep 17 00:00:00 2001
From: Roman Golyshev <roman.golyshev@jetbrains.com>
Date: Mon, 29 Mar 2021 16:34:53 +0300
Subject: [PATCH] Refactor `KotlinFirCompletionContributor`

Remove some code duplication
---
 .../KotlinFirCompletionContributor.kt         | 26 +++++++++----------
 1 file changed, 12 insertions(+), 14 deletions(-)

diff --git a/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt b/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt
index 23c99f70338..e62e0811763 100644
--- a/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt
+++ b/idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt
@@ -177,13 +177,8 @@ private class KotlinCommonCompletionProvider(
         val typeOfPossibleReceiver = explicitReceiver.getKtType()
         val possibleReceiverScope = typeOfPossibleReceiver.getTypeScope() ?: return
 
-        val nonExtensionMembers = possibleReceiverScope
-            .getCallableSymbols(scopeNameFilter)
-            .filterNot { it.isExtension }
-
-        val extensionNonMembers = implicitScopes
-            .getCallableSymbols(scopeNameFilter)
-            .filter { it.isExtension && it.hasSuitableExtensionReceiver() }
+        val nonExtensionMembers = possibleReceiverScope.collectNonExtensions()
+        val extensionNonMembers = implicitScopes.collectSuitableExtensions(hasSuitableExtensionReceiver)
 
         nonExtensionMembers.forEach { addSymbolToCompletion(result, expectedType, it) }
         extensionNonMembers.forEach { addSymbolToCompletion(result, expectedType, it) }
@@ -201,13 +196,8 @@ private class KotlinCommonCompletionProvider(
     ) {
         val (implicitScopes, implicitReceiversTypes) = implicitScopesContext
 
-        val availableNonExtensions = implicitScopes
-            .getCallableSymbols(scopeNameFilter)
-            .filterNot { it.isExtension }
-
-        val extensionsWhichCanBeCalled = implicitScopes
-            .getCallableSymbols(scopeNameFilter)
-            .filter { it.isExtension && it.hasSuitableExtensionReceiver() }
+        val availableNonExtensions = implicitScopes.collectNonExtensions()
+        val extensionsWhichCanBeCalled = implicitScopes.collectSuitableExtensions(hasSuitableExtensionReceiver)
 
         availableNonExtensions.forEach { addSymbolToCompletion(result, expectedType, it) }
         extensionsWhichCanBeCalled.forEach { addSymbolToCompletion(result, expectedType, it) }
@@ -237,6 +227,14 @@ private class KotlinCommonCompletionProvider(
             .filter(hasSuitableExtensionReceiver)
     }
 
+    private fun KtScope.collectNonExtensions(): Sequence<KtCallableSymbol> =
+        getCallableSymbols(scopeNameFilter).filterNot { it.isExtension }
+
+    private fun KtScope.collectSuitableExtensions(
+        hasSuitableExtensionReceiver: KtCallableSymbol.() -> Boolean
+    ): Sequence<KtCallableSymbol> =
+        getCallableSymbols(scopeNameFilter).filter { it.isExtension && it.hasSuitableExtensionReceiver() }
+
     private fun KtAnalysisSession.findAllNamesOfTypes(implicitReceiversTypes: List<KtType>) =
         implicitReceiversTypes.flatMapTo(hashSetOf()) { with(typeNamesProvider) { findAllNames(it) } }
 }