Refactored to be more type safe

idea/idea-live-templates/src/org/jetbrains/kotlin/idea/liveTemplates/macro/AnyVariableMacro.kt

@@ -17,12 +17,16 @@
 package org.jetbrains.kotlin.idea.liveTemplates.macro
 
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.UserDataHolder
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.resolve.BindingContext
 
-class AnyVariableMacro : BaseKotlinVariableMacro() {
+class AnyVariableMacro : BaseKotlinVariableMacro<Unit>() {
     override fun getName() = "kotlinAnyVariable"
     override fun getPresentableName() = "kotlinAnyVariable()"
 
-    override fun isSuitable(variableDescriptor: VariableDescriptor, project: Project, userData: UserDataHolder) = true
+    override fun initState(contextElement: KtElement, bindingContext: BindingContext) {
+    }
+
+    override fun isSuitable(variableDescriptor: VariableDescriptor, project: Project, state: Unit) = true
 }

idea/idea-live-templates/src/org/jetbrains/kotlin/idea/liveTemplates/macro/BaseKotlinVariableMacro.kt

@@ -19,7 +19,6 @@ package org.jetbrains.kotlin.idea.liveTemplates.macro
 import com.intellij.codeInsight.lookup.LookupElement
 import com.intellij.codeInsight.template.*
 import com.intellij.openapi.project.Project
-import com.intellij.openapi.util.UserDataHolder
 import com.intellij.psi.PsiDocumentManager
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptorWithVisibility
@@ -38,9 +37,8 @@ import org.jetbrains.kotlin.renderer.render
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.scopes.DescriptorKindFilter
-import org.jetbrains.kotlin.util.slicedMap.UserDataHolderImpl
 
-abstract class BaseKotlinVariableMacro : Macro() {
+abstract class BaseKotlinVariableMacro<TState> : Macro() {
     private fun getVariables(params: Array<Expression>, context: ExpressionContext): Collection<VariableDescriptor> {
         if (params.size != 0) return emptyList()
 
@@ -60,23 +58,21 @@ abstract class BaseKotlinVariableMacro : Macro() {
             return descriptor !is DeclarationDescriptorWithVisibility || descriptor.isVisible(contextElement, null, bindingContext, resolutionFacade)
         }
 
-        val userData = UserDataHolderImpl()
+        val state = initState(contextElement, bindingContext)
-        initUserData(userData, contextElement, bindingContext)
 
         val helper = ReferenceVariantsHelper(bindingContext, resolutionFacade, resolutionFacade.moduleDescriptor, ::isVisible)
         return helper
                 .getReferenceVariants(contextElement, CallTypeAndReceiver.DEFAULT, DescriptorKindFilter.VARIABLES, { true })
                 .map { it as VariableDescriptor }
-                .filter { isSuitable(it, project, userData) }
+                .filter { isSuitable(it, project, state) }
     }
 
-    protected open fun initUserData(userData: UserDataHolder, contextElement: KtElement, bindingContext: BindingContext) {
+    protected abstract fun initState(contextElement: KtElement, bindingContext: BindingContext): TState
-    }
 
     protected abstract fun isSuitable(
             variableDescriptor: VariableDescriptor,
             project: Project,
-            userData: UserDataHolder): Boolean
+            state: TState): Boolean
 
     override fun calculateResult(params: Array<Expression>, context: ExpressionContext): Result? {
         val vars = getVariables(params, context)

idea/idea-live-templates/src/org/jetbrains/kotlin/idea/liveTemplates/macro/SuitableVariableMacro.kt

@@ -17,8 +17,6 @@
 package org.jetbrains.kotlin.idea.liveTemplates.macro
 
 import com.intellij.openapi.project.Project
-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.ExpectedInfo
@@ -31,36 +29,32 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtNameReferenceExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 
-class SuitableVariableMacro : BaseKotlinVariableMacro() {
+class SuitableVariableMacro : BaseKotlinVariableMacro<SuitableVariableMacro.State?>() {
-    private companion object {
+    class State(val expectedInfos: Collection<ExpectedInfo>, val smartCastCalculator: SmartCastCalculator)
-        val EXPECTED_INFOS_KEY = Key<Collection<ExpectedInfo>>("EXPECTED_INFOS_KEY")
-        val SMART_CAST_CALCULATOR_KEY = Key<SmartCastCalculator>("SMART_CAST_CALCULATOR_KEY")
-    }
 
     override fun getName() = "kotlinVariable"
     override fun getPresentableName() = "kotlinVariable()"
 
-    override fun initUserData(userData: UserDataHolder, contextElement: KtElement, bindingContext: BindingContext) {
+    override fun initState(contextElement: KtElement, bindingContext: BindingContext): State? {
         val resolutionFacade = contextElement.getResolutionFacade()
         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)
+                    return State(expectedInfos, smartCastCalculator)
                 }
             }
         }
+
+        return null
     }
 
-    override fun isSuitable(variableDescriptor: VariableDescriptor, project: Project, userData: UserDataHolder): Boolean {
+    override fun isSuitable(variableDescriptor: VariableDescriptor, project: Project, state: State?): Boolean {
-        val expectedInfos = userData.getUserData(EXPECTED_INFOS_KEY) ?: return true
+        if (state == null) return true
-        val smartCastCalculator = userData.getUserData(SMART_CAST_CALCULATOR_KEY)!!
+        val types = state.smartCastCalculator.types(variableDescriptor)
-        val types = smartCastCalculator.types(variableDescriptor)
+        return state.expectedInfos.any { expectedInfo -> types.any { expectedInfo.filter.matchingSubstitutor(FuzzyType(it, emptyList())) != null } }
-        return expectedInfos.any { expectedInfo -> types.any { expectedInfo.filter.matchingSubstitutor(FuzzyType(it, emptyList())) != null } }
     }
 }