Refactoring

idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/CallableUsageReplacementStrategy.kt

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.resolve.calls.model.isReallySuccess
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 
 class CallableUsageReplacementStrategy(
-        private val replacement: ReplacementExpression
+        private val replacement: ReplacementCode
 ) : UsageReplacementStrategy {
 
     override fun createReplacer(usage: KtSimpleNameExpression): (() -> KtElement)? {
@@ -44,9 +44,8 @@ class CallableUsageReplacementStrategy(
                 createReplacer(nameExpression)!!.invoke()
             }
             else {
-                // copy replacement expression because it is modified by performCallReplacement
                 @Suppress("UNCHECKED_CAST")
-                ReplacementEngine.performCallReplacement(usage, bindingContext, resolvedCall, callElement, replacement.copy())
+                ReplacementEngine.performCallReplacement(usage, bindingContext, resolvedCall, callElement, replacement)
             }
         }
     }

idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/ClassUsageReplacementStrategy.kt

@@ -25,7 +25,7 @@ import org.jetbrains.kotlin.utils.addToStdlib.check
 
 class ClassUsageReplacementStrategy(
         typeReplacement: KtUserType?,
-        constructorReplacement: ReplacementExpression?,
+        constructorReplacement: ReplacementCode?,
         project: Project
 ) : UsageReplacementStrategy {
 

idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/ConstructedExpressionHolder.kt

@@ -36,35 +36,27 @@ import org.jetbrains.kotlin.types.ErrorUtils
 import org.jetbrains.kotlin.types.KotlinType
 import java.util.*
 
-internal abstract class ConstructedCodeHolder<TElement>(
-        var expression: KtExpression,
+internal abstract class ConstructedCodeHolder<TElement : KtElement>(
+        val replacement: MutableReplacementCode,
         val elementToBeReplaced: TElement,
         val bindingContext: BindingContext
 ) {
-    val psiFactory = KtPsiFactory(expression)
-
-    fun replaceExpression(oldExpression: KtExpression, newExpression: KtExpression): KtExpression {
-        assert(expression.isAncestor(oldExpression))
-        if (oldExpression == expression) {
-            expression = newExpression
-            return expression
-        }
-        else {
-            return oldExpression.replace(newExpression) as KtExpression
-        }
-    }
+    val psiFactory = KtPsiFactory(elementToBeReplaced)
 
     abstract fun finish(postProcessing: (PsiChildRange) -> PsiChildRange): TElement
 }
 
 internal class ConstructedAnnotationEntryHolder(
-        expression: KtExpression,
+        replacement: MutableReplacementCode,
         elementToBeReplaced: KtAnnotationEntry,
         bindingContext: BindingContext
-) : ConstructedCodeHolder<KtAnnotationEntry>(expression, elementToBeReplaced, bindingContext) {
+) : ConstructedCodeHolder<KtAnnotationEntry>(replacement, elementToBeReplaced, bindingContext) {
 
     override fun finish(postProcessing: (PsiChildRange) -> PsiChildRange): KtAnnotationEntry {
-        val dummyAnnotationEntry = createByPattern("@Dummy($0)", expression) { psiFactory.createAnnotationEntry(it) }
+        assert(replacement.mainExpression != null)
+        assert(replacement.statementsBefore.isEmpty())
+
+        val dummyAnnotationEntry = createByPattern("@Dummy($0)", replacement.mainExpression!!) { psiFactory.createAnnotationEntry(it) }
         val replaced = elementToBeReplaced.replace(dummyAnnotationEntry)
         var range = PsiChildRange.singleElement(replaced)
         range = postProcessing(range)
@@ -78,10 +70,10 @@ internal class ConstructedAnnotationEntryHolder(
 }
 
 internal class ConstructedExpressionHolder(
-        expression: KtExpression,
+        replacement: MutableReplacementCode,
         expressionToBeReplaced: KtExpression,
         bindingContext: BindingContext
-) : ConstructedCodeHolder<KtExpression>(expression, expressionToBeReplaced, bindingContext) {
+) : ConstructedCodeHolder<KtExpression>(replacement, expressionToBeReplaced, bindingContext) {
 
     private data class StatementToInsert<TStatement : KtExpression>(val statement: TStatement, val postProcessing: (TStatement) -> Unit)
 
@@ -104,7 +96,9 @@ internal class ConstructedExpressionHolder(
             (toInsert.postProcessing as (KtExpression) -> Unit).invoke(inserted)
         }
 
-        val replaced = elementToBeReplaced.replace(expression)
+        val replaced = elementToBeReplaced.replace(replacement.mainExpression!!) //TODO: support null here
+
+        //TODO: support code.statementsBefore
 
         var range = if (insertedStatements.isEmpty())
             PsiChildRange.singleElement(replaced)
@@ -123,7 +117,7 @@ internal class ConstructedExpressionHolder(
             nameSuggestion: String? = null,
             safeCall: Boolean = false
     ) {
-        assert(usages.all { expression.isAncestor(it, strict = true) })
+        assert(usages.all { replacement.containsStrictlyInside(it) })
 
         fun replaceUsages(name: Name) {
             val nameInCode = psiFactory.createExpression(name.render())
@@ -140,8 +134,8 @@ internal class ConstructedExpressionHolder(
             return Name.identifier(name)
         }
 
-        // checks that name is used (without receiver) inside expression being constructed but not inside usages that will be replaced
-        fun isNameUsed(name: String) = collectNameUsages(expression, name).any { nameUsage -> usages.none { it.isAncestor(nameUsage) } }
+        // checks that name is used (without receiver) inside code being constructed but not inside usages that will be replaced
+        fun isNameUsed(name: String) = collectNameUsages(replacement, name).any { nameUsage -> usages.none { it.isAncestor(nameUsage) } }
 
         if (!safeCall) {
             val block = elementToBeReplaced.parent as? KtBlockExpression
@@ -182,19 +176,24 @@ internal class ConstructedExpressionHolder(
             }
         }
 
+        //TODO: handle mainExpression == null and statementsBefore!
+
         val dot = if (safeCall) "?." else "."
 
-        expression = if (!isNameUsed("it")) {
+        replacement.mainExpression = if (!isNameUsed("it")) {
             replaceUsages(Name.identifier("it"))
-            psiFactory.createExpressionByPattern("$0${dot}let { $1 }", value, expression)
+            psiFactory.createExpressionByPattern("$0${dot}let { $1 }", value, replacement.mainExpression!!)
         }
         else {
             val name = suggestName { !isNameUsed(it) }
             replaceUsages(name)
-            psiFactory.createExpressionByPattern("$0${dot}let { $1 -> $2 }", value, name, expression)
+            psiFactory.createExpressionByPattern("$0${dot}let { $1 -> $2 }", value, name, replacement.mainExpression!!)
         }
     }
 
-    private fun collectNameUsages(scope: KtExpression, name: String)
-            = scope.collectDescendantsOfType<KtSimpleNameExpression> { it.getReceiverExpression() == null && it.getReferencedName() == name }
+    private fun collectNameUsages(scope: MutableReplacementCode, name: String): List<KtSimpleNameExpression> {
+        return scope.expressions.flatMap { expression ->
+            expression.collectDescendantsOfType<KtSimpleNameExpression> { it.getReceiverExpression() == null && it.getReferencedName() == name }
+        }
+    }
 }

idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/MutableReplacementCode.kt

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2010-2016 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.replacement
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.idea.core.copied
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
+import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
+import org.jetbrains.kotlin.psi.psiUtil.isAncestor
+import org.jetbrains.kotlin.utils.addToStdlib.singletonOrEmptyList
+
+internal class MutableReplacementCode(
+        var mainExpression: KtExpression?,
+        val statementsBefore: MutableList<KtExpression>,
+        val fqNamesToImport: MutableCollection<FqName>
+) {
+    fun replaceExpression(oldExpression: KtExpression, newExpression: KtExpression): KtExpression {
+        assert(oldExpression in this)
+
+        if (oldExpression == mainExpression) {
+            mainExpression = newExpression
+            return newExpression
+        }
+
+        val index = statementsBefore.indexOf(oldExpression)
+        if (index >= 0) {
+            statementsBefore[index] = newExpression
+            return newExpression
+        }
+
+        return oldExpression.replace(newExpression) as KtExpression
+    }
+
+    val expressions: Collection<KtExpression>
+        get() = statementsBefore + mainExpression.singletonOrEmptyList()
+
+    operator fun contains(element: PsiElement): Boolean {
+        return expressions.any { it.isAncestor(element) }
+    }
+
+    fun containsStrictlyInside(element: PsiElement): Boolean {
+        return expressions.any { it.isAncestor(element, strict = true) }
+    }
+}
+
+internal fun ReplacementCode.toMutable(): MutableReplacementCode {
+    return MutableReplacementCode(
+            mainExpression?.copied(),
+            statementsBefore.map { it.copied() }.toMutableList(),
+            fqNamesToImport.toMutableSet())
+}
+
+internal fun MutableReplacementCode.toNonMutable(): ReplacementCode {
+    return ReplacementCode(mainExpression, statementsBefore, fqNamesToImport)
+}
+
+internal inline fun <reified T : PsiElement> MutableReplacementCode.collectDescendantsOfType(noinline predicate: (T) -> Boolean = { true }): List<T> {
+    return expressions.flatMap { it.collectDescendantsOfType<T>({ true }, predicate) }
+}
+
+internal inline fun <reified T : PsiElement> MutableReplacementCode.forEachDescendantOfType(noinline action: (T) -> Unit) {
+    expressions.forEach { it.forEachDescendantOfType<T>(action) }
+}
+

idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/ReplacementBuilder.kt

@@ -19,49 +19,41 @@ package org.jetbrains.kotlin.idea.replacement
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
 import org.jetbrains.kotlin.descriptors.ValueParameterDescriptor
-import org.jetbrains.kotlin.idea.analysis.analyzeInContext
 import org.jetbrains.kotlin.idea.core.asExpression
-import org.jetbrains.kotlin.idea.core.copied
-import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.imports.importableFqName
 import org.jetbrains.kotlin.idea.intentions.InsertExplicitTypeArgumentsIntention
 import org.jetbrains.kotlin.idea.references.canBeResolvedViaImport
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
+import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
 import org.jetbrains.kotlin.psi.psiUtil.getReceiverExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.isReallySuccess
 import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
-import org.jetbrains.kotlin.resolve.descriptorUtil.module
-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.TypeUtils
-import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 import java.util.*
 
 class ReplacementBuilder(
         private val targetCallable: CallableDescriptor,
         private val resolutionFacade: ResolutionFacade
 ) {
-    fun buildReplacementExpression(
-            expression: KtExpression,
-            resolutionScope: LexicalScope,
-            expectedType: KotlinType = TypeUtils.NO_EXPECTED_TYPE,
-            importFqNames: Collection<FqName> = emptyList(),
-            copyExpression: Boolean = true
-    ): ReplacementExpression {
-        @Suppress("NAME_SHADOWING")
-        var expression = if (copyExpression) expression.copied() else expression
+    private val psiFactory = KtPsiFactory(resolutionFacade.project)
 
-        var bindingContext = analyzeInContext(expression, resolutionScope, expectedType)
+    fun buildReplacementCode(
+            mainExpression: KtExpression?,
+            statementsBefore: List<KtExpression>,
+            analyze: () -> BindingContext,
+            importFqNames: Collection<FqName> = emptyList()
+    ): ReplacementCode {
+        var bindingContext = analyze()
+
+        val result = MutableReplacementCode(mainExpression, statementsBefore.toMutableList(), importFqNames.toMutableSet())
 
         val typeArgsToAdd = ArrayList<Pair<KtCallExpression, KtTypeArgumentList>>()
-        expression.forEachDescendantOfType<KtCallExpression> {
+        result.forEachDescendantOfType<KtCallExpression> {
             if (InsertExplicitTypeArgumentsIntention.isApplicableTo(it, bindingContext)) {
                 typeArgsToAdd.add(it to InsertExplicitTypeArgumentsIntention.createTypeArguments(it, bindingContext)!!)
             }
@@ -73,28 +65,25 @@ class ReplacementBuilder(
             }
 
             // reanalyze expression - new usages of type parameters may be added
-            bindingContext = analyzeInContext(expression, resolutionScope, expectedType)
+            bindingContext = analyze()
         }
 
         val receiversToAdd = ArrayList<Pair<KtExpression, KtExpression>>()
-        val resultImportFqNames = importFqNames.toMutableSet()
 
-        val psiFactory = KtPsiFactory(expression)
-
-        expression.forEachDescendantOfType<KtSimpleNameExpression> { expression ->
+        result.forEachDescendantOfType<KtSimpleNameExpression> { expression ->
             val target = bindingContext[BindingContext.REFERENCE_TARGET, expression] ?: return@forEachDescendantOfType
 
             //TODO: other types of references ('[]' etc)
             if (expression.mainReference.canBeResolvedViaImport(target)) {
-                resultImportFqNames.add(target.importableFqName!!)
+                result.fqNamesToImport.add(target.importableFqName!!)
             }
 
             if (expression.getReceiverExpression() == null) {
                 if (target is ValueParameterDescriptor && target.containingDeclaration == targetCallable) {
-                    expression.putCopyableUserData(ReplacementExpression.PARAMETER_USAGE_KEY, target.name)
+                    expression.putCopyableUserData(ReplacementCode.PARAMETER_USAGE_KEY, target.name)
                 }
                 else if (target is TypeParameterDescriptor && target.containingDeclaration == targetCallable) {
-                    expression.putCopyableUserData(ReplacementExpression.TYPE_PARAMETER_USAGE_KEY, target.name)
+                    expression.putCopyableUserData(ReplacementCode.TYPE_PARAMETER_USAGE_KEY, target.name)
                 }
 
                 val resolvedCall = expression.getResolvedCall(bindingContext)
@@ -104,6 +93,7 @@ class ReplacementBuilder(
                     else
                         resolvedCall.dispatchReceiver
                     if (receiver is ImplicitReceiver) {
+                        val resolutionScope = expression.getResolutionScope(bindingContext, resolutionFacade)
                         val receiverExpression = receiver.asExpression(resolutionScope, psiFactory)
                         if (receiverExpression != null) {
                             receiversToAdd.add(expression to receiverExpression)
@@ -116,24 +106,10 @@ class ReplacementBuilder(
         // add receivers in reverse order because arguments of a call were processed after the callee's name
         for ((expr, receiverExpression) in receiversToAdd.asReversed()) {
             val expressionToReplace = expr.parent as? KtCallExpression ?: expr
-            val newExpr = expressionToReplace.replaced(psiFactory.createExpressionByPattern("$0.$1", receiverExpression, expressionToReplace))
-            if (expressionToReplace == expression) {
-                expression = newExpr
-            }
+            result.replaceExpression(expressionToReplace,
+                                              psiFactory.createExpressionByPattern("$0.$1", receiverExpression, expressionToReplace))
         }
 
-        return ReplacementExpression(expression, resultImportFqNames)
-    }
-
-    private fun analyzeInContext(expression: KtExpression, scope: LexicalScope, expectedType: KotlinType): BindingContext {
-        val module = scope.ownerDescriptor.module
-        val frontendService = if (module.builtIns.builtInsModule == module) {
-            // TODO: doubtful place, do we require this module or not? Built-ins module doesn't have some necessary components...
-            resolutionFacade.getFrontendService(ExpressionTypingServices::class.java)
-        }
-        else {
-            resolutionFacade.getFrontendService(module, ExpressionTypingServices::class.java)
-        }
-        return expression.analyzeInContext(scope, expectedType = expectedType, expressionTypingServices = frontendService)
+        return result.toNonMutable()
     }
 }

idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/ReplacementCode.kt

@@ -17,9 +17,8 @@
 package org.jetbrains.kotlin.idea.replacement
 
 import com.intellij.openapi.util.Key
-import org.jetbrains.kotlin.idea.core.copied
-import org.jetbrains.kotlin.idea.replacement.ReplacementExpression.Companion.PARAMETER_USAGE_KEY
-import org.jetbrains.kotlin.idea.replacement.ReplacementExpression.Companion.TYPE_PARAMETER_USAGE_KEY
+import org.jetbrains.kotlin.idea.replacement.ReplacementCode.Companion.PARAMETER_USAGE_KEY
+import org.jetbrains.kotlin.idea.replacement.ReplacementCode.Companion.TYPE_PARAMETER_USAGE_KEY
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtExpression
@@ -31,16 +30,15 @@ import org.jetbrains.kotlin.psi.KtExpression
  * * All external symbols to be imported should be either referenced via fully-qualified form or included into [fqNamesToImport]
  * * All usages of value parameters (of our callable) should be marked with [PARAMETER_USAGE_KEY] copyable user data (holds the name of the corresponding parameter)
  * * All usages of type parameters (of our callable) should be marked with [TYPE_PARAMETER_USAGE_KEY] copyable user data (holds the name of the corresponding type parameter)
- * Use [ReplacementBuilder.buildReplacementExpression].
+ * Use [ReplacementBuilder.buildReplacementCode].
  */
-class ReplacementExpression(
-        val expression: KtExpression,
+class ReplacementCode(
+        val mainExpression: KtExpression?,
+        val statementsBefore: List<KtExpression>,
         val fqNamesToImport: Collection<FqName>
 ) {
-    fun copy() = ReplacementExpression(expression.copied(), fqNamesToImport)
-
     companion object {
         val PARAMETER_USAGE_KEY: Key<Name> = Key("PARAMETER_USAGE")
         val TYPE_PARAMETER_USAGE_KEY: Key<Name> = Key("TYPE_PARAMETER_USAGE")
     }
-}
+}

idea/idea-replacement-engine/src/org/jetbrains/kotlin/idea/replacement/ReplacementEngine.kt

@@ -52,8 +52,11 @@ object ReplacementEngine {
             bindingContext: BindingContext,
             resolvedCall: ResolvedCall<out CallableDescriptor>,
             callElement: TCallElement,
-            replacement: ReplacementExpression
+            replacement: ReplacementCode
     ): KtElement {
+        @Suppress("NAME_SHADOWING")
+        val replacement = replacement.toMutable()
+
         val project = element.project
         val psiFactory = KtPsiFactory(project)
         val descriptor = resolvedCall.resultingDescriptor
@@ -82,34 +85,34 @@ object ReplacementEngine {
         receiver?.mark(RECEIVER_VALUE_KEY)
 
         val codeHolder = when (elementToBeReplaced) {
-            is KtExpression -> ConstructedExpressionHolder(replacement.expression, elementToBeReplaced, bindingContext)
-            is KtAnnotationEntry -> ConstructedAnnotationEntryHolder(replacement.expression, elementToBeReplaced, bindingContext)
+            is KtExpression -> ConstructedExpressionHolder(replacement, elementToBeReplaced, bindingContext)
+            is KtAnnotationEntry -> ConstructedAnnotationEntryHolder(replacement, elementToBeReplaced, bindingContext)
             else -> error("Unsupported element")
         }
 
         //TODO: this@
-        for (thisExpression in replacement.expression.collectDescendantsOfType<KtThisExpression>()) {
+        for (thisExpression in replacement.collectDescendantsOfType<KtThisExpression>()) {
             if (receiver != null) {
-                codeHolder.replaceExpression(thisExpression, receiver)
+                replacement.replaceExpression(thisExpression, receiver)
             }
             else {
                 thisExpression.mark(RECEIVER_VALUE_KEY)
             }
         }
 
-        val introduceValuesForParameters = codeHolder.processValueParameterUsages(resolvedCall, project)
+        val introduceValuesForParameters = processValueParameterUsages(replacement, resolvedCall, bindingContext, project)
 
-        codeHolder.processTypeParameterUsages(resolvedCall)
+        processTypeParameterUsages(replacement, resolvedCall)
 
         if (qualifiedExpression is KtSafeQualifiedExpression) {
-            (codeHolder as ConstructedExpressionHolder).wrapExpressionForSafeCall(receiver!!, receiverType)
+            (codeHolder as ConstructedExpressionHolder).wrapCodeForSafeCall(receiver!!, receiverType)
         }
         else if (callElement is KtBinaryExpression && callElement.operationToken == KtTokens.IDENTIFIER) {
-            codeHolder.keepInfixFormIfPossible()
+            keepInfixFormIfPossible(replacement)
         }
 
         if (receiver != null && codeHolder is ConstructedExpressionHolder) {
-            val thisReplaced = codeHolder.expression.collectDescendantsOfType<KtExpression> { it[RECEIVER_VALUE_KEY] }
+            val thisReplaced = replacement.collectDescendantsOfType<KtExpression> { it[RECEIVER_VALUE_KEY] }
             if (receiver.shouldKeepValue(thisReplaced.size)) {
                 codeHolder.introduceValue(receiver, receiverType, thisReplaced)
             }
@@ -117,7 +120,7 @@ object ReplacementEngine {
 
         if (codeHolder is ConstructedExpressionHolder) {
             for ((parameter, value, valueType) in introduceValuesForParameters) {
-                val usagesReplaced = codeHolder.expression.collectDescendantsOfType<KtExpression> { it[PARAMETER_VALUE_KEY] == parameter }
+                val usagesReplaced = replacement.collectDescendantsOfType<KtExpression> { it[PARAMETER_VALUE_KEY] == parameter }
                 codeHolder.introduceValue(value, valueType, usagesReplaced, nameSuggestion = parameter.name.asString())
             }
         }
@@ -133,8 +136,10 @@ object ReplacementEngine {
         }
     }
 
-    private fun ConstructedCodeHolder<*>.processValueParameterUsages(
+    private fun processValueParameterUsages(
+            replacement: MutableReplacementCode,
             resolvedCall: ResolvedCall<out CallableDescriptor>,
+            bindingContext: BindingContext,
             project: Project
     ): Collection<IntroduceValueForParameter> {
         val introduceValuesForParameters = ArrayList<IntroduceValueForParameter>()
@@ -147,8 +152,8 @@ object ReplacementEngine {
             argument.expression.put(PARAMETER_VALUE_KEY, parameter)
 
             val parameterName = parameter.name
-            val usages = expression.collectDescendantsOfType<KtExpression> {
-                it[ReplacementExpression.PARAMETER_USAGE_KEY] == parameterName
+            val usages = replacement.collectDescendantsOfType<KtExpression> {
+                it[ReplacementCode.PARAMETER_USAGE_KEY] == parameterName
             }
             usages.forEach {
                 val usageArgument = it.parent as? KtValueArgument
@@ -158,7 +163,7 @@ object ReplacementEngine {
                 if (argument.isDefaultValue) {
                     usageArgument?.mark(DEFAULT_PARAMETER_VALUE_KEY)
                 }
-                replaceExpression(it, argument.expression)
+                replacement.replaceExpression(it, argument.expression)
             }
 
             //TODO: sometimes we need to add explicit type arguments here because we don't have expected type in the new context
@@ -176,7 +181,7 @@ object ReplacementEngine {
             val value: KtExpression,
             val valueType: KotlinType?)
 
-    private fun ConstructedCodeHolder<*>.processTypeParameterUsages(resolvedCall: ResolvedCall<out CallableDescriptor>) {
+    private fun processTypeParameterUsages(replacement: MutableReplacementCode, resolvedCall: ResolvedCall<out CallableDescriptor>) {
         val typeParameters = resolvedCall.resultingDescriptor.original.typeParameters
 
         val callElement = resolvedCall.call.callElement
@@ -186,8 +191,8 @@ object ReplacementEngine {
 
         for ((index, typeParameter) in typeParameters.withIndex()) {
             val parameterName = typeParameter.name
-            val usages = expression.collectDescendantsOfType<KtExpression> {
-                it[ReplacementExpression.TYPE_PARAMETER_USAGE_KEY] == parameterName
+            val usages = replacement.collectDescendantsOfType<KtExpression> {
+                it[ReplacementCode.TYPE_PARAMETER_USAGE_KEY] == parameterName
             }
 
             val factory = KtPsiFactory(callElement)
@@ -209,7 +214,7 @@ object ReplacementEngine {
                     val arguments =
                             if (typeElement is KtUserType && KotlinBuiltIns.isArray(type)) typeElement.typeArgumentList?.text.orEmpty()
                             else ""
-                    replaceExpression(usage, KtPsiFactory(usage).createExpression(
+                    replacement.replaceExpression(usage, KtPsiFactory(usage).createExpression(
                             IdeDescriptorRenderers.SOURCE_CODE.renderClassifierName(typeClassifier) + arguments
                     ))
                 }
@@ -218,36 +223,54 @@ object ReplacementEngine {
                 }
                 else {
                     //TODO: tests for this?
-                    replaceExpression(usage, KtPsiFactory(usage).createExpression(typeElement.text))
+                    replacement.replaceExpression(usage, KtPsiFactory(usage).createExpression(typeElement.text))
                 }
             }
         }
     }
 
-    private fun ConstructedExpressionHolder.wrapExpressionForSafeCall(receiver: KtExpression, receiverType: KotlinType?) {
-        val qualified = expression as? KtQualifiedExpression
-        if (qualified != null) {
-            if (qualified.receiverExpression[RECEIVER_VALUE_KEY]) {
-                if (qualified is KtSafeQualifiedExpression) return // already safe
-                val selector = qualified.selectorExpression
-                if (selector != null) {
-                    expression = psiFactory.createExpressionByPattern("$0?.$1", receiver, selector)
-                    return
+    private fun ConstructedExpressionHolder.wrapCodeForSafeCall(receiver: KtExpression, receiverType: KotlinType?) {
+        if (replacement.statementsBefore.isEmpty()) {
+            val qualified = replacement.mainExpression as? KtQualifiedExpression
+            if (qualified != null) {
+                if (qualified.receiverExpression[RECEIVER_VALUE_KEY]) {
+                    if (qualified is KtSafeQualifiedExpression) return // already safe
+                    val selector = qualified.selectorExpression
+                    if (selector != null) {
+                        replacement.mainExpression = KtPsiFactory(receiver).createExpressionByPattern("$0?.$1", receiver, selector)
+                        return
+                    }
                 }
             }
         }
 
         if (elementToBeReplaced.isUsedAsExpression(bindingContext)) {
-            val thisReplaced = expression.collectDescendantsOfType<KtExpression> { it[RECEIVER_VALUE_KEY] }
+            val thisReplaced = replacement.collectDescendantsOfType<KtExpression> { it[RECEIVER_VALUE_KEY] }
             introduceValue(receiver, receiverType, thisReplaced, safeCall = true)
         }
         else {
-            expression = psiFactory.createExpressionByPattern("if ($0 != null) { $1 }", receiver, expression)
+            val ifExpression = KtPsiFactory(receiver).buildExpression {
+                appendFixedText("if (")
+                appendExpression(receiver)
+                appendFixedText("!=null) {")
+                replacement.statementsBefore.forEach {
+                    appendExpression(it)
+                    appendFixedText("\n")
+                }
+                replacement.mainExpression?.let {
+                    appendExpression(it)
+                    appendFixedText("\n")
+                }
+                appendFixedText("}")
+            }
+            replacement.mainExpression = ifExpression
+            replacement.statementsBefore.clear()
         }
     }
 
-    private fun ConstructedCodeHolder<*>.keepInfixFormIfPossible() {
-        val dotQualified = expression as? KtDotQualifiedExpression ?: return
+    private fun keepInfixFormIfPossible(replacement: MutableReplacementCode) {
+        if (replacement.statementsBefore.isNotEmpty()) return
+        val dotQualified = replacement.mainExpression as? KtDotQualifiedExpression ?: return
         val receiver = dotQualified.receiverExpression
         if (!receiver[RECEIVER_VALUE_KEY]) return
         val call = dotQualified.selectorExpression as? KtCallExpression ?: return
@@ -255,7 +278,7 @@ object ReplacementEngine {
         val argument = call.valueArguments.singleOrNull() ?: return
         if (argument.getArgumentName() != null) return
         val argumentExpression = argument.getArgumentExpression() ?: return
-        expression = psiFactory.createExpressionByPattern("$0 ${nameExpression.text} $1", receiver, argumentExpression)
+        replacement.mainExpression = KtPsiFactory(receiver).createExpressionByPattern("$0 ${nameExpression.text} $1", receiver, argumentExpression)
     }
 
     private fun KtExpression?.shouldKeepValue(usageCount: Int): Boolean {
@@ -305,13 +328,13 @@ object ReplacementEngine {
                 val (expression, parameterUsages) = defaultValue
 
                 for ((param, usages) in parameterUsages) {
-                    usages.forEach { it.put(ReplacementExpression.PARAMETER_USAGE_KEY, param.name) }
+                    usages.forEach { it.put(ReplacementCode.PARAMETER_USAGE_KEY, param.name) }
                 }
 
                 val expressionCopy = expression.copied()
 
                 // clean up user data in original
-                expression.forEachDescendantOfType<KtExpression> { it.clear(ReplacementExpression.PARAMETER_USAGE_KEY) }
+                expression.forEachDescendantOfType<KtExpression> { it.clear(ReplacementCode.PARAMETER_USAGE_KEY) }
 
                 return Argument(expressionCopy, null/*TODO*/, isDefaultValue = true)
             }
@@ -382,8 +405,8 @@ object ReplacementEngine {
             // clean up user data
             it.forEachDescendantOfType<KtExpression> {
                 it.clear(USER_CODE_KEY)
-                it.clear(ReplacementExpression.PARAMETER_USAGE_KEY)
-                it.clear(ReplacementExpression.TYPE_PARAMETER_USAGE_KEY)
+                it.clear(ReplacementCode.PARAMETER_USAGE_KEY)
+                it.clear(ReplacementCode.TYPE_PARAMETER_USAGE_KEY)
                 it.clear(PARAMETER_VALUE_KEY)
                 it.clear(RECEIVER_VALUE_KEY)
                 it.clear(WAS_FUNCTION_LITERAL_ARGUMENT_KEY)

idea/src/org/jetbrains/kotlin/idea/quickfix/replaceWith/ReplaceWithAnnotationAnalyzer.kt

@@ -17,11 +17,12 @@
 package org.jetbrains.kotlin.idea.quickfix.replaceWith
 
 import org.jetbrains.kotlin.descriptors.*
+import org.jetbrains.kotlin.idea.analysis.analyzeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.resolveImportReference
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.idea.replacement.ReplacementBuilder
-import org.jetbrains.kotlin.idea.replacement.ReplacementExpression
+import org.jetbrains.kotlin.idea.replacement.ReplacementCode
 import org.jetbrains.kotlin.idea.resolve.ResolutionFacade
 import org.jetbrains.kotlin.idea.resolve.frontendService
 import org.jetbrains.kotlin.name.FqName
@@ -39,6 +40,7 @@ import org.jetbrains.kotlin.resolve.scopes.*
 import org.jetbrains.kotlin.resolve.scopes.utils.chainImportingScopes
 import org.jetbrains.kotlin.resolve.scopes.utils.memberScopeAsImportingScope
 import org.jetbrains.kotlin.storage.LockBasedStorageManager
+import org.jetbrains.kotlin.types.expressions.ExpressionTypingServices
 import java.util.*
 
 data class ReplaceWith(val pattern: String, val imports: List<String>)
@@ -48,7 +50,7 @@ object ReplaceWithAnnotationAnalyzer {
             annotation: ReplaceWith,
             symbolDescriptor: CallableDescriptor,
             resolutionFacade: ResolutionFacade
-    ): ReplacementExpression? {
+    ): ReplacementCode? {
         val originalDescriptor = (if (symbolDescriptor is CallableMemberDescriptor)
             DescriptorUtils.unwrapFakeOverride(symbolDescriptor)
         else
@@ -60,7 +62,7 @@ object ReplaceWithAnnotationAnalyzer {
             annotation: ReplaceWith,
             symbolDescriptor: CallableDescriptor,
             resolutionFacade: ResolutionFacade
-    ): ReplacementExpression? {
+    ): ReplacementCode? {
         val psiFactory = KtPsiFactory(resolutionFacade.project)
         val expression = try {
             psiFactory.createExpression(annotation.pattern)
@@ -75,8 +77,20 @@ object ReplaceWithAnnotationAnalyzer {
         val scope = getResolutionScope(symbolDescriptor, symbolDescriptor,
                                        listOf(explicitImportsScope) + defaultImportsScopes) ?: return null
 
+        val expressionTypingServices = if (module.builtIns.builtInsModule == module) {
+            // TODO: doubtful place, do we require this module or not? Built-ins module doesn't have some necessary components...
+            resolutionFacade.getFrontendService(ExpressionTypingServices::class.java)
+        }
+        else {
+            resolutionFacade.getFrontendService(module, ExpressionTypingServices::class.java)
+        }
+
+        fun analyzeExpression(): BindingContext {
+            return expression.analyzeInContext(scope, expressionTypingServices = expressionTypingServices)
+        }
+
         return ReplacementBuilder(symbolDescriptor, resolutionFacade)
-                .buildReplacementExpression(expression, scope, importFqNames = importFqNames(annotation), copyExpression = false)
+                .buildReplacementCode(expression, emptyList(), ::analyzeExpression, importFqNames = importFqNames(annotation))
     }
 
     fun analyzeClassReplacement(

idea/src/org/jetbrains/kotlin/idea/refactoring/inline/KotlinInlineFunctionHandler.kt

@@ -25,13 +25,16 @@ import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.refactoring.RefactoringBundle
 import org.jetbrains.kotlin.descriptors.SimpleFunctionDescriptor
 import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.idea.analysis.analyzeInContext
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToDescriptor
+import org.jetbrains.kotlin.idea.core.copied
 import org.jetbrains.kotlin.idea.replacement.CallableUsageReplacementStrategy
 import org.jetbrains.kotlin.idea.replacement.ReplacementBuilder
 import org.jetbrains.kotlin.idea.replacement.replaceUsagesInWholeProject
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.types.TypeUtils
 
 class KotlinInlineFunctionHandler: InlineActionHandler() {
@@ -53,8 +56,15 @@ class KotlinInlineFunctionHandler: InlineActionHandler() {
             descriptor.returnType ?: TypeUtils.NO_EXPECTED_TYPE
         else
             TypeUtils.NO_EXPECTED_TYPE
+
+        val expression = bodyExpression.copied()
+
+        fun analyzeExpression(): BindingContext {
+            return expression.analyzeInContext(bodyExpression.getResolutionScope(), contextExpression = bodyExpression, expectedType = expectedType)
+        }
+
         val replacement = ReplacementBuilder(descriptor, element.getResolutionFacade())
-                .buildReplacementExpression(bodyExpression, bodyExpression.getResolutionScope(), expectedType)
+                .buildReplacementCode(expression, emptyList(), ::analyzeExpression)
 
         val commandName = RefactoringBundle.message("inline.command", element.name)
         CallableUsageReplacementStrategy(replacement)