Extraction Engine: Get rid of explicit element offset manipulation and use copyable user data instead

compiler/frontend/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt

@@ -560,6 +560,8 @@ public class KtPsiFactory(private val project: Project) {
             return this
         }
 
+        public fun transform(f: StringBuilder.() -> Unit) = sb.f()
+
         public fun asString(): String {
             if (state != State.DONE) {
                 state = State.DONE

compiler/frontend/src/org/jetbrains/kotlin/psi/userDataUtil.kt

@@ -21,13 +21,13 @@ import com.intellij.openapi.util.UserDataHolder
 import com.intellij.psi.PsiElement
 import kotlin.reflect.KProperty
 
-public class UserDataProperty<in R: UserDataHolder, T : Any>(val key: Key<T>) {
+public class UserDataProperty<in R : UserDataHolder, T : Any>(val key: Key<T>) {
     operator fun getValue(thisRef: R, desc: KProperty<*>) = thisRef.getUserData(key)
 
     operator fun setValue(thisRef: R, desc: KProperty<*>, value: T?) = thisRef.putUserData(key, value)
 }
 
-public class NotNullableUserDataProperty<in R: UserDataHolder, T : Any>(val key: Key<T>, val defaultValue: T) {
+public class NotNullableUserDataProperty<in R : UserDataHolder, T : Any>(val key: Key<T>, val defaultValue: T) {
     operator fun getValue(thisRef: R, desc: KProperty<*>) = thisRef.getUserData(key) ?: defaultValue
 
     operator fun setValue(thisRef: R, desc: KProperty<*>, value: T) {
@@ -35,16 +35,25 @@ public class NotNullableUserDataProperty<in R: UserDataHolder, T : Any>(val key:
     }
 }
 
-public class CopyableUserDataProperty<in R: PsiElement, T : Any>(val key: Key<T>) {
+public class CopyableUserDataProperty<in R : PsiElement, T : Any>(val key: Key<T>) {
     operator fun getValue(thisRef: R, property: KProperty<*>) = thisRef.getCopyableUserData(key)
 
     operator fun setValue(thisRef: R, property: KProperty<*>, value: T?) = thisRef.putCopyableUserData(key, value)
 }
 
-public class NotNullableCopyableUserDataProperty<in R: PsiElement, T : Any>(val key: Key<T>, val defaultValue: T) {
+public class NotNullableCopyableUserDataProperty<in R : PsiElement, T : Any>(val key: Key<T>, val defaultValue: T) {
     operator fun getValue(thisRef: R, property: KProperty<*>) = thisRef.getCopyableUserData(key) ?: defaultValue
 
     operator fun setValue(thisRef: R, property: KProperty<*>, value: T) {
         thisRef.putCopyableUserData(key, if (value != defaultValue) value else null)
     }
 }
+
+public class NotNullableCopyableUserDataPropertyWithLazyDefault<in R : PsiElement, T : Any>(val key: Key<T>,
+                                                                                            val computeDefaultValue: () -> T) {
+    private val delegate by lazy { NotNullableCopyableUserDataProperty<R, T>(key, computeDefaultValue()) }
+
+    operator fun getValue(thisRef: R, property: KProperty<*>) = delegate.getValue(thisRef, property)
+
+    operator fun setValue(thisRef: R, property: KProperty<*>, value: T) = delegate.setValue(thisRef, property, value)
+}

idea/src/org/jetbrains/kotlin/idea/refactoring/changeSignature/KotlinChangeSignatureUsageProcessor.kt

@@ -671,7 +671,7 @@ class KotlinChangeSignatureUsageProcessor : ChangeSignatureUsageProcessor {
             }
 
             val psiFactory = KtPsiFactory(callable.project)
-            val tempFile = (callable.containingFile as KtFile).createTempCopy { it }
+            val tempFile = (callable.containingFile as KtFile).createTempCopy()
             val functionWithReceiver = tempFile.findElementAt(callable.textOffset)?.getNonStrictParentOfType<KtNamedFunction>() ?: return
             val receiverTypeRef = psiFactory.createType(newReceiverInfo.currentTypeText)
             functionWithReceiver.setReceiverTypeReference(receiverTypeRef)

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

@@ -113,9 +113,7 @@ public class KotlinExtractFunctionDialog extends DialogWrapper {
 
         setOKActionEnabled(checkNames());
         signaturePreviewField.setText(
-                ExtractorUtilKt.getDeclarationText(getCurrentConfiguration(),
+                ExtractorUtilKt.getSignaturePreview(getCurrentConfiguration(), IdeDescriptorRenderers.SOURCE_CODE_SHORT_NAMES_IN_TYPES)
-                                                   false,
-                                                   IdeDescriptorRenderers.SOURCE_CODE_SHORT_NAMES_IN_TYPES)
         );
     }
 
@@ -285,59 +283,12 @@ public class KotlinExtractFunctionDialog extends DialogWrapper {
         for (KotlinParameterTablePanel.ParameterInfo parameterInfo : newParameterInfos) {
             oldToNewParameters.put(parameterInfo.getOriginalParameter(), parameterInfo.toParameter());
         }
-        ArrayList<Parameter> newParameters = ContainerUtil.newArrayList(oldToNewParameters.values());
-
         Parameter originalReceiver = originalDescriptor.getReceiverParameter();
         Parameter newReceiver = newReceiverInfo != null ? newReceiverInfo.toParameter() : null;
         if (originalReceiver != null && newReceiver != null) {
             oldToNewParameters.put(originalReceiver, newReceiver);
         }
 
-        ControlFlow controlFlow = originalDescriptor.getControlFlow();
+        return ExtractableCodeDescriptorKt.copy(originalDescriptor, newName, newVisibility, oldToNewParameters, newReceiver, returnType);
-        List<OutputValue> outputValues = new ArrayList<OutputValue>(controlFlow.getOutputValues());
-        for (int i = 0; i < outputValues.size(); i++) {
-            OutputValue outputValue = outputValues.get(i);
-            if (outputValue instanceof OutputValue.ParameterUpdate) {
-                OutputValue.ParameterUpdate parameterUpdate = (OutputValue.ParameterUpdate) outputValue;
-                outputValues.set(i, new OutputValue.ParameterUpdate(oldToNewParameters.get(parameterUpdate.getParameter()), parameterUpdate.getOriginalExpressions()));
-            }
-        }
-        controlFlow = new ControlFlow(outputValues, controlFlow.getBoxerFactory(), controlFlow.getDeclarationsToCopy());
-
-        MultiMap<Integer, Replacement> replacementMap = MultiMap.<Integer, Replacement>create();
-        for (Map.Entry<Integer, Collection<Replacement>> e : originalDescriptor.getReplacementMap().entrySet()) {
-            Integer offset = e.getKey();
-            Collection<Replacement> replacements = e.getValue();
-
-            for (Replacement replacement : replacements) {
-                if (replacement instanceof ParameterReplacement) {
-                    ParameterReplacement parameterReplacement = (ParameterReplacement) replacement;
-                    Parameter parameter = parameterReplacement.getParameter();
-
-                    Parameter newParameter = oldToNewParameters.get(parameter);
-                    if (newParameter != null) {
-                        replacementMap.remove(offset, replacement);
-                        replacementMap.putValue(offset, parameterReplacement.copy(newParameter));
-                    }
-                }
-                else {
-                    replacementMap.put(offset, replacements);
-                }
-
-            }
-        }
-
-        return new ExtractableCodeDescriptor(
-                originalDescriptor.getExtractionData(),
-                originalDescriptor.getOriginalContext(),
-                Collections.singletonList(newName),
-                newVisibility,
-                newParameters,
-                newReceiver,
-                originalDescriptor.getTypeParameters(),
-                replacementMap,
-                controlFlow,
-                returnType != null ? returnType : originalDescriptor.getReturnType()
-        );
     }
 }

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

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine
 
+import com.intellij.openapi.Disposable
 import com.intellij.openapi.util.text.StringUtil
 import com.intellij.psi.PsiElement
 import com.intellij.util.containers.ContainerUtil
@@ -97,7 +98,7 @@ class RenameReplacement(override val parameter: Parameter): ParameterReplacement
 class WrapInWithReplacement(override val parameter: Parameter): ParameterReplacement {
     override fun invoke(descriptor: ExtractableCodeDescriptor, e: KtElement): KtElement {
         val call = e.parents.firstIsInstance<KtCallExpression>().getQualifiedExpressionForSelectorOrThis()
-        val replacingExpression = KtPsiFactory(e).createExpressionByPattern("with($0) { $1 }", parameter.name, call.text)
+        val replacingExpression = KtPsiFactory(e).createExpressionByPattern("with($0) { $1 }", parameter.name, call)
         val replace = call.replace(replacingExpression)
         return (replace as KtCallExpression).functionLiteralArguments.first().getFunctionLiteral().bodyExpression!!.statements.first()
     }
@@ -112,7 +113,7 @@ class AddPrefixReplacement(override val parameter: Parameter): ParameterReplacem
         if (descriptor.receiverParameter == parameter) return e
 
         val selector = (e.parent as? KtCallExpression) ?: e
-        val replacingExpression = KtPsiFactory(e).createExpression("${parameter.nameForRef}.${selector.text}")
+        val replacingExpression = KtPsiFactory(e).createExpressionByPattern("${parameter.nameForRef}.$0", selector)
         val newExpr = (selector.replace(replacingExpression) as KtQualifiedExpression).selectorExpression!!
         return (newExpr as? KtCallExpression)?.calleeExpression ?: newExpr
     }
@@ -170,13 +171,13 @@ abstract class OutputValueBoxer(val outputValues: List<OutputValue>) {
 
     abstract val returnType: KotlinType
 
-    protected abstract fun getBoxingExpressionText(arguments: List<String>): String?
+    protected abstract fun getBoxingExpressionPattern(arguments: List<KtExpression>): String?
 
     abstract val boxingRequired: Boolean
 
-    fun getReturnExpression(arguments: List<String>, psiFactory: KtPsiFactory): KtReturnExpression? {
+    fun getReturnExpression(arguments: List<KtExpression>, psiFactory: KtPsiFactory): KtReturnExpression? {
-        val expressionText = getBoxingExpressionText(arguments) ?: return null
+        val expressionPattern = getBoxingExpressionPattern(arguments) ?: return null
-        return psiFactory.createExpression("return $expressionText") as KtReturnExpression
+        return psiFactory.createExpressionByPattern("return $expressionPattern", *arguments.toTypedArray()) as KtReturnExpression
     }
 
     protected abstract fun extractExpressionByIndex(boxedExpression: KtExpression, index: Int): KtExpression?
@@ -230,13 +231,13 @@ abstract class OutputValueBoxer(val outputValues: List<OutputValue>) {
 
         override val boxingRequired: Boolean = outputValues.size() > 1
 
-        override fun getBoxingExpressionText(arguments: List<String>): String? {
+        override fun getBoxingExpressionPattern(arguments: List<KtExpression>): String? {
-            return when (arguments.size()) {
+            return when (arguments.size) {
                 0 -> null
-                1 -> arguments.first()
+                1 -> "$0"
                 else -> {
-                    val constructorName = DescriptorUtils.getFqName(returnType.getConstructor().getDeclarationDescriptor()!!).asString()
+                    val constructorName = DescriptorUtils.getFqName(returnType.constructor.declarationDescriptor!!).asString()
-                    return arguments.joinToString(prefix = "$constructorName(", separator = ", ", postfix = ")")
+                    return arguments.indices.joinToString(prefix = "$constructorName(", separator = ", ", postfix = ")") { "\$$it" }
                 }
             }
         }
@@ -270,9 +271,9 @@ abstract class OutputValueBoxer(val outputValues: List<OutputValue>) {
 
         override val boxingRequired: Boolean = outputValues.size() > 0
 
-        override fun getBoxingExpressionText(arguments: List<String>): String? {
+        override fun getBoxingExpressionPattern(arguments: List<KtExpression>): String? {
             if (arguments.isEmpty()) return null
-            return arguments.joinToString(prefix = "kotlin.listOf(", separator = ", ", postfix = ")")
+            return arguments.indices.joinToString(prefix = "kotlin.listOf(", separator = ", ", postfix = ")") { "\$$it" }
         }
 
         override fun extractExpressionByIndex(boxedExpression: KtExpression, index: Int): KtExpression? {
@@ -324,7 +325,14 @@ val ControlFlow.possibleReturnTypes: List<KotlinType>
     }
 
 fun ControlFlow.toDefault(): ControlFlow =
-        copy(outputValues = outputValues.filterNot { it is OutputValue.Jump || it is OutputValue.ExpressionValue })
+        copy(outputValues = outputValues.filterNot { it is Jump || it is ExpressionValue })
+
+fun ControlFlow.copy(oldToNewParameters: Map<Parameter, Parameter>): ControlFlow {
+    val newOutputValues = outputValues.map {
+        if (it is ParameterUpdate) ParameterUpdate(oldToNewParameters[it.parameter]!!, it.originalExpressions) else it
+    }
+    return copy(outputValues = newOutputValues)
+}
 
 data class ExtractableCodeDescriptor(
         val extractionData: ExtractionData,
@@ -334,7 +342,7 @@ data class ExtractableCodeDescriptor(
         val parameters: List<Parameter>,
         val receiverParameter: Parameter?,
         val typeParameters: List<TypeParameter>,
-        val replacementMap: MultiMap<Int, Replacement>,
+        val replacementMap: MultiMap<KtSimpleNameExpression, Replacement>,
         val controlFlow: ControlFlow,
         val returnType: KotlinType
 ) {
@@ -342,6 +350,39 @@ data class ExtractableCodeDescriptor(
     val duplicates: List<DuplicateInfo> by lazy { findDuplicates() }
 }
 
+fun ExtractableCodeDescriptor.copy(
+        newName: String,
+        newVisibility: String,
+        oldToNewParameters: Map<Parameter, Parameter>,
+        newReceiver: Parameter?,
+        returnType: KotlinType?
+): ExtractableCodeDescriptor {
+    val newReplacementMap = MultiMap.create<KtSimpleNameExpression, Replacement>()
+    for ((ref, replacements) in replacementMap.entrySet()) {
+        val newReplacements = replacements.map {
+            if (it is ParameterReplacement) {
+                val parameter = it.parameter
+                val newParameter = oldToNewParameters[parameter] ?: return@map it
+                it.copy(newParameter)
+            }
+            else it
+        }
+        newReplacementMap.putValues(ref, newReplacements)
+    }
+
+    return ExtractableCodeDescriptor(
+            extractionData,
+            originalContext,
+            listOf(newName),
+            newVisibility,
+            parameters.map { oldToNewParameters[it]!! },
+            newReceiver,
+            typeParameters,
+            newReplacementMap,
+            controlFlow.copy(oldToNewParameters),
+            returnType ?: this.returnType)
+}
+
 enum class ExtractionTarget(val targetName: String) {
     FUNCTION("function") {
         override fun isAvailable(descriptor: ExtractableCodeDescriptor) = true
@@ -431,9 +472,10 @@ data class ExtractionGeneratorConfiguration(
 data class ExtractionResult(
         val config: ExtractionGeneratorConfiguration,
         val declaration: KtNamedDeclaration,
-        val duplicateReplacers: Map<KotlinPsiRange, () -> Unit>,
+        val duplicateReplacers: Map<KotlinPsiRange, () -> Unit>
-        val nameByOffset: MultiMap<Int, KtElement>
+) : Disposable {
-)
+    override fun dispose() = unmarkReferencesInside(declaration)
+}
 
 class AnalysisResult (
         val descriptor: ExtractableCodeDescriptor?,

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

@@ -16,7 +16,9 @@
 
 package org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine
 
+import com.intellij.openapi.Disposable
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiNameIdentifierOwner
 import com.intellij.psi.util.PsiTreeUtil
@@ -72,20 +74,21 @@ data class ResolveResult(
 
 data class ResolvedReferenceInfo(
         val refExpr: KtSimpleNameExpression,
-        val offsetInBody: Int,
         val resolveResult: ResolveResult,
         val smartCast: KotlinType?,
         val possibleTypes: Set<KotlinType>
 )
 
+internal var KtSimpleNameExpression.resolveResult: ResolveResult? by CopyableUserDataProperty(Key.create("RESOLVE_RESULT"))
+
 data class ExtractionData(
         val originalFile: KtFile,
         val originalRange: KotlinPsiRange,
         val targetSibling: PsiElement,
         val duplicateContainer: PsiElement? = null,
         val options: ExtractionOptions = ExtractionOptions.DEFAULT
-) {
+) : Disposable {
-    val project: Project = originalFile.getProject()
+    val project: Project = originalFile.project
     val originalElements: List<PsiElement> = originalRange.elements
     val physicalElements = originalElements.map { it.substringContextOrThis }
 
@@ -108,9 +111,6 @@ data class ExtractionData(
         }
     }
 
-    val originalStartOffset: Int?
-        get() = originalElements.firstOrNull()?.let { e -> e.getTextRange()!!.getStartOffset() }
-
     val commonParent = PsiTreeUtil.findCommonParent(physicalElements) as KtElement
 
     val bindingContext: BindingContext? by lazy { commonParent.getContextForContainingDeclarationBody() }
@@ -118,65 +118,63 @@ data class ExtractionData(
     private val itFakeDeclaration by lazy { KtPsiFactory(originalFile).createParameter("it: Any?") }
     private val synthesizedInvokeDeclaration by lazy { KtPsiFactory(originalFile).createFunction("fun invoke() {}") }
 
-    val refOffsetToDeclaration by lazy {
+    init {
-        fun isExtractableIt(descriptor: DeclarationDescriptor, context: BindingContext): Boolean {
+        markReferences()
-            if (!(descriptor is ValueParameterDescriptor && (context[BindingContext.AUTO_CREATED_IT, descriptor] ?: false))) return false
+    }
-            val function = DescriptorToSourceUtils.descriptorToDeclaration(descriptor.getContainingDeclaration()) as? KtFunctionLiteral
+
-            return function == null || !function.isInsideOf(physicalElements)
+    private fun isExtractableIt(descriptor: DeclarationDescriptor, context: BindingContext): Boolean {
+        if (!(descriptor is ValueParameterDescriptor && (context[BindingContext.AUTO_CREATED_IT, descriptor] ?: false))) return false
+        val function = DescriptorToSourceUtils.descriptorToDeclaration(descriptor.containingDeclaration) as? KtFunctionLiteral
+        return function == null || !function.isInsideOf(physicalElements)
+    }
+
+    private tailrec fun getDeclaration(descriptor: DeclarationDescriptor, context: BindingContext): PsiNameIdentifierOwner? {
+        (DescriptorToSourceUtilsIde.getAnyDeclaration(project, descriptor) as? PsiNameIdentifierOwner)?.let { return it }
+
+        return when {
+            isExtractableIt(descriptor, context) -> itFakeDeclaration
+            isSynthesizedInvoke(descriptor) -> synthesizedInvokeDeclaration
+            descriptor is SyntheticJavaPropertyDescriptor -> getDeclaration(descriptor.getMethod, context)
+            else -> null
         }
+    }
 
-        tailrec fun getDeclaration(descriptor: DeclarationDescriptor, context: BindingContext): PsiNameIdentifierOwner? {
+    private fun markReferences() {
-            (DescriptorToSourceUtilsIde.getAnyDeclaration(project, descriptor) as? PsiNameIdentifierOwner)?.let { return it }
+        val context = bindingContext ?: return
-
+        val visitor = object : KtTreeVisitorVoid() {
-            return when {
+            override fun visitQualifiedExpression(expression: KtQualifiedExpression) {
-                isExtractableIt(descriptor, context) -> itFakeDeclaration
+                if (context[BindingContext.SMARTCAST, expression] != null) {
-                isSynthesizedInvoke(descriptor) -> synthesizedInvokeDeclaration
+                    expression.selectorExpression?.accept(this)
-                descriptor is SyntheticJavaPropertyDescriptor -> getDeclaration(descriptor.getMethod, context)
+                    return
-                else -> null
-            }
-        }
-
-        val originalStartOffset = originalStartOffset
-        val context = bindingContext
-
-        if (originalStartOffset != null && context != null) {
-            val resultMap = HashMap<Int, ResolveResult>()
-
-            val visitor = object: KtTreeVisitorVoid() {
-                override fun visitQualifiedExpression(expression: KtQualifiedExpression) {
-                    if (context[BindingContext.SMARTCAST, expression] != null) {
-                        expression.getSelectorExpression()?.accept(this)
-                        return
-                    }
-
-                    super.visitQualifiedExpression(expression)
                 }
 
-                override fun visitSimpleNameExpression(ref: KtSimpleNameExpression) {
+                super.visitQualifiedExpression(expression)
-                    if (ref.getParent() is KtValueArgumentName) return
+            }
 
-                    val physicalRef = substringInfo?.let {
+            override fun visitSimpleNameExpression(ref: KtSimpleNameExpression) {
-                        // If substring contains some references it must be extracted as a string template
+                if (ref.parent is KtValueArgumentName) return
-                        val physicalExpression = expressions.single() as KtStringTemplateExpression
-                        val extractedContentOffset = physicalExpression.getContentRange().startOffset + physicalExpression.startOffset
-                        val offsetInExtracted = ref.startOffset - extractedContentOffset
-                        val offsetInTemplate = it.relativeContentRange.startOffset + offsetInExtracted
-                        it.template.findElementAt(offsetInTemplate)!!.getStrictParentOfType<KtSimpleNameExpression>()
-                    } ?: ref
 
-                    val resolvedCall = physicalRef.getResolvedCall(context)
+                val physicalRef = substringInfo?.let {
-                    val descriptor = context[BindingContext.REFERENCE_TARGET, physicalRef] ?: return
+                    // If substring contains some references it must be extracted as a string template
-                    val declaration = getDeclaration(descriptor, context) ?: return
+                    val physicalExpression = expressions.single() as KtStringTemplateExpression
+                    val extractedContentOffset = physicalExpression.getContentRange().startOffset + physicalExpression.startOffset
+                    val offsetInExtracted = ref.startOffset - extractedContentOffset
+                    val offsetInTemplate = it.relativeContentRange.startOffset + offsetInExtracted
+                    it.template.findElementAt(offsetInTemplate)!!.getStrictParentOfType<KtSimpleNameExpression>()
+                } ?: ref
 
-                    val offset = ref.getTextRange()!!.getStartOffset() - originalStartOffset
+                val resolvedCall = physicalRef.getResolvedCall(context)
-                    resultMap[offset] = ResolveResult(physicalRef, declaration, descriptor, resolvedCall)
+                val descriptor = context[BindingContext.REFERENCE_TARGET, physicalRef] ?: return
+                val declaration = getDeclaration(descriptor, context) ?: return
+
+                val resolveResult = ResolveResult(physicalRef, declaration, descriptor, resolvedCall)
+                physicalRef.resolveResult = resolveResult
+                if (ref != physicalRef) {
+                    ref.resolveResult = resolveResult
                 }
             }
-            expressions.forEach { it.accept(visitor) }
-
-            resultMap
         }
-        else Collections.emptyMap<Int, ResolveResult>()
+        expressions.forEach { it.accept(visitor) }
     }
 
     fun getPossibleTypes(expression: KtExpression, resolvedCall: ResolvedCall<*>?, context: BindingContext): Set<KotlinType> {
@@ -195,75 +193,72 @@ data class ExtractionData(
     fun getBrokenReferencesInfo(body: KtBlockExpression): List<ResolvedReferenceInfo> {
         val originalContext = bindingContext ?: return listOf()
 
-        val startOffset = body.getBlockContentOffset()
+        val newReferences = body.collectDescendantsOfType<KtSimpleNameExpression> { it.resolveResult != null }
 
         val referencesInfo = ArrayList<ResolvedReferenceInfo>()
         val refToContextMap = KotlinFileReferencesResolver.resolve(body)
-        for ((ref, context) in refToContextMap) {
+        for (newRef in newReferences) {
-            if (ref !is KtSimpleNameExpression) continue
+            val context = refToContextMap[newRef] ?: continue
-
+            val originalResolveResult = newRef.resolveResult ?: continue
-            val offset = ref.getTextRange()!!.getStartOffset() - startOffset
-            val originalResolveResult = refOffsetToDeclaration[offset] ?: continue
 
             val smartCast: KotlinType?
             val possibleTypes: Set<KotlinType>
 
             // Qualified property reference: a.b
-            val qualifiedExpression = ref.getQualifiedExpressionForSelector()
+            val qualifiedExpression = newRef.getQualifiedExpressionForSelector()
             if (qualifiedExpression != null) {
-                val smartCastTarget = originalResolveResult.originalRefExpr.getParent() as KtExpression
+                val smartCastTarget = originalResolveResult.originalRefExpr.parent as KtExpression
                 smartCast = originalContext[BindingContext.SMARTCAST, smartCastTarget]
                 possibleTypes = getPossibleTypes(smartCastTarget, originalResolveResult.resolvedCall, originalContext)
                 val receiverDescriptor =
-                        (originalResolveResult.resolvedCall?.getDispatchReceiver() as? ImplicitReceiver)?.declarationDescriptor
+                        (originalResolveResult.resolvedCall?.dispatchReceiver as? ImplicitReceiver)?.declarationDescriptor
                 if (smartCast == null
                     && !DescriptorUtils.isCompanionObject(receiverDescriptor)
-                    && qualifiedExpression.getReceiverExpression() !is KtSuperExpression) continue
+                    && qualifiedExpression.receiverExpression !is KtSuperExpression) continue
             }
             else {
                 smartCast = originalContext[BindingContext.SMARTCAST, originalResolveResult.originalRefExpr]
                 possibleTypes = getPossibleTypes(originalResolveResult.originalRefExpr, originalResolveResult.resolvedCall, originalContext)
             }
 
-            val parent = ref.getParent()
+            val parent = newRef.parent
 
             // Skip P in type references like 'P.Q'
-            if (parent is KtUserType && (parent.getParent() as? KtUserType)?.getQualifier() == parent) continue
+            if (parent is KtUserType && (parent.parent as? KtUserType)?.qualifier == parent) continue
 
-            val descriptor = context[BindingContext.REFERENCE_TARGET, ref]
+            val descriptor = context[BindingContext.REFERENCE_TARGET, newRef]
             val isBadRef = !(compareDescriptors(project, originalResolveResult.descriptor, descriptor)
-                             && originalContext.diagnostics.forElement(originalResolveResult.originalRefExpr) == context.diagnostics.forElement(ref))
+                             && originalContext.diagnostics.forElement(originalResolveResult.originalRefExpr) == context.diagnostics.forElement(newRef))
                            || smartCast != null
             if (isBadRef && !originalResolveResult.declaration.isInsideOf(physicalElements)) {
                 val originalResolvedCall = originalResolveResult.resolvedCall as? VariableAsFunctionResolvedCall
                 val originalFunctionCall = originalResolvedCall?.functionCall
                 val originalVariableCall = originalResolvedCall?.variableCall
-                val invokeDescriptor = originalFunctionCall?.getResultingDescriptor()
+                val invokeDescriptor = originalFunctionCall?.resultingDescriptor
                 if (invokeDescriptor != null && isSynthesizedInvoke(invokeDescriptor) && invokeDescriptor.isExtension) {
                     val variableResolveResult = originalResolveResult.copy(resolvedCall = originalVariableCall!!,
-                                                                           descriptor = originalVariableCall.getResultingDescriptor())
+                                                                           descriptor = originalVariableCall.resultingDescriptor)
                     val functionResolveResult = originalResolveResult.copy(resolvedCall = originalFunctionCall!!,
-                                                                           descriptor = originalFunctionCall.getResultingDescriptor(),
+                                                                           descriptor = originalFunctionCall.resultingDescriptor,
                                                                            declaration = synthesizedInvokeDeclaration)
-                    referencesInfo.add(ResolvedReferenceInfo(ref, offset, variableResolveResult, smartCast, possibleTypes))
+                    referencesInfo.add(ResolvedReferenceInfo(newRef, variableResolveResult, smartCast, possibleTypes))
-                    referencesInfo.add(ResolvedReferenceInfo(ref, offset, functionResolveResult, smartCast, possibleTypes))
+                    referencesInfo.add(ResolvedReferenceInfo(newRef, functionResolveResult, smartCast, possibleTypes))
                 }
                 else {
-                    referencesInfo.add(ResolvedReferenceInfo(ref, offset, originalResolveResult, smartCast, possibleTypes))
+                    referencesInfo.add(ResolvedReferenceInfo(newRef, originalResolveResult, smartCast, possibleTypes))
                 }
             }
         }
 
         return referencesInfo
     }
+
+    override fun dispose() {
+        expressions.forEach { unmarkReferencesInside(it) }
+    }
 }
 
-// Hack:
+fun unmarkReferencesInside(root: PsiElement) {
-// we can't get first element offset through getStatement()/getChildren() since they skip comments and whitespaces
+    if (!root.isValid) return
-// So we take offset of the left brace instead and increase it by 2 (which is length of "{\n" separating block start and its first element)
+    root.forEachDescendantOfType<KtSimpleNameExpression> { it.resolveResult = null }
-internal fun KtExpression.getBlockContentOffset(): Int {
-    (this as? KtBlockExpression)?.getLBrace()?.let {
-        return it.getTextRange()!!.getStartOffset() + 2
-    }
-    return getTextRange()!!.getStartOffset()
 }

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

@@ -16,21 +16,18 @@
 
 package org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine
 
-import com.intellij.psi.*
+import com.intellij.openapi.application.ApplicationManager
-import org.jetbrains.kotlin.psi.*
+import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.editor.*
+import com.intellij.openapi.project.Project
-import org.jetbrains.kotlin.idea.util.psi.patternMatching.*
+import com.intellij.openapi.ui.MessageType
-import com.intellij.openapi.application.*
+import com.intellij.openapi.ui.popup.Balloon
-import com.intellij.refactoring.*
+import com.intellij.openapi.ui.popup.JBPopupFactory
-import org.jetbrains.kotlin.idea.util.application.*
+import com.intellij.refactoring.BaseRefactoringProcessor
-import org.jetbrains.kotlin.idea.refactoring.*
+import com.intellij.ui.awt.RelativePoint
-import org.jetbrains.kotlin.idea.refactoring.introduce.*
-import com.intellij.ui.awt.*
-import com.intellij.openapi.ui.popup.*
-import com.intellij.openapi.ui.*
-import javax.swing.event.*
-import com.intellij.openapi.project.*
 import org.jetbrains.kotlin.idea.core.refactoring.checkConflictsInteractively
+import org.jetbrains.kotlin.idea.refactoring.introduce.showErrorHint
+import org.jetbrains.kotlin.idea.util.application.executeWriteCommand
+import javax.swing.event.HyperlinkEvent
 
 public abstract class ExtractionEngineHelper(val operationName: String) {
     open fun adjustExtractionData(data: ExtractionData): ExtractionData = data
@@ -66,7 +63,15 @@ public class ExtractionEngine(
         fun validateAndRefactor() {
             val validationResult = analysisResult.descriptor!!.validate()
             project.checkConflictsInteractively(validationResult.conflicts) {
-                helper.configureAndRun(project, editor, validationResult, onFinish)
+                helper.configureAndRun(project, editor, validationResult) {
+                    try {
+                        onFinish(it)
+                    }
+                    finally {
+                        it.dispose()
+                        extractionData.dispose()
+                    }
+                }
             }
         }
 

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

@@ -19,6 +19,7 @@ package org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiNamedElement
+import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.refactoring.util.RefactoringUIUtil
 import com.intellij.util.containers.MultiMap
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
@@ -55,10 +56,7 @@ import org.jetbrains.kotlin.idea.util.approximateWithResolvableType
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.idea.util.isResolvableInScope
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.getElementTextWithContext
+import org.jetbrains.kotlin.psi.psiUtil.*
-import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
-import org.jetbrains.kotlin.psi.psiUtil.isInsideOf
 import org.jetbrains.kotlin.renderer.DescriptorRenderer
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.DescriptorToSourceUtils
@@ -374,39 +372,35 @@ private fun ExtractionData.analyzeControlFlow(
     return controlFlow to null
 }
 
-fun ExtractionData.createTemporaryDeclaration(functionText: String): KtNamedDeclaration {
+fun ExtractionData.createTemporaryDeclaration(pattern: String): KtNamedDeclaration {
-    val textRange = targetSibling.textRange!!
+    val targetSiblingMarker = Any()
+    PsiTreeUtil.mark(targetSibling, targetSiblingMarker)
+    val tmpFile = originalFile.createTempCopy("")
+    tmpFile.deleteChildRange(tmpFile.firstChild, tmpFile.lastChild)
+    tmpFile.addRange(originalFile.firstChild, originalFile.lastChild)
+    val newTargetSibling = PsiTreeUtil.releaseMark(tmpFile, targetSiblingMarker)!!
+    val newTargetParent = newTargetSibling.parent
 
-    val insertText: String
+    val declaration = KtPsiFactory(originalFile).createDeclarationByPattern<KtNamedDeclaration>(
-    val insertPosition: Int
+            pattern,
-    val lookupPosition: Int
+            PsiChildRange(originalElements.firstOrNull(), originalElements.lastOrNull())
-    if (insertBefore) {
+    )
-        insertPosition = textRange.startOffset
+    return if (insertBefore) {
-        lookupPosition = insertPosition
+        newTargetParent.addBefore(declaration, newTargetSibling) as KtNamedDeclaration
-        insertText = functionText
     }
     else {
-        insertPosition = textRange.endOffset
+        newTargetParent.addAfter(declaration, newTargetSibling) as KtNamedDeclaration
-        lookupPosition = insertPosition + 1
-        insertText = "\n$functionText"
     }
-
-    val tmpFile = originalFile.createTempCopy { text ->
-        StringBuilder(text).insert(insertPosition, insertText).toString()
-    }
-    return tmpFile.findElementAt(lookupPosition)?.getNonStrictParentOfType<KtNamedDeclaration>()!!
 }
 
 internal fun ExtractionData.createTemporaryCodeBlock(): KtBlockExpression =
-        (createTemporaryDeclaration("fun() {\n$codeFragmentText\n}\n") as KtNamedFunction).bodyExpression as KtBlockExpression
+        (createTemporaryDeclaration("fun() {\n$0\n}\n") as KtNamedFunction).bodyExpression as KtBlockExpression
 
 private fun KotlinType.collectReferencedTypes(processTypeArguments: Boolean): List<KotlinType> {
     if (!processTypeArguments) return Collections.singletonList(this)
     return DFS.dfsFromNode(
             this,
-            object: Neighbors<KotlinType> {
+            Neighbors<KotlinType> { current -> current.arguments.map { it.type } },
-                override fun getNeighbors(current: KotlinType): Iterable<KotlinType> = current.arguments.map { it.type }
-            },
             VisitedWithSet(),
             object: CollectingNodeHandler<KotlinType, KotlinType, ArrayList<KotlinType>>(ArrayList()) {
                 override fun afterChildren(current: KotlinType) {
@@ -568,17 +562,6 @@ private class DelegatingParameter(
     override fun getParameterType(allowSpecialClassNames: Boolean) = parameterType
 }
 
-internal class ParametersInfo {
-    var errorMessage: ErrorMessage? = null
-    val replacementMap: MultiMap<Int, Replacement> = MultiMap.create()
-    val originalRefToParameter: MultiMap<KtSimpleNameExpression, MutableParameter> = MultiMap.create()
-    val parameters: MutableSet<MutableParameter> = HashSet()
-    val typeParameters: MutableSet<TypeParameter> = HashSet()
-    val nonDenotableTypes: MutableSet<KotlinType> = HashSet()
-}
-
-
-
 private fun ExtractionData.checkDeclarationsMovingOutOfScope(
         enclosingDeclaration: KtDeclaration,
         controlFlow: ControlFlow,
@@ -773,7 +756,8 @@ fun ExtractableCodeDescriptor.validate(): ExtractableCodeDescriptorWithConflicts
     val body = result.declaration.getGeneratedBody()
     val bindingContext = body.analyzeFully()
 
-    fun processReference(resolveResult: ResolveResult, currentRefExpr: KtReferenceExpression) {
+    fun processReference(currentRefExpr: KtSimpleNameExpression) {
+        val resolveResult = currentRefExpr.resolveResult ?: return
         if (currentRefExpr.parent is KtThisExpression) return
 
         val diagnostics = bindingContext.diagnostics.forElement(currentRefExpr)
@@ -808,18 +792,6 @@ fun ExtractableCodeDescriptor.validate(): ExtractableCodeDescriptorWithConflicts
         }
     }
 
-    fun validateBody() {
-        for ((originalOffset, resolveResult) in extractionData.refOffsetToDeclaration) {
-            if (resolveResult.declaration.isInsideOf(extractionData.physicalElements)) continue
-
-            val currentRefExprs = result.nameByOffset[originalOffset].mapNotNull {
-                (it as? KtThisExpression)?.instanceReference ?: it as? KtSimpleNameExpression
-            }
-
-            currentRefExprs.forEach { processReference(resolveResult, it) }
-        }
-    }
-
     result.declaration.accept(
             object : KtTreeVisitorVoid() {
                 override fun visitUserType(userType: KtUserType) {
@@ -831,12 +803,8 @@ fun ExtractableCodeDescriptor.validate(): ExtractableCodeDescriptorWithConflicts
                     }
                 }
 
-                override fun visitKtElement(element: KtElement) {
+                override fun visitSimpleNameExpression(expression: KtSimpleNameExpression) {
-                    if (element == body) {
+                    processReference(expression)
-                        validateBody()
-                        return
-                    }
-                    super.visitKtElement(element)
                 }
             }
     )

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

@@ -16,12 +16,12 @@
 
 package org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine
 
+import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiElement
-import com.intellij.psi.PsiFile
+import com.intellij.psi.codeStyle.CodeStyleManager
 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.*
@@ -56,8 +56,65 @@ import org.jetbrains.kotlin.types.isFlexible
 import org.jetbrains.kotlin.types.typeUtil.builtIns
 import java.util.*
 
-fun ExtractionGeneratorConfiguration.getDeclarationText(
+private fun buildSignature(config: ExtractionGeneratorConfiguration, renderer: DescriptorRenderer): CallableBuilder {
-        withBody: Boolean = true,
+    val extractionTarget = config.generatorOptions.target
+    if (!extractionTarget.isAvailable(config.descriptor)) {
+        val message = "Can't generate ${extractionTarget.targetName}: ${config.descriptor.extractionData.codeFragmentText}"
+        throw BaseRefactoringProcessor.ConflictsInTestsException(listOf(message))
+    }
+
+    val builderTarget = when (extractionTarget) {
+        ExtractionTarget.FUNCTION, ExtractionTarget.FAKE_LAMBDALIKE_FUNCTION -> CallableBuilder.Target.FUNCTION
+        else -> CallableBuilder.Target.READ_ONLY_PROPERTY
+    }
+    return CallableBuilder(builderTarget).apply {
+        modifier(config.descriptor.visibility)
+
+        typeParams(
+                config.descriptor.typeParameters.map {
+                    val typeParameter = it.originalDeclaration
+                    val bound = typeParameter.extendsBound
+                    typeParameter.name + (bound?.let { " : " + it.text } ?: "")
+                }
+        )
+
+        fun KotlinType.typeAsString(): String {
+            return if (config.descriptor.extractionData.options.allowSpecialClassNames && isSpecial()) {
+                DEBUG_TYPE_REFERENCE_STRING
+            } else {
+                renderer.renderType(this)
+            }
+        }
+
+        config.descriptor.receiverParameter?.let {
+            val receiverType = it.getParameterType(config.descriptor.extractionData.options.allowSpecialClassNames)
+            val receiverTypeAsString = receiverType.typeAsString()
+            val isFunctionType = KotlinBuiltIns.isExactFunctionOrExtensionFunctionType(receiverType)
+            receiver(if (isFunctionType) "($receiverTypeAsString)" else receiverTypeAsString)
+        }
+
+        name(config.generatorOptions.dummyName ?: config.descriptor.name)
+
+        config.descriptor.parameters.forEach { parameter ->
+            param(parameter.name,
+                       parameter.getParameterType(config.descriptor.extractionData.options.allowSpecialClassNames).typeAsString())
+        }
+
+        with(config.descriptor.returnType) {
+            if (isDefault() || isError || extractionTarget == ExtractionTarget.PROPERTY_WITH_INITIALIZER) {
+                noReturnType()
+            } else {
+                returnType(typeAsString())
+            }
+        }
+
+        typeConstraints(config.descriptor.typeParameters.flatMap { it.originalConstraints }.map { it.text!! })
+    }
+}
+
+fun ExtractionGeneratorConfiguration.getSignaturePreview(renderer: DescriptorRenderer) = buildSignature(this, renderer).asString()
+
+fun ExtractionGeneratorConfiguration.getDeclarationPattern(
         descriptorRenderer: DescriptorRenderer = if (generatorOptions.flexibleTypesAllowed)
                                                     DescriptorRenderer.FLEXIBLE_TYPES_FOR_CODE
                                                  else IdeDescriptorRenderers.SOURCE_CODE
@@ -67,58 +124,20 @@ fun ExtractionGeneratorConfiguration.getDeclarationText(
         throw BaseRefactoringProcessor.ConflictsInTestsException(listOf("Can't generate ${extractionTarget.targetName}: ${descriptor.extractionData.codeFragmentText}"))
     }
 
-    val builderTarget = when (extractionTarget) {
+    return buildSignature(this, descriptorRenderer).let { builder ->
-        ExtractionTarget.FUNCTION, ExtractionTarget.FAKE_LAMBDALIKE_FUNCTION -> CallableBuilder.Target.FUNCTION
+        builder.transform {
-        else -> CallableBuilder.Target.READ_ONLY_PROPERTY
+            for (i in sequence(indexOf('$')) { indexOf('$', it + 2) }) {
-    }
+                if (i < 0) break
-    return CallableBuilder(builderTarget).let { builder ->
+                insert(i + 1, '$')
-        builder.modifier(descriptor.visibility)
-
-        builder.typeParams(
-                descriptor.typeParameters.map {
-                    val typeParameter = it.originalDeclaration
-                    val bound = typeParameter.extendsBound
-                    typeParameter.name + (bound?.let { " : " + it.text } ?: "")
-                }
-        )
-
-        fun KotlinType.typeAsString(): String {
-            return if (descriptor.extractionData.options.allowSpecialClassNames && isSpecial()) DEBUG_TYPE_REFERENCE_STRING else descriptorRenderer.renderType(this)
-        }
-
-        descriptor.receiverParameter?.let {
-            val receiverType = it.getParameterType(descriptor.extractionData.options.allowSpecialClassNames)
-            val receiverTypeAsString = receiverType.typeAsString()
-            val isFunctionType = KotlinBuiltIns.isExactFunctionOrExtensionFunctionType(receiverType)
-            builder.receiver(if (isFunctionType) "($receiverTypeAsString)" else receiverTypeAsString)
-        }
-
-        builder.name(generatorOptions.dummyName ?: descriptor.name)
-
-        descriptor.parameters.forEach { parameter ->
-            builder.param(parameter.name,
-                          parameter.getParameterType(descriptor.extractionData.options.allowSpecialClassNames).typeAsString())
-        }
-
-        with(descriptor.returnType) {
-            if (isDefault() || isError || extractionTarget == ExtractionTarget.PROPERTY_WITH_INITIALIZER) {
-                builder.noReturnType()
-            } else {
-                builder.returnType(typeAsString())
             }
         }
 
-        builder.typeConstraints(descriptor.typeParameters.flatMap { it.originalConstraints }.map { it.text!! })
+        when (extractionTarget) {
-
+            ExtractionTarget.FUNCTION,
-        if (withBody) {
+            ExtractionTarget.FAKE_LAMBDALIKE_FUNCTION,
-            val bodyText = descriptor.extractionData.codeFragmentText
+            ExtractionTarget.PROPERTY_WITH_GETTER -> builder.blockBody("$0")
-            when (extractionTarget) {
+            ExtractionTarget.PROPERTY_WITH_INITIALIZER -> builder.initializer("$0")
-                ExtractionTarget.FUNCTION,
+            ExtractionTarget.LAZY_PROPERTY -> builder.lazyBody("$0")
-                ExtractionTarget.FAKE_LAMBDALIKE_FUNCTION,
-                ExtractionTarget.PROPERTY_WITH_GETTER -> builder.blockBody(bodyText)
-                ExtractionTarget.PROPERTY_WITH_INITIALIZER -> builder.initializer(bodyText)
-                ExtractionTarget.LAZY_PROPERTY -> builder.lazyBody(bodyText)
-            }
         }
 
         builder.asString()
@@ -131,22 +150,7 @@ fun KotlinType.isSpecial(): Boolean {
 }
 
 fun createNameCounterpartMap(from: KtElement, to: KtElement): Map<KtSimpleNameExpression, KtSimpleNameExpression> {
-    val map = HashMap<KtSimpleNameExpression, KtSimpleNameExpression>()
+    return from.collectDescendantsOfType<KtSimpleNameExpression>().zip(to.collectDescendantsOfType<KtSimpleNameExpression>()).toMap()
-
-    val fromOffset = from.textRange!!.startOffset
-    from.accept(
-            object : KtTreeVisitorVoid() {
-                override fun visitSimpleNameExpression(expression: KtSimpleNameExpression) {
-                    val offset = expression.textRange!!.startOffset - fromOffset
-                    val newExpression = to.findElementAt(offset)?.getNonStrictParentOfType<KtSimpleNameExpression>()
-                    assert(newExpression != null) { "Couldn't find expression at $offset in '${to.text}'" }
-
-                    map[expression] = newExpression!!
-                }
-            }
-    )
-
-    return map
 }
 
 class DuplicateInfo(
@@ -408,131 +412,94 @@ private fun makeCall(
     }
 }
 
+private var KtExpression.isJumpElementToReplace: Boolean
+        by NotNullableCopyableUserDataProperty(Key.create("IS_JUMP_ELEMENT_TO_REPLACE"), false)
+
+private var KtReturnExpression.isReturnForLabelRemoval: Boolean
+        by NotNullableCopyableUserDataProperty(Key.create("IS_RETURN_FOR_LABEL_REMOVAL"), false)
+
 fun ExtractionGeneratorConfiguration.generateDeclaration(
         declarationToReplace: KtNamedDeclaration? = null
 ): ExtractionResult{
     val psiFactory = KtPsiFactory(descriptor.extractionData.originalFile)
-    val nameByOffset = MultiMap<Int, KtElement>()
+
+    fun getReturnsForLabelRemoval() = descriptor.controlFlow.outputValues
+            .flatMapTo(ArrayList<KtReturnExpression>()) { it.originalExpressions.filterIsInstance<KtReturnExpression>() }
 
     fun createDeclaration(): KtNamedDeclaration {
+        descriptor.controlFlow.jumpOutputValue?.elementsToReplace?.forEach { it.isJumpElementToReplace = true }
+        getReturnsForLabelRemoval().forEach { it.isReturnForLabelRemoval = true }
+
         return with(descriptor.extractionData) {
             if (generatorOptions.inTempFile) {
-                createTemporaryDeclaration("${getDeclarationText()}\n")
+                createTemporaryDeclaration("${getDeclarationPattern()}\n")
             }
             else {
-                psiFactory.createDeclaration(getDeclarationText())
+                psiFactory.createDeclarationByPattern(
+                        getDeclarationPattern(),
+                        PsiChildRange(originalElements.firstOrNull(), originalElements.lastOrNull())
+                )
             }
         }
     }
 
-    fun getReturnArguments(resultExpression: KtExpression?): List<String> {
+    fun getReturnArguments(resultExpression: KtExpression?): List<KtExpression> {
         return descriptor.controlFlow.outputValues
                 .mapNotNull {
                     when (it) {
-                        is ExpressionValue -> resultExpression?.text
+                        is ExpressionValue -> resultExpression
-                        is Jump -> if (it.conditional) "false" else null
+                        is Jump -> if (it.conditional) psiFactory.createExpression("false") else null
-                        is ParameterUpdate -> it.parameter.nameForRef
+                        is ParameterUpdate -> psiFactory.createExpression(it.parameter.nameForRef)
-                        is Initializer -> it.initializedDeclaration.name
+                        is Initializer -> psiFactory.createExpression(it.initializedDeclaration.name!!)
                         else -> throw IllegalArgumentException("Unknown output value: $it")
                     }
                 }
     }
 
-    fun replaceWithReturn(
+    fun KtExpression.replaceWithReturn(replacingExpression: KtReturnExpression) {
-            originalExpression: KtExpression,
+        descriptor.controlFlow.defaultOutputValue?.let {
-            replacingExpression: KtReturnExpression,
+            val boxedExpression = replaced(replacingExpression).returnedExpression!!
-            expressionToUnifyWith: KtExpression?
-    ) {
-        val currentResultExpression =
-                (if (originalExpression is KtReturnExpression) originalExpression.returnedExpression else originalExpression) ?: return
-
-        val newResultExpression = descriptor.controlFlow.defaultOutputValue?.let {
-            val boxedExpression = originalExpression.replaced(replacingExpression).returnedExpression!!
             descriptor.controlFlow.outputValueBoxer.extractExpressionByValue(boxedExpression, it)
         }
-
-        @Suppress
-        if (newResultExpression == null) {
-            throw AssertionError("Can' replace '${originalExpression.text}' with '${replacingExpression.text}'")
-        }
-
-        val counterpartMap = createNameCounterpartMap(currentResultExpression, expressionToUnifyWith ?: newResultExpression)
-        counterpartMap.entries.forEach {
-            val (cmOriginalExpr, cmNewExpr) = it
-            nameByOffset.entrySet().find { cmOriginalExpr in it.value }?.let {
-                nameByOffset.remove(it.key, cmOriginalExpr)
-                nameByOffset.putValue(it.key, cmNewExpr)
-            }
-        }
-    }
-
-    fun getCounterparts<T : KtExpression>(originalExpressions: Collection<T>,
-                                                                   body: KtExpression,
-                                                                   bodyOffset: Int,
-                                                                   file: PsiFile): List<T> {
-        return originalExpressions.map { originalExpression ->
-            val offsetInBody = originalExpression.textRange!!.startOffset - descriptor.extractionData.originalStartOffset!!
-            file.findElementAt(bodyOffset + offsetInBody)?.getNonStrictParentOfType(originalExpression.javaClass)
-            ?: throw AssertionError("Couldn't find expression at $offsetInBody in '${body.text}'")
-        }
     }
 
     fun adjustDeclarationBody(declaration: KtNamedDeclaration) {
         val body = declaration.getGeneratedBody()
 
-        val exprReplacementMap = MultiMap<KtElement, (ExtractableCodeDescriptor, KtElement) -> KtElement>()
+        val jumpValue = descriptor.controlFlow.jumpOutputValue
-        val originalOffsetByExpr = LinkedHashMap<KtElement, Int>()
+        if (jumpValue != null) {
+            val replacingReturn = psiFactory.createExpression(if (jumpValue.conditional) "return true" else "return")
+            body.collectDescendantsOfType<KtExpression> { it.isJumpElementToReplace }.forEach {
+                it.replace(replacingReturn)
+                it.isJumpElementToReplace = false
+            }
+        }
 
-        val bodyOffset = body.getBlockContentOffset()
+        body.collectDescendantsOfType<KtReturnExpression> { it.isReturnForLabelRemoval }.forEach {
-        val file = body.containingFile!!
+            it.getTargetLabel()?.delete()
+            it.isReturnForLabelRemoval = false
+        }
 
         /*
          * Sort by descending position so that internals of value/type arguments in calls and qualified types are replaced
          * before calls/types themselves
          */
-        for ((offsetInBody, resolveResult) in descriptor.extractionData.refOffsetToDeclaration.entries.sortedByDescending { it.key }) {
+        val currentRefs = body
-            val expr = file.findElementAt(bodyOffset + offsetInBody)?.getNonStrictParentOfType<KtSimpleNameExpression>()
+                .collectDescendantsOfType<KtSimpleNameExpression> { it.resolveResult != null }
-            assert(expr != null) { "Couldn't find expression at $offsetInBody in '${body.text}'" }
+                .sortedByDescending { it.startOffset }
 
-            originalOffsetByExpr[expr!!] = offsetInBody
+        currentRefs.forEach {
-
+            val resolveResult = it.resolveResult!!
-            descriptor.replacementMap[offsetInBody].let { replacements ->
+            val currentRef = if (it.isValid) {
-                replacements.forEach {
+                it
-                    exprReplacementMap.putValue(expr, it)
-                }
             }
-        }
+            else {
-
+                body.findDescendantOfType<KtSimpleNameExpression> { it.resolveResult == resolveResult } ?: return@forEach
-        val replacingReturn: KtExpression?
-        val expressionsToReplaceWithReturn: List<KtElement>
-
-        val returnsForLabelRemoval = descriptor.controlFlow.outputValues
-                .flatMapTo(ArrayList<KtReturnExpression>()) { it.originalExpressions.filterIsInstance<KtReturnExpression>() }
-
-        val jumpValue = descriptor.controlFlow.jumpOutputValue
-        if (jumpValue != null) {
-            replacingReturn = psiFactory.createExpression(if (jumpValue.conditional) "return true" else "return")
-            returnsForLabelRemoval.removeAll(jumpValue.elementsToReplace)
-            expressionsToReplaceWithReturn = getCounterparts(jumpValue.elementsToReplace, body, bodyOffset, file)
-        }
-        else {
-            replacingReturn = null
-            expressionsToReplaceWithReturn = Collections.emptyList()
-        }
-
-        if (replacingReturn != null) {
-            for (expr in expressionsToReplaceWithReturn) {
-                expr.replace(replacingReturn)
-            }
-        }
-
-        getCounterparts(returnsForLabelRemoval, body, bodyOffset, file).forEach { it.getTargetLabel()?.delete() }
-
-        for ((expr, originalOffset) in originalOffsetByExpr) {
-            if (expr.isValid) {
-                val replacements = exprReplacementMap[expr].mapNotNull { it?.invoke(descriptor, expr) }
-                nameByOffset.put(originalOffset, if (replacements.isEmpty()) arrayListOf(expr) else replacements)
             }
+            val originalRef = resolveResult.originalRefExpr
+            val newRef = descriptor.replacementMap[originalRef]
+                    .fold(currentRef as KtElement) { currentRef, replacement -> replacement(descriptor, currentRef) }
+            (newRef as? KtSimpleNameExpression)?.resolveResult = resolveResult
         }
 
         if (generatorOptions.target == ExtractionTarget.PROPERTY_WITH_INITIALIZER) return
@@ -554,20 +521,19 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
         val lastExpression = body.statements.lastOrNull()
         if (lastExpression is KtReturnExpression) return
 
-        val (defaultExpression, expressionToUnifyWith) =
+        val defaultExpression =
                 if (!generatorOptions.inTempFile && defaultValue != null && descriptor.controlFlow.outputValueBoxer.boxingRequired && lastExpression!!.isMultiLine()) {
                     val varNameValidator = NewDeclarationNameValidator(body, lastExpression, NewDeclarationNameValidator.Target.VARIABLES)
                     val resultVal = KotlinNameSuggester.suggestNamesByType(defaultValue.valueType, varNameValidator, null).first()
-                    val newDecl = body.addBefore(psiFactory.createDeclaration("val $resultVal = ${lastExpression.text}"), lastExpression) as KtProperty
+                    body.addBefore(psiFactory.createDeclaration("val $resultVal = ${lastExpression.text}"), lastExpression)
                     body.addBefore(psiFactory.createNewLine(), lastExpression)
-                    psiFactory.createExpression(resultVal) to newDecl.initializer!!
+                    psiFactory.createExpression(resultVal)
-                }
-                else {
-                    lastExpression to null
                 }
+                else lastExpression
 
         val returnExpression = descriptor.controlFlow.outputValueBoxer.getReturnExpression(getReturnArguments(defaultExpression), psiFactory) ?: return
 
+        @Suppress("NON_EXHAUSTIVE_WHEN")
         when(generatorOptions.target) {
             ExtractionTarget.LAZY_PROPERTY, ExtractionTarget.FAKE_LAMBDALIKE_FUNCTION -> {
                 // In the case of lazy property absence of default value means that output values are of OutputValue.Initializer type
@@ -580,11 +546,8 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
         }
 
         when {
-            defaultValue == null ->
+            defaultValue == null -> body.appendElement(returnExpression)
-                body.appendElement(returnExpression)
+            !defaultValue.callSiteReturn -> lastExpression!!.replaceWithReturn(returnExpression)
-
-            !defaultValue.callSiteReturn ->
-                replaceWithReturn(lastExpression!!, returnExpression, expressionToUnifyWith)
         }
 
         if (generatorOptions.allowExpressionBody) {
@@ -657,7 +620,7 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
         }
     }
 
-    if (generatorOptions.inTempFile) return ExtractionResult(this, declaration, Collections.emptyMap(), nameByOffset)
+    if (generatorOptions.inTempFile) return ExtractionResult(this, declaration, Collections.emptyMap())
 
     val replaceInitialOccurrence = {
         val arguments = descriptor.parameters.map { it.argumentText }
@@ -670,7 +633,7 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
         ShortenReferences.DEFAULT.process(declaration)
     }
 
-    if (generatorOptions.inTempFile) return ExtractionResult(this, declaration, emptyMap(), nameByOffset)
+    if (generatorOptions.inTempFile) return ExtractionResult(this, declaration, emptyMap())
 
     val duplicateReplacers = HashMap<KotlinPsiRange, () -> Unit>().apply {
         if (generatorOptions.delayInitialOccurrenceReplacement) {
@@ -688,5 +651,7 @@ fun ExtractionGeneratorConfiguration.generateDeclaration(
         }
     }
 
-    return ExtractionResult(this, declaration, duplicateReplacers, nameByOffset)
+    CodeStyleManager.getInstance(descriptor.extractionData.project).reformat(declaration)
+
+    return ExtractionResult(this, declaration, duplicateReplacers)
 }

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

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.refactoring.introduce.extractionEngine
 
 import com.intellij.psi.PsiElement
+import com.intellij.util.containers.MultiMap
 import org.jetbrains.kotlin.cfg.pseudocode.Pseudocode
 import org.jetbrains.kotlin.cfg.pseudocode.SingleType
 import org.jetbrains.kotlin.cfg.pseudocode.getElementValuesRecursively
@@ -51,6 +52,15 @@ import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 import java.util.*
 
+internal class ParametersInfo {
+    var errorMessage: AnalysisResult.ErrorMessage? = null
+    val originalRefToParameter = MultiMap.create<KtSimpleNameExpression, MutableParameter>()
+    val parameters = HashSet<MutableParameter>()
+    val typeParameters = HashSet<TypeParameter>()
+    val nonDenotableTypes = HashSet<KotlinType>()
+    val replacementMap = MultiMap.create<KtSimpleNameExpression, Replacement>()
+}
+
 internal fun ExtractionData.inferParametersInfo(
         commonParent: PsiElement,
         pseudocode: Pseudocode,
@@ -158,7 +168,7 @@ private fun ExtractionData.extractReceiver(
         )) return
 
         if (referencedClassifierDescriptor is ClassDescriptor) {
-            info.replacementMap.putValue(refInfo.offsetInBody, FqNameReplacement(originalDescriptor.getImportableDescriptor().fqNameSafe))
+            info.replacementMap.putValue(originalRef, FqNameReplacement(originalDescriptor.getImportableDescriptor().fqNameSafe))
         }
     }
     else {
@@ -244,12 +254,12 @@ private fun ExtractionData.extractReceiver(
                 }
             }
 
-            info.replacementMap.putValue(refInfo.offsetInBody,
+            val replacement = when {
-                    when {
+                isMemberExtensionFunction -> WrapInWithReplacement(parameter)
-                        isMemberExtensionFunction -> WrapInWithReplacement(parameter)
+                hasThisReceiver && extractThis -> AddPrefixReplacement(parameter)
-                        hasThisReceiver && extractThis -> AddPrefixReplacement(parameter)
+                else -> RenameReplacement(parameter)
-                        else -> RenameReplacement(parameter)
+            }
-                    })
+            info.replacementMap.putValue(originalRef, replacement)
         }
     }
 }
@@ -281,7 +291,6 @@ private fun suggestParameterType(
                    ?: if (receiverToExtract.exists()) receiverToExtract.type else null
 
                receiverToExtract is ImplicitReceiver -> {
-                   val calleeExpression = resolvedCall!!.call.calleeExpression
                    val typeByDataFlowInfo = if (useSmartCastsIfPossible) {
                        val dataFlowInfo = bindingContext.getDataFlowInfo(resolvedCall!!.call.callElement)
                        val possibleTypes = dataFlowInfo.getPossibleTypes(DataFlowValueFactory.createDataFlowValueForStableReceiver(receiverToExtract))

idea/src/org/jetbrains/kotlin/idea/refactoring/jetRefactoringUtil.kt

@@ -130,9 +130,9 @@ public fun PsiElement.getUsageContext(): PsiElement {
 public fun PsiElement.isInJavaSourceRoot(): Boolean =
         !JavaProjectRootsUtil.isOutsideJavaSourceRoot(getContainingFile())
 
-public inline fun KtFile.createTempCopy(textTransform: (String) -> String): KtFile {
+public fun KtFile.createTempCopy(text: String? = null): KtFile {
-    val tmpFile = KtPsiFactory(this).createAnalyzableFile(getName(), textTransform(getText() ?: ""), this)
+    val tmpFile = KtPsiFactory(this).createAnalyzableFile(name, text ?: this.text ?: "", this)
-    tmpFile.setOriginalFile(this)
+    tmpFile.originalFile = this
     tmpFile.suppressDiagnosticsInDebugMode = suppressDiagnosticsInDebugMode
     return tmpFile
 }

idea/testData/refactoring/extractFunction/basic/fragmentWithMultilineComment.kt.after

@@ -8,8 +8,8 @@ fun foo(a: Int): Int {
 
 private fun b(a: Int): Boolean {
     /*
-        test
+    test
-        */
+    */
     println(a)
     if (a > 0) return true
     return false