Make parameter MultiMap

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractFunction/ui/KotlinExtractFunctionDialog.java

@@ -42,10 +42,8 @@ import javax.swing.*;
 import java.awt.*;
 import java.awt.event.ItemEvent;
 import java.awt.event.ItemListener;
-import java.util.ArrayList;
+import java.util.*;
-import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 
 public class KotlinExtractFunctionDialog extends DialogWrapper {
     private JPanel contentPane;
@@ -306,22 +304,26 @@ public class KotlinExtractFunctionDialog extends DialogWrapper {
         }
         controlFlow = new ControlFlow(outputValues, controlFlow.getBoxerFactory(), controlFlow.getDeclarationsToCopy());
 
-        Map<Integer, Replacement> replacementMap = ContainerUtil.newHashMap();
+        MultiMap<Integer, Replacement> replacementMap = MultiMap.<Integer, Replacement>create();
-        for (Map.Entry<Integer, Replacement> e : originalDescriptor.getReplacementMap().entrySet()) {
+        for (Map.Entry<Integer, Collection<Replacement>> e : originalDescriptor.getReplacementMap().entrySet()) {
             Integer offset = e.getKey();
-            Replacement replacement = e.getValue();
+            Collection<Replacement> replacements = e.getValue();
 
-            if (replacement instanceof ParameterReplacement) {
+            for (Replacement replacement : replacements) {
-                ParameterReplacement parameterReplacement = (ParameterReplacement) replacement;
+                if (replacement instanceof ParameterReplacement) {
-                Parameter parameter = parameterReplacement.getParameter();
+                    ParameterReplacement parameterReplacement = (ParameterReplacement) replacement;
+                    Parameter parameter = parameterReplacement.getParameter();
 
-                Parameter newParameter = oldToNewParameters.get(parameter);
+                    Parameter newParameter = oldToNewParameters.get(parameter);
-                if (newParameter != null) {
+                    if (newParameter != null) {
-                    replacementMap.put(offset, parameterReplacement.copy(newParameter));
+                        replacementMap.remove(offset, replacement);
+                        replacementMap.putValue(offset, parameterReplacement.copy(newParameter));
+                    }
                 }
-            }
+                else {
-            else {
+                    replacementMap.put(offset, replacements);
-                replacementMap.put(offset, replacement);
+                }
+
             }
         }
 

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/ExtractableCodeDescriptor.kt

@@ -28,9 +28,7 @@ import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.refactoring.KotlinRefactoringBundle
 import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.AnalysisResult.ErrorMessage
 import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.AnalysisResult.Status
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.ExpressionValue
+import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.*
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.Initializer
-import org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine.OutputValue.Jump
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference.ShorteningMode
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.idea.util.approximateFlexibleTypes
@@ -49,7 +47,7 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.resolveTopLevelClass
 import org.jetbrains.kotlin.types.*
 import org.jetbrains.kotlin.types.typeUtil.builtIns
 import org.jetbrains.kotlin.types.typeUtil.isUnit
-import java.util.Collections
+import java.util.*
 
 interface Parameter {
     val argumentText: String
@@ -326,7 +324,7 @@ data class ExtractableCodeDescriptor(
         val parameters: List<Parameter>,
         val receiverParameter: Parameter?,
         val typeParameters: List<TypeParameter>,
-        val replacementMap: Map<Int, Replacement>,
+        val replacementMap: MultiMap<Int, Replacement>,
         val controlFlow: ControlFlow,
         val returnType: KotlinType
 ) {
@@ -424,7 +422,7 @@ data class ExtractionResult(
         val config: ExtractionGeneratorConfiguration,
         val declaration: KtNamedDeclaration,
         val duplicateReplacers: Map<KotlinPsiRange, () -> Unit>,
-        val nameByOffset: Map<Int, KtElement>
+        val nameByOffset: MultiMap<Int, KtElement>
 )
 
 class AnalysisResult (

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/extractableAnalysisUtil.kt

@@ -570,8 +570,8 @@ private class DelegatingParameter(
 
 internal class ParametersInfo {
     var errorMessage: ErrorMessage? = null
-    val replacementMap: MutableMap<Int, Replacement> = HashMap()
+    val replacementMap: MultiMap<Int, Replacement> = MultiMap.create()
-    val originalRefToParameter: MutableMap<KtSimpleNameExpression, MutableParameter> = HashMap()
+    val originalRefToParameter: MultiMap<KtSimpleNameExpression, MutableParameter> = MultiMap.create()
     val parameters: MutableSet<MutableParameter> = HashSet()
     val typeParameters: MutableSet<TypeParameter> = HashSet()
     val nonDenotableTypes: MutableSet<KotlinType> = HashSet()
@@ -688,7 +688,7 @@ fun ExtractionData.performAnalysis(): AnalysisResult {
     controlFlow.jumpOutputValue?.elementToInsertAfterCall?.accept(
             object : KtTreeVisitorVoid() {
                 override fun visitSimpleNameExpression(expression: KtSimpleNameExpression) {
-                    paramsInfo.originalRefToParameter[expression]?.let { it.refCount-- }
+                    paramsInfo.originalRefToParameter[expression].firstOrNull()?.let { it.refCount-- }
                 }
             }
     )
@@ -773,46 +773,50 @@ fun ExtractableCodeDescriptor.validate(): ExtractableCodeDescriptorWithConflicts
     val body = result.declaration.getGeneratedBody()
     val bindingContext = body.analyzeFully()
 
+    fun processReference(resolveResult: ResolveResult, currentRefExpr: KtReferenceExpression) {
+        if (currentRefExpr.parent is KtThisExpression) return
+
+        val diagnostics = bindingContext.diagnostics.forElement(currentRefExpr)
+
+        val currentDescriptor = bindingContext[BindingContext.REFERENCE_TARGET, currentRefExpr]
+        val currentTarget =
+                currentDescriptor?.let { DescriptorToSourceUtilsIde.getAnyDeclaration(extractionData.project, it) } as? PsiNamedElement
+        if (currentTarget is KtParameter && currentTarget.parent == valueParameterList) return
+        if (currentTarget is KtTypeParameter && currentTarget.parent == typeParameterList) return
+        if (currentDescriptor is LocalVariableDescriptor
+            && parameters.any { it.mirrorVarName == currentDescriptor.name.asString() }) return
+
+        if (diagnostics.any { it.factory in Errors.UNRESOLVED_REFERENCE_DIAGNOSTICS }
+            || (currentDescriptor != null
+                && !ErrorUtils.isError(currentDescriptor)
+                && !compareDescriptors(extractionData.project, currentDescriptor, resolveResult.descriptor))) {
+            conflicts.putValue(
+                    resolveResult.originalRefExpr,
+                    getDeclarationMessage(resolveResult.declaration, "0.will.no.longer.be.accessible.after.extraction")
+            )
+            return
+        }
+
+        diagnostics.firstOrNull { it.factory in Errors.INVISIBLE_REFERENCE_DIAGNOSTICS }?.let {
+            val message = when (it.factory) {
+                Errors.INVISIBLE_SETTER ->
+                    getDeclarationMessage(resolveResult.declaration, "setter.of.0.will.become.invisible.after.extraction", false)
+                else ->
+                    getDeclarationMessage(resolveResult.declaration, "0.will.become.invisible.after.extraction")
+            }
+            conflicts.putValue(resolveResult.originalRefExpr, message)
+        }
+    }
+
     fun validateBody() {
         for ((originalOffset, resolveResult) in extractionData.refOffsetToDeclaration) {
             if (resolveResult.declaration.isInsideOf(extractionData.originalElements)) continue
 
-            val currentRefExpr = result.nameByOffset[originalOffset]?.let {
+            val currentRefExprs = result.nameByOffset[originalOffset].map {
                 (it as? KtThisExpression)?.instanceReference ?: it as? KtSimpleNameExpression
-            } ?: continue
+            }.filterNotNull()
 
-            if (currentRefExpr.parent is KtThisExpression) continue
+            currentRefExprs.forEach { processReference(resolveResult, it) }
-
-            val diagnostics = bindingContext.diagnostics.forElement(currentRefExpr)
-
-            val currentDescriptor = bindingContext[BindingContext.REFERENCE_TARGET, currentRefExpr]
-            val currentTarget =
-                    currentDescriptor?.let { DescriptorToSourceUtilsIde.getAnyDeclaration(extractionData.project, it) } as? PsiNamedElement
-            if (currentTarget is KtParameter && currentTarget.parent == valueParameterList) continue
-            if (currentTarget is KtTypeParameter && currentTarget.parent == typeParameterList) continue
-            if (currentDescriptor is LocalVariableDescriptor
-                && parameters.any { it.mirrorVarName == currentDescriptor.name.asString() }) continue
-
-            if (diagnostics.any { it.factory in Errors.UNRESOLVED_REFERENCE_DIAGNOSTICS }
-                || (currentDescriptor != null
-                    && !ErrorUtils.isError(currentDescriptor)
-                    && !compareDescriptors(extractionData.project, currentDescriptor, resolveResult.descriptor))) {
-                conflicts.putValue(
-                        resolveResult.originalRefExpr,
-                        getDeclarationMessage(resolveResult.declaration, "0.will.no.longer.be.accessible.after.extraction")
-                )
-                continue
-            }
-
-            diagnostics.firstOrNull { it.factory in Errors.INVISIBLE_REFERENCE_DIAGNOSTICS }?.let {
-                val message = when (it.factory) {
-                    Errors.INVISIBLE_SETTER ->
-                        getDeclarationMessage(resolveResult.declaration, "setter.of.0.will.become.invisible.after.extraction", false)
-                    else ->
-                        getDeclarationMessage(resolveResult.declaration, "0.will.become.invisible.after.extraction")
-                }
-                conflicts.putValue(resolveResult.originalRefExpr, message)
-            }
         }
     }
 

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/extractorUtil.kt

@@ -21,6 +21,7 @@ import com.intellij.psi.PsiFile
 import com.intellij.psi.search.LocalSearchScope
 import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.refactoring.BaseRefactoringProcessor
+import com.intellij.util.containers.MultiMap
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
 import org.jetbrains.kotlin.idea.core.*
@@ -405,7 +406,7 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
         declarationToReplace: KtNamedDeclaration? = null
 ): ExtractionResult{
     val psiFactory = KtPsiFactory(descriptor.extractionData.originalFile)
-    val nameByOffset = HashMap<Int, KtElement>()
+    val nameByOffset = MultiMap<Int, KtElement>()
 
     fun createDeclaration(): KtNamedDeclaration {
         return with(descriptor.extractionData) {
@@ -451,7 +452,13 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
         }
 
         val counterpartMap = createNameCounterpartMap(currentResultExpression, expressionToUnifyWith ?: newResultExpression)
-        nameByOffset.entries.forEach { e -> counterpartMap.getRaw(e.value)?.let { e.setValue(it) } }
+        counterpartMap.entries.forEach {
+            val (cmOriginalExpr, cmNewExpr) = it
+            nameByOffset.entrySet().find { it.value.containsRaw(cmOriginalExpr) }?.let {
+                nameByOffset.remove(it.key, cmOriginalExpr)
+                nameByOffset.putValue(it.key, cmNewExpr)
+            }
+        }
     }
 
     fun getCounterparts<T : KtExpression>(originalExpressions: Collection<T>,
@@ -468,7 +475,7 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
     fun adjustDeclarationBody(declaration: KtNamedDeclaration) {
         val body = declaration.getGeneratedBody()
 
-        val exprReplacementMap = HashMap<KtElement, (ExtractableCodeDescriptor, KtElement) -> KtElement>()
+        val exprReplacementMap = MultiMap<KtElement, (ExtractableCodeDescriptor, KtElement) -> KtElement>()
         val originalOffsetByExpr = LinkedHashMap<KtElement, Int>()
 
         val bodyOffset = body.getBlockContentOffset()
@@ -484,8 +491,10 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
 
             originalOffsetByExpr[expr!!] = offsetInBody
 
-            descriptor.replacementMap[offsetInBody]?.let { replacement ->
+            descriptor.replacementMap[offsetInBody].let { replacements ->
-                exprReplacementMap[expr] = replacement
+                replacements.forEach {
+                    exprReplacementMap.putValue(expr, it)
+                }
             }
         }
 
@@ -516,7 +525,8 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
 
         for ((expr, originalOffset) in originalOffsetByExpr) {
             if (expr.isValid) {
-                nameByOffset.put(originalOffset, exprReplacementMap[expr]?.invoke(descriptor, expr) ?: expr)
+                val replacements = exprReplacementMap[expr].map { it?.invoke(descriptor, expr) }.filterNotNull()
+                nameByOffset.put(originalOffset, if (replacements.isEmpty()) arrayListOf(expr) else replacements)
             }
         }
 

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/inferParameterInfo.kt

@@ -165,7 +165,7 @@ private fun extractReceiver(
         )) return
 
         if (referencedClassifierDescriptor is ClassDescriptor) {
-            info.replacementMap[refInfo.offsetInBody] = FqNameReplacement(originalDescriptor.getImportableDescriptor().fqNameSafe)
+            info.replacementMap.putValue(refInfo.offsetInBody, FqNameReplacement(originalDescriptor.getImportableDescriptor().fqNameSafe))
         }
     }
     else {
@@ -230,7 +230,7 @@ private fun extractReceiver(
             }
 
             parameter.refCount++
-            info.originalRefToParameter[originalRef] = parameter
+            info.originalRefToParameter.putValue(originalRef, parameter)
 
             parameter.addDefaultType(parameterType)
 
@@ -251,11 +251,11 @@ private fun extractReceiver(
                 }
             }
 
-            info.replacementMap[refInfo.offsetInBody] =
+            info.replacementMap.putValue(refInfo.offsetInBody,
                     when {
                         hasThisReceiver && extractThis -> AddPrefixReplacement(parameter)
                         else -> RenameReplacement(parameter)
-                    }
+                    })
         }
     }
 }