IterableVariableMacro to honor smart casts

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/smart/MultipleArgumentsItemProvider.kt

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.idea.KotlinDescriptorIconProvider
 import org.jetbrains.kotlin.idea.completion.ArgumentPositionData
 import org.jetbrains.kotlin.idea.completion.ExpectedInfo
-import org.jetbrains.kotlin.idea.completion.SmartCastCalculator
+import org.jetbrains.kotlin.idea.core.SmartCastCalculator
 import org.jetbrains.kotlin.idea.completion.Tail
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.util.getResolutionScope

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/smart/SmartCompletion.kt

@@ -27,6 +27,7 @@ import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
 import org.jetbrains.kotlin.idea.completion.*
 import org.jetbrains.kotlin.idea.completion.handlers.WithTailInsertHandler
+import org.jetbrains.kotlin.idea.core.SmartCastCalculator
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.util.CallTypeAndReceiver
 import org.jetbrains.kotlin.idea.util.FuzzyType
@@ -81,7 +82,7 @@ class SmartCompletion(
     private val callableTypeExpectedInfo = expectedInfos.filterCallableExpected()
 
     public val smartCastCalculator: SmartCastCalculator by lazy(LazyThreadSafetyMode.NONE) {
-        SmartCastCalculator(bindingContext, resolutionFacade.moduleDescriptor, expression, resolutionFacade)
+        SmartCastCalculator(bindingContext, resolutionFacade.moduleDescriptor, expression, callTypeAndReceiver.receiver as? KtExpression, resolutionFacade)
     }
 
     public val descriptorFilter: ((DeclarationDescriptor, AbstractLookupElementFactory) -> Collection<LookupElement>)? =

idea/idea-completion/src/org/jetbrains/kotlin/idea/completion/smart/Utils.kt

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.descriptors.impl.LocalVariableDescriptor
 import org.jetbrains.kotlin.idea.completion.*
 import org.jetbrains.kotlin.idea.completion.handlers.WithExpressionPrefixInsertHandler
 import org.jetbrains.kotlin.idea.completion.handlers.WithTailInsertHandler
+import org.jetbrains.kotlin.idea.core.SmartCastCalculator
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.util.*
 import org.jetbrains.kotlin.resolve.callableReferences.getReflectionTypeForCandidateDescriptor

idea/idea-core/src/org/jetbrains/kotlin/idea/core/SmartCastCalculator.kt

@@ -14,19 +14,18 @@
  * limitations under the License.
  */
 
-package org.jetbrains.kotlin.idea.completion
+package org.jetbrains.kotlin.idea.core
 
 import com.intellij.openapi.util.Pair
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.descriptors.ClassifierDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.ReceiverParameterDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
-import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.util.getImplicitReceiversWithInstance
+import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtSimpleNameExpression
-import org.jetbrains.kotlin.psi.psiUtil.getReceiverExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
@@ -43,15 +42,14 @@ import java.util.*
 class SmartCastCalculator(
         val bindingContext: BindingContext,
         val containingDeclarationOrModule: DeclarationDescriptor,
-        expression: KtExpression,
+        contextElement: PsiElement,
+        receiver: KtExpression?,
         resolutionFacade: ResolutionFacade
 ) {
-    private val receiver = if (expression is KtSimpleNameExpression) expression.getReceiverExpression() else null
-
     // keys are VariableDescriptor's and ThisReceiver's
     private val entityToSmartCastInfo: Map<Any, SmartCastInfo> = processDataFlowInfo(
-            bindingContext.getDataFlowInfo(expression),
+            bindingContext.getDataFlowInfo(contextElement),
-            expression.getResolutionScope(bindingContext, resolutionFacade),
+            contextElement.getResolutionScope(bindingContext, resolutionFacade),
             receiver)
 
     fun types(descriptor: VariableDescriptor): Collection<KotlinType> {
@@ -114,7 +112,7 @@ class SmartCastCalculator(
 
         val entityToInfo = HashMap<Any, SmartCastInfo>()
 
-        for ((dataFlowValue, types) in dataFlowInfo.completeTypeInfo.asMap().entrySet()) {
+        for ((dataFlowValue, types) in dataFlowInfo.completeTypeInfo.asMap().entries) {
             val entity = dataFlowValueToEntity.invoke(dataFlowValue)
             if (entity != null) {
                 entityToInfo[entity] = SmartCastInfo(types, false)

idea/src/org/jetbrains/kotlin/idea/liveTemplates/macro/IterableVariableMacro.kt

@@ -23,6 +23,7 @@ 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.SmartCastCalculator
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.resolve.BindingContext
@@ -30,6 +31,7 @@ import org.jetbrains.kotlin.resolve.BindingContext
 class IterableVariableMacro : BaseKotlinVariableMacro() {
     private companion object {
         val ITERABLE_TYPES_DETECTOR_KEY = Key<IterableTypesDetector>("ITERABLE_TYPES_DETECTOR_KEY")
+        val SMART_CAST_CALCULATOR_KEY = Key<SmartCastCalculator>("SMART_CAST_CALCULATOR_KEY")
     }
 
     override fun getName() = "kotlinIterableVariable"
@@ -39,12 +41,15 @@ class IterableVariableMacro : BaseKotlinVariableMacro() {
         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)
     }
 
     override fun isSuitable(variableDescriptor: VariableDescriptor, project: Project, userData: UserDataHolder): Boolean {
         val detector = userData.getUserData(ITERABLE_TYPES_DETECTOR_KEY)!!
-        //TODO: smart-casts
+        val smartCastCalculator = userData.getUserData(SMART_CAST_CALCULATOR_KEY)!!
-        return detector.isIterable(variableDescriptor.type, null)
+        val types = smartCastCalculator.types(variableDescriptor)
+        return types.any { detector.isIterable(it, null) }
     }
 }

idea/testData/templates/iter.exp.kt

@@ -7,12 +7,14 @@ class MyClass {
     public var collection : HashSet<Int>? = null
     private var isAlive : Boolean = false
 
-    fun main(args : Array<String>, v : Int) {
+    fun main(args : Array<String>, v : Int, o: Any) {
         var str = ""
         val myList = ArrayList<String>()
         val stream = FileInputStream(".")
-        for (collection in args) {
+        if (o is String) {
-            <caret>
+            for (collection in args) {
+                <caret>
+            }
         }
     }
 }

idea/testData/templates/iter.kt

@@ -7,10 +7,12 @@ class MyClass {
     public var collection : HashSet<Int>? = null
     private var isAlive : Boolean = false
 
-    fun main(args : Array<String>, v : Int) {
+    fun main(args : Array<String>, v : Int, o: Any) {
         var str = ""
         val myList = ArrayList<String>()
         val stream = FileInputStream(".")
-        <caret>
+        if (o is String) {
+            <caret>
+        }
     }
 }

idea/tests/org/jetbrains/kotlin/idea/codeInsight/LiveTemplatesTest.java

@@ -198,7 +198,7 @@ public class LiveTemplatesTest extends KotlinLightCodeInsightFixtureTestCase {
     public void testIter() {
         start();
 
-        assertStringItems("args", "myList", "str", "stream");
+        assertStringItems("args", "myList", "o", "str", "stream");
         type("args");
         nextTab(2);