From d4ba15fd906d5ac0b388d9eff5416ee6d88d852c Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Thu, 19 Nov 2015 10:28:29 +0100
Subject: [PATCH] "kotlinVariable()" macro to filter by expected type, dropped
 IterableVariableMacro

---
 .../kotlin/idea/core/ExpectedInfos.kt         |  8 ++--
 idea/resources/liveTemplates/Kotlin.xml       |  2 +-
 idea/src/META-INF/plugin.xml                  |  2 +-
 .../liveTemplates/macro/AnyVariableMacro.kt   |  4 +-
 ...iableMacro.kt => SuitableVariableMacro.kt} | 37 ++++++++++++-------
 5 files changed, 32 insertions(+), 21 deletions(-)
 rename idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/{IterableVariableMacro.kt => SuitableVariableMacro.kt} (52%)

diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ExpectedInfos.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ExpectedInfos.kt
index 0fe5a127b57..31a5019c7ab 100644
--- a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ExpectedInfos.kt
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ExpectedInfos.kt
@@ -139,10 +139,10 @@ class WhenEntryAdditionalData(val whenWithSubject: Boolean) : ExpectedInfo.Addit
 object IfConditionAdditionalData : ExpectedInfo.AdditionalData
 
 class ExpectedInfos(
-        val bindingContext: BindingContext,
-        val resolutionFacade: ResolutionFacade,
-        val useHeuristicSignatures: Boolean = true,
-        val useOuterCallsExpectedTypeCount: Int = 0
+        private val bindingContext: BindingContext,
+        private val resolutionFacade: ResolutionFacade,
+        private val useHeuristicSignatures: Boolean = true,
+        private val useOuterCallsExpectedTypeCount: Int = 0
 ) {
     public fun calculate(expressionWithType: KtExpression): Collection<ExpectedInfo> {
         val expectedInfos = calculateForArgument(expressionWithType)
diff --git a/idea/resources/liveTemplates/Kotlin.xml b/idea/resources/liveTemplates/Kotlin.xml
index 701800f396f..cf5bc18248e 100644
--- a/idea/resources/liveTemplates/Kotlin.xml
+++ b/idea/resources/liveTemplates/Kotlin.xml
@@ -42,7 +42,7 @@
     <template resource-bundle="org.jetbrains.kotlin.idea.KotlinBundle" key="livetemplate.description.iter"
               name="iter" toReformat="true" toShortenFQNames="true"
               value="for ($VAR$ in $ITERABLE$) {&#10;$END$&#10;}">
-        <variable alwaysStopAt="true" defaultValue="&quot;iterable&quot;" expression="kotlinIterableVariable()" name="ITERABLE" />
+        <variable alwaysStopAt="true" defaultValue="&quot;iterable&quot;" expression="kotlinVariable()" name="ITERABLE" />
         <variable alwaysStopAt="true" defaultValue="&quot;i&quot;" expression="kotlinSuggestVariableName()" name="VAR" />
         <context>
             <option name="KOTLIN_STATEMENT" value="true" />
diff --git a/idea/src/META-INF/plugin.xml b/idea/src/META-INF/plugin.xml
index aa073396e5d..41ac385fc6a 100644
--- a/idea/src/META-INF/plugin.xml
+++ b/idea/src/META-INF/plugin.xml
@@ -426,9 +426,9 @@
     <liveTemplateContext implementation="org.jetbrains.kotlin.idea.liveTemplates.KotlinTemplateContextType$Comment"/>
     <defaultLiveTemplatesProvider implementation="org.jetbrains.kotlin.idea.liveTemplates.KotlinLiveTemplatesProvider"/>
     <liveTemplateMacro implementation="org.jetbrains.kotlin.idea.liveTemplates.macro.AnyVariableMacro"/>
+    <liveTemplateMacro implementation="org.jetbrains.kotlin.idea.liveTemplates.macro.SuitableVariableMacro"/>
     <liveTemplateMacro implementation="org.jetbrains.kotlin.idea.liveTemplates.macro.FunctionParametersMacro"/>
     <liveTemplateMacro implementation="org.jetbrains.kotlin.idea.liveTemplates.macro.AnonymousSuperMacro"/>
-    <liveTemplateMacro implementation="org.jetbrains.kotlin.idea.liveTemplates.macro.IterableVariableMacro"/>
     <liveTemplateMacro implementation="org.jetbrains.kotlin.idea.liveTemplates.macro.SuggestVariableNameMacro"/>
     <liveTemplateOptionalProcessor implementation="org.jetbrains.kotlin.idea.liveTemplates.KotlinShortenFQNamesProcessor"/>
 
diff --git a/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/AnyVariableMacro.kt b/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/AnyVariableMacro.kt
index f167fe22202..54baa45e6aa 100644
--- a/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/AnyVariableMacro.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/AnyVariableMacro.kt
@@ -21,8 +21,8 @@ import com.intellij.openapi.util.UserDataHolder
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 
 class AnyVariableMacro : BaseKotlinVariableMacro() {
-    override fun getName() = "kotlinVariable"
-    override fun getPresentableName() = "kotlinVariable()"
+    override fun getName() = "kotlinAnyVariable"
+    override fun getPresentableName() = "kotlinAnyVariable()"
 
     override fun isSuitable(variableDescriptor: VariableDescriptor, project: Project, userData: UserDataHolder) = true
 }
diff --git a/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/IterableVariableMacro.kt b/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/SuitableVariableMacro.kt
similarity index 52%
rename from idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/IterableVariableMacro.kt
rename to idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/SuitableVariableMacro.kt
index a9e06cd4986..736be8162a8 100644
--- a/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/IterableVariableMacro.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/SuitableVariableMacro.kt
@@ -21,35 +21,46 @@ import com.intellij.openapi.util.Key
 import com.intellij.openapi.util.UserDataHolder
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
-import org.jetbrains.kotlin.idea.core.IterableTypesDetection
-import org.jetbrains.kotlin.idea.core.IterableTypesDetector
+import org.jetbrains.kotlin.idea.core.ExpectedInfo
+import org.jetbrains.kotlin.idea.core.ExpectedInfos
 import org.jetbrains.kotlin.idea.core.SmartCastCalculator
+import org.jetbrains.kotlin.idea.util.CallTypeAndReceiver
+import org.jetbrains.kotlin.idea.util.FuzzyType
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtNameReferenceExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 
-class IterableVariableMacro : BaseKotlinVariableMacro() {
+class SuitableVariableMacro : BaseKotlinVariableMacro() {
     private companion object {
-        val ITERABLE_TYPES_DETECTOR_KEY = Key<IterableTypesDetector>("ITERABLE_TYPES_DETECTOR_KEY")
+        val EXPECTED_INFOS_KEY = Key<Collection<ExpectedInfo>>("EXPECTED_INFOS_KEY")
         val SMART_CAST_CALCULATOR_KEY = Key<SmartCastCalculator>("SMART_CAST_CALCULATOR_KEY")
     }
 
-    override fun getName() = "kotlinIterableVariable"
-    override fun getPresentableName() = "kotlinIterableVariable()"
+    override fun getName() = "kotlinVariable"
+    override fun getPresentableName() = "kotlinVariable()"
 
     override fun initUserData(userData: UserDataHolder, contextElement: KtElement, bindingContext: BindingContext) {
         val resolutionFacade = contextElement.getResolutionFacade()
-        val scope = contextElement.getResolutionScope(bindingContext, resolutionFacade)
-        val detector = resolutionFacade.getIdeService(IterableTypesDetection::class.java).createDetector(scope)
-        val smartCastCalculator = SmartCastCalculator(bindingContext, scope.ownerDescriptor, contextElement, null, resolutionFacade)
-        userData.putUserData(ITERABLE_TYPES_DETECTOR_KEY, detector)
-        userData.putUserData(SMART_CAST_CALCULATOR_KEY, smartCastCalculator)
+        if (contextElement is KtNameReferenceExpression) {
+            val callTypeAndReceiver = CallTypeAndReceiver.detect(contextElement)
+            if (callTypeAndReceiver is CallTypeAndReceiver.DEFAULT) {
+                val expectedInfos = ExpectedInfos(bindingContext, resolutionFacade).calculate(contextElement)
+                if (expectedInfos.isNotEmpty()) {
+                    userData.putUserData(EXPECTED_INFOS_KEY, expectedInfos)
+
+                    val scope = contextElement.getResolutionScope(bindingContext, resolutionFacade)
+                    val smartCastCalculator = SmartCastCalculator(bindingContext, scope.ownerDescriptor, contextElement, null, resolutionFacade)
+                    userData.putUserData(SMART_CAST_CALCULATOR_KEY, smartCastCalculator)
+                }
+            }
+        }
     }
 
     override fun isSuitable(variableDescriptor: VariableDescriptor, project: Project, userData: UserDataHolder): Boolean {
-        val detector = userData.getUserData(ITERABLE_TYPES_DETECTOR_KEY)!!
+        val expectedInfos = userData.getUserData(EXPECTED_INFOS_KEY) ?: return true
         val smartCastCalculator = userData.getUserData(SMART_CAST_CALCULATOR_KEY)!!
         val types = smartCastCalculator.types(variableDescriptor)
-        return types.any { detector.isIterable(it, null) }
+        return expectedInfos.any { expectedInfo -> types.any { expectedInfo.filter.matchingSubstitutor(FuzzyType(it, emptyList())) != null } }
     }
 }