IterableVariableMacro to honor smart casts

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

@@ -1,139 +0,0 @@
-/*
- * Copyright 2010-2015 JetBrains s.r.o.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package org.jetbrains.kotlin.idea.completion
-
-import com.intellij.openapi.util.Pair
-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.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
-import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue
-import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
-import org.jetbrains.kotlin.resolve.calls.smartcasts.Nullability
-import org.jetbrains.kotlin.resolve.scopes.LexicalScope
-import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitReceiver
-import org.jetbrains.kotlin.types.KotlinType
-import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
-import org.jetbrains.kotlin.types.typeUtil.makeNotNullable
-import java.util.*
-
-class SmartCastCalculator(
-        val bindingContext: BindingContext,
-        val containingDeclarationOrModule: DeclarationDescriptor,
-        expression: 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),
-            expression.getResolutionScope(bindingContext, resolutionFacade),
-            receiver)
-
-    fun types(descriptor: VariableDescriptor): Collection<KotlinType> {
-        val type = descriptor.returnType ?: return emptyList()
-        return entityType(descriptor, type)
-    }
-
-    fun types(thisReceiverParameter: ReceiverParameterDescriptor): Collection<KotlinType> {
-        val type = thisReceiverParameter.type
-        val thisReceiver = thisReceiverParameter.value as? ImplicitReceiver ?: return listOf(type)
-        return entityType(thisReceiver, type)
-    }
-
-    private fun entityType(entity: Any, ownType: KotlinType): Collection<KotlinType> {
-        val smartCastInfo = entityToSmartCastInfo[entity] ?: return listOf(ownType)
-
-        var types = smartCastInfo.types + ownType
-
-        if (smartCastInfo.notNull) {
-            types = types.map { it.makeNotNullable() }
-        }
-
-        return types
-    }
-
-    private data class SmartCastInfo(var types: Collection<KotlinType>, var notNull: Boolean) {
-        constructor() : this(emptyList(), false)
-    }
-
-    private fun processDataFlowInfo(dataFlowInfo: DataFlowInfo, resolutionScope: LexicalScope?, receiver: KtExpression?): Map<Any, SmartCastInfo> {
-        if (dataFlowInfo == DataFlowInfo.EMPTY) return emptyMap()
-
-        val dataFlowValueToEntity: (DataFlowValue) -> Any?
-        if (receiver != null) {
-            val receiverType = bindingContext.getType(receiver) ?: return emptyMap()
-            val receiverId = DataFlowValueFactory.createDataFlowValue(receiver, receiverType, bindingContext, containingDeclarationOrModule).id
-            dataFlowValueToEntity = { value ->
-                val id = value.id
-                if (id is Pair<*, *> && id.first == receiverId) id.second as? VariableDescriptor else null
-            }
-        }
-        else {
-            dataFlowValueToEntity = fun (value: DataFlowValue): Any? {
-                val id = value.id
-                when(id) {
-                    is VariableDescriptor, is ImplicitReceiver -> return id
-
-                    is Pair<*, *> -> {
-                        val first = id.first
-                        val second = id.second
-                        if (first !is ImplicitReceiver || second !is VariableDescriptor) return null
-                        if (resolutionScope?.findNearestReceiverForVariable(second)?.value != first) return null
-                        return second
-                    }
-
-                    else -> return null
-                }
-            }
-        }
-
-        val entityToInfo = HashMap<Any, SmartCastInfo>()
-
-        for ((dataFlowValue, types) in dataFlowInfo.completeTypeInfo.asMap().entrySet()) {
-            val entity = dataFlowValueToEntity.invoke(dataFlowValue)
-            if (entity != null) {
-                entityToInfo[entity] = SmartCastInfo(types, false)
-            }
-        }
-
-        for ((dataFlowValue, nullability) in dataFlowInfo.completeNullabilityInfo) {
-            if (nullability == Nullability.NOT_NULL) {
-                val entity = dataFlowValueToEntity(dataFlowValue) ?: continue
-                entityToInfo.getOrPut(entity, { SmartCastInfo() }).notNull = true
-            }
-        }
-
-        return entityToInfo
-    }
-
-    private fun LexicalScope.findNearestReceiverForVariable(variableDescriptor: VariableDescriptor): ReceiverParameterDescriptor? {
-        val classifier = variableDescriptor.containingDeclaration as? ClassifierDescriptor ?: return null
-        val type = classifier.defaultType
-        return getImplicitReceiversWithInstance().firstOrNull { it.type.isSubtypeOf(type) }
-    }
-}

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