Very basic version of inline function for block body

2016-10-21 19:13:27 +03:00
parent 66c815968b
commit a4aa9bab8d
16 changed files with 388 additions and 230 deletions
@@ -1,199 +0,0 @@
-/*
- * 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 org.jetbrains.kotlin.idea.analysis.computeTypeInContext
-import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
-import org.jetbrains.kotlin.idea.core.KotlinNameSuggester
-import org.jetbrains.kotlin.idea.core.replaced
-import org.jetbrains.kotlin.idea.core.setType
-import org.jetbrains.kotlin.idea.util.getResolutionScope
-import org.jetbrains.kotlin.name.Name
-import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
-import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
-import org.jetbrains.kotlin.psi.psiUtil.getReceiverExpression
-import org.jetbrains.kotlin.psi.psiUtil.isAncestor
-import org.jetbrains.kotlin.renderer.render
-import org.jetbrains.kotlin.resolve.BindingContext
-import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfo
-import org.jetbrains.kotlin.resolve.scopes.utils.findLocalVariable
-import org.jetbrains.kotlin.types.ErrorUtils
-import org.jetbrains.kotlin.types.KotlinType
-import java.util.*
-
-internal abstract class ConstructedCodeHolder<TElement : KtElement>(
-        val replacement: MutableReplacementCode,
-        val elementToBeReplaced: TElement,
-        val bindingContext: BindingContext
-) {
-    val psiFactory = KtPsiFactory(elementToBeReplaced)
-
-    abstract fun finish(postProcessing: (PsiChildRange) -> PsiChildRange): TElement
-}
-
-internal class ConstructedAnnotationEntryHolder(
-        replacement: MutableReplacementCode,
-        elementToBeReplaced: KtAnnotationEntry,
-        bindingContext: BindingContext
-) : ConstructedCodeHolder<KtAnnotationEntry>(replacement, elementToBeReplaced, bindingContext) {
-
-    override fun finish(postProcessing: (PsiChildRange) -> PsiChildRange): KtAnnotationEntry {
-        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)
-
-        assert(range.first == range.last)
-        assert(range.first is KtAnnotationEntry)
-        val annotationEntry = range.first as KtAnnotationEntry
-        val text = annotationEntry.valueArguments.single().getArgumentExpression()!!.text
-        return annotationEntry.replaced(psiFactory.createAnnotationEntry("@" + text))
-    }
-}
-
-internal class ConstructedExpressionHolder(
-        replacement: MutableReplacementCode,
-        expressionToBeReplaced: KtExpression,
-        bindingContext: BindingContext
-) : ConstructedCodeHolder<KtExpression>(replacement, expressionToBeReplaced, bindingContext) {
-
-    private data class StatementToInsert<TStatement : KtExpression>(val statement: TStatement, val postProcessing: (TStatement) -> Unit)
-
-    private val statementsToInsert = ArrayList<StatementToInsert<*>>()
-
-    private fun <TStatement : KtExpression> addStatementToInsert(statement: TStatement, postProcessing: (TStatement) -> Unit = {}) {
-        statementsToInsert.add(StatementToInsert(statement, postProcessing))
-    }
-
-    override fun finish(postProcessing: (PsiChildRange) -> PsiChildRange): KtExpression {
-        val insertedStatements = ArrayList<KtExpression>()
-        for (toInsert in statementsToInsert.asReversed()) { //TODO: do we need it?
-            val block = elementToBeReplaced.parent as KtBlockExpression //TODO
-
-            val inserted = block.addBefore(toInsert.statement, elementToBeReplaced) as KtExpression
-            block.addBefore(psiFactory.createNewLine(), elementToBeReplaced)
-            insertedStatements.add(inserted)
-
-            @Suppress("UNCHECKED_CAST")
-            (toInsert.postProcessing as (KtExpression) -> Unit).invoke(inserted)
-        }
-
-        val replaced = elementToBeReplaced.replace(replacement.mainExpression!!) //TODO: support null here
-
-        //TODO: support code.statementsBefore
-
-        var range = if (insertedStatements.isEmpty())
-            PsiChildRange.singleElement(replaced)
-        else
-            PsiChildRange(insertedStatements.first(), replaced)
-
-        range = postProcessing(range)
-
-        return range.last as KtExpression
-    }
-
-    fun introduceValue(
-            value: KtExpression,
-            valueType: KotlinType?,
-            usages: Collection<KtExpression>,
-            nameSuggestion: String? = null,
-            safeCall: Boolean = false
-    ) {
-        assert(usages.all { replacement.containsStrictlyInside(it) })
-
-        fun replaceUsages(name: Name) {
-            val nameInCode = psiFactory.createExpression(name.render())
-            for (usage in usages) {
-                usage.replace(nameInCode)
-            }
-        }
-
-        fun suggestName(validator: (String) -> Boolean): Name {
-            val name = if (nameSuggestion != null)
-                KotlinNameSuggester.suggestNameByName(nameSuggestion, validator)
-            else
-                KotlinNameSuggester.suggestNamesByExpressionOnly(value, bindingContext, validator, "t").first()
-            return Name.identifier(name)
-        }
-
-        // 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
-            if (block != null) {
-                val resolutionScope = elementToBeReplaced.getResolutionScope(bindingContext, elementToBeReplaced.getResolutionFacade())
-
-                if (usages.isNotEmpty()) {
-                    var explicitType: KotlinType? = null
-                    if (valueType != null && !ErrorUtils.containsErrorType(valueType)) {
-                        val valueTypeWithoutExpectedType = value.computeTypeInContext(
-                                resolutionScope,
-                                elementToBeReplaced,
-                                dataFlowInfo = bindingContext.getDataFlowInfo(elementToBeReplaced)
-                        )
-                        if (valueTypeWithoutExpectedType == null || ErrorUtils.containsErrorType(valueTypeWithoutExpectedType)) {
-                            explicitType = valueType
-                        }
-                    }
-
-                    val name = suggestName { name ->
-                        resolutionScope.findLocalVariable(Name.identifier(name)) == null && !isNameUsed(name)
-                    }
-
-                    val declaration = psiFactory.createDeclarationByPattern<KtVariableDeclaration>("val $0 = $1", name, value)
-                    addStatementToInsert(declaration,
-                                         postProcessing = {
-                                             if (explicitType != null) {
-                                                 it.setType(explicitType!!)
-                                             }
-                                         })
-
-                    replaceUsages(name)
-                }
-                else {
-                    addStatementToInsert(value)
-                }
-                return
-            }
-        }
-
-        //TODO: handle mainExpression == null and statementsBefore!
-
-        val dot = if (safeCall) "?." else "."
-
-        replacement.mainExpression = if (!isNameUsed("it")) {
-            replaceUsages(Name.identifier("it"))
-            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, replacement.mainExpression!!)
-        }
-    }
-
-    private fun collectNameUsages(scope: MutableReplacementCode, name: String): List<KtSimpleNameExpression> {
-        return scope.expressions.flatMap { expression ->
-            expression.collectDescendantsOfType<KtSimpleNameExpression> { it.getReceiverExpression() == null && it.getReferencedName() == name }
-        }
-    }
-}
@@ -23,13 +23,26 @@ 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.SmartList
 import org.jetbrains.kotlin.utils.addToStdlib.singletonOrEmptyList
+import java.util.*

 internal class MutableReplacementCode(
        var mainExpression: KtExpression?,
        val statementsBefore: MutableList<KtExpression>,
        val fqNamesToImport: MutableCollection<FqName>
 ) {
+    private val _postInsertionActions = HashMap<KtExpression, SmartList<(KtExpression) -> KtExpression>>()
+
+    val postInsertionActions: Map<KtExpression, List<(KtExpression) -> KtExpression>>
+        get() = _postInsertionActions
+
+    fun <TStatement : KtExpression> addPostInsertionAction(statementBefore: TStatement, action: (TStatement) -> TStatement) {
+        assert(statementBefore in statementsBefore)
+        @Suppress("UNCHECKED_CAST")
+        _postInsertionActions.getOrPut(statementBefore) { SmartList() }.add(action as (KtExpression) -> KtExpression)
+    }
+
    fun replaceExpression(oldExpression: KtExpression, newExpression: KtExpression): KtExpression {
        assert(oldExpression in this)

@@ -62,7 +62,6 @@ object ReplacementEngine {
        val descriptor = resolvedCall.resultingDescriptor
        val file = element.getContainingKtFile()

-        val qualifiedExpression = callElement.getQualifiedExpressionForSelector()
        val elementToBeReplaced = when (callElement) {
            is KtExpression -> callElement.getQualifiedExpressionForSelectorOrThis()
            else -> callElement
@@ -84,12 +83,6 @@ object ReplacementEngine {

        receiver?.mark(RECEIVER_VALUE_KEY)

-        val codeHolder = when (elementToBeReplaced) {
-            is KtExpression -> ConstructedExpressionHolder(replacement, elementToBeReplaced, bindingContext)
-            is KtAnnotationEntry -> ConstructedAnnotationEntryHolder(replacement, elementToBeReplaced, bindingContext)
-            else -> error("Unsupported element")
-        }
-
        //TODO: this@
        for (thisExpression in replacement.collectDescendantsOfType<KtThisExpression>()) {
            if (receiver != null) {
@@ -104,24 +97,24 @@ object ReplacementEngine {

        processTypeParameterUsages(replacement, resolvedCall)

-        if (qualifiedExpression is KtSafeQualifiedExpression) {
-            (codeHolder as ConstructedExpressionHolder).wrapCodeForSafeCall(receiver!!, receiverType)
+        if (elementToBeReplaced is KtSafeQualifiedExpression) {
+            wrapCodeForSafeCall(replacement, receiver!!, receiverType, elementToBeReplaced, bindingContext)
        }
        else if (callElement is KtBinaryExpression && callElement.operationToken == KtTokens.IDENTIFIER) {
            keepInfixFormIfPossible(replacement)
        }

-        if (receiver != null && codeHolder is ConstructedExpressionHolder) {
-            val thisReplaced = replacement.collectDescendantsOfType<KtExpression> { it[RECEIVER_VALUE_KEY] }
-            if (receiver.shouldKeepValue(thisReplaced.size)) {
-                codeHolder.introduceValue(receiver, receiverType, thisReplaced)
+        if (elementToBeReplaced is KtExpression) {
+            if (receiver != null) {
+                val thisReplaced = replacement.collectDescendantsOfType<KtExpression> { it[RECEIVER_VALUE_KEY] }
+                if (receiver.shouldKeepValue(thisReplaced.size)) {
+                    replacement.introduceValue(receiver, receiverType, thisReplaced, elementToBeReplaced)
+                }
            }
-        }

-        if (codeHolder is ConstructedExpressionHolder) {
            for ((parameter, value, valueType) in introduceValuesForParameters) {
                val usagesReplaced = replacement.collectDescendantsOfType<KtExpression> { it[PARAMETER_VALUE_KEY] == parameter }
-                codeHolder.introduceValue(value, valueType, usagesReplaced, nameSuggestion = parameter.name.asString())
+                replacement.introduceValue(value, valueType, usagesReplaced, elementToBeReplaced, nameSuggestion = parameter.name.asString())
            }
        }

@@ -129,11 +122,17 @@ object ReplacementEngine {
                .flatMap { file.resolveImportReference(it) }
                .forEach { ImportInsertHelper.getInstance(project).importDescriptor(file, it) }

-        return codeHolder.finish { range ->
+        val replacementPerformer = when (elementToBeReplaced) {
+            is KtExpression -> ExpressionReplacementPerformer(replacement, elementToBeReplaced)
+            is KtAnnotationEntry -> AnnotationEntryReplacementPerformer(replacement, elementToBeReplaced)
+            else -> error("Unsupported element")
+        }
+
+        return replacementPerformer.doIt(postProcessing = { range ->
            val newRange = postProcessInsertedCode(range)
            commentSaver.restore(newRange)
            newRange
-        }
+        })
    }

    private fun processValueParameterUsages(
@@ -229,7 +228,13 @@ object ReplacementEngine {
        }
    }

-    private fun ConstructedExpressionHolder.wrapCodeForSafeCall(receiver: KtExpression, receiverType: KotlinType?) {
+    private fun wrapCodeForSafeCall(
+            replacement: MutableReplacementCode,
+            receiver: KtExpression,
+            receiverType: KotlinType?,
+            expressionToBeReplaced: KtExpression,
+            bindingContext: BindingContext
+    ) {
        if (replacement.statementsBefore.isEmpty()) {
            val qualified = replacement.mainExpression as? KtQualifiedExpression
            if (qualified != null) {
@@ -244,9 +249,9 @@ object ReplacementEngine {
            }
        }

-        if (elementToBeReplaced.isUsedAsExpression(bindingContext)) {
+        if (expressionToBeReplaced.isUsedAsExpression(bindingContext)) {
            val thisReplaced = replacement.collectDescendantsOfType<KtExpression> { it[RECEIVER_VALUE_KEY] }
-            introduceValue(receiver, receiverType, thisReplaced, safeCall = true)
+            replacement.introduceValue(receiver, receiverType, thisReplaced, expressionToBeReplaced, safeCall = true)
        }
        else {
            val ifExpression = KtPsiFactory(receiver).buildExpression {
@@ -0,0 +1,105 @@
+/*
+ * 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 org.jetbrains.kotlin.idea.core.replaced
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.PsiChildRange
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
+import java.util.*
+
+internal abstract class ReplacementPerformer<TElement : KtElement>(
+        protected val replacement: MutableReplacementCode,
+        protected var elementToBeReplaced: TElement
+) {
+    protected val psiFactory = KtPsiFactory(elementToBeReplaced)
+
+    abstract fun doIt(postProcessing: (PsiChildRange) -> PsiChildRange): TElement
+}
+
+internal class AnnotationEntryReplacementPerformer(
+        replacement: MutableReplacementCode,
+        elementToBeReplaced: KtAnnotationEntry
+) : ReplacementPerformer<KtAnnotationEntry>(replacement, elementToBeReplaced) {
+
+    override fun doIt(postProcessing: (PsiChildRange) -> PsiChildRange): KtAnnotationEntry {
+        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)
+
+        assert(range.first == range.last)
+        assert(range.first is KtAnnotationEntry)
+        val annotationEntry = range.first as KtAnnotationEntry
+        val text = annotationEntry.valueArguments.single().getArgumentExpression()!!.text
+        return annotationEntry.replaced(psiFactory.createAnnotationEntry("@" + text))
+    }
+}
+
+internal class ExpressionReplacementPerformer(
+        replacement: MutableReplacementCode,
+        expressionToBeReplaced: KtExpression
+) : ReplacementPerformer<KtExpression>(replacement, expressionToBeReplaced) {
+
+    override fun doIt(postProcessing: (PsiChildRange) -> PsiChildRange): KtExpression {
+        val insertedStatements = ArrayList<KtExpression>()
+        val toInsertedStatementMap = HashMap<KtExpression, KtExpression>()
+        for (statement in replacement.statementsBefore) {
+            val anchor = findOrCreateBlockToInsertStatement()
+            val block = anchor.parent as KtBlockExpression
+
+            val inserted = block.addBefore(statement, anchor) as KtExpression
+            block.addBefore(psiFactory.createNewLine(), anchor)
+            toInsertedStatementMap.put(statement, inserted)
+            insertedStatements.add(inserted)
+        }
+
+        val replaced = elementToBeReplaced.replace(replacement.mainExpression!!) //TODO: support null here
+
+        for ((statement, actions) in replacement.postInsertionActions) {
+            val inserted = toInsertedStatementMap[statement]!!
+            actions.forEach { it(inserted) }
+        }
+
+        var range = if (insertedStatements.isEmpty()) {
+            PsiChildRange.singleElement(replaced)
+        }
+        else {
+            val statement = insertedStatements.first()
+            PsiChildRange(statement, replaced.parentsWithSelf.first { it.parent == statement.parent })
+        }
+
+        range = postProcessing(range)
+
+        return range.last as KtExpression //TODO: return value not correct!
+    }
+
+    /**
+     * Returns statement in a block to insert statement before it
+     */
+    private fun findOrCreateBlockToInsertStatement(): KtExpression {
+        //TODO: 1. There can be no block above at all
+        //TODO: 2. Create block for control statements without block
+        //TODO: 3. Sometimes it's not correct because of side effects
+        return elementToBeReplaced.parentsWithSelf
+                .filterIsInstance<KtExpression>()
+                .firstOrNull { it.parent is KtBlockExpression }!!
+    }
+}
@@ -0,0 +1,126 @@
+/*
+ * 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 org.jetbrains.kotlin.idea.analysis.computeTypeInContext
+import org.jetbrains.kotlin.idea.caches.resolve.analyze
+import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
+import org.jetbrains.kotlin.idea.core.KotlinNameSuggester
+import org.jetbrains.kotlin.idea.core.setType
+import org.jetbrains.kotlin.idea.util.getResolutionScope
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.collectDescendantsOfType
+import org.jetbrains.kotlin.psi.psiUtil.getReceiverExpression
+import org.jetbrains.kotlin.psi.psiUtil.isAncestor
+import org.jetbrains.kotlin.renderer.render
+import org.jetbrains.kotlin.resolve.bindingContextUtil.getDataFlowInfo
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
+import org.jetbrains.kotlin.resolve.scopes.utils.findLocalVariable
+import org.jetbrains.kotlin.types.ErrorUtils
+import org.jetbrains.kotlin.types.KotlinType
+
+internal fun MutableReplacementCode.introduceValue(
+        value: KtExpression,
+        valueType: KotlinType?,
+        usages: Collection<KtExpression>,
+        expressionToBeReplaced: KtExpression,
+        nameSuggestion: String? = null,
+        safeCall: Boolean = false
+) {
+    assert(usages.all { this.containsStrictlyInside(it) })
+
+    val psiFactory = KtPsiFactory(value)
+
+    val bindingContext = expressionToBeReplaced.analyze(BodyResolveMode.FULL)
+
+    fun replaceUsages(name: Name) {
+        val nameInCode = psiFactory.createExpression(name.render())
+        for (usage in usages) {
+            usage.replace(nameInCode)
+        }
+    }
+
+    fun suggestName(validator: (String) -> Boolean): Name {
+        val name = if (nameSuggestion != null)
+            KotlinNameSuggester.suggestNameByName(nameSuggestion, validator)
+        else
+            KotlinNameSuggester.suggestNamesByExpressionOnly(value, bindingContext, validator, "t").first()
+        return Name.identifier(name)
+    }
+
+    // checks that name is used (without receiver) inside code being constructed but not inside usages that will be replaced
+    fun isNameUsed(name: String) = collectNameUsages(this, name).any { nameUsage -> usages.none { it.isAncestor(nameUsage) } }
+
+    if (!safeCall) {
+        val block = expressionToBeReplaced.parent as? KtBlockExpression
+        if (block != null) {
+            val resolutionScope = expressionToBeReplaced.getResolutionScope(bindingContext, expressionToBeReplaced.getResolutionFacade())
+
+            if (usages.isNotEmpty()) {
+                var explicitType: KotlinType? = null
+                if (valueType != null && !ErrorUtils.containsErrorType(valueType)) {
+                    val valueTypeWithoutExpectedType = value.computeTypeInContext(
+                            resolutionScope,
+                            expressionToBeReplaced,
+                            dataFlowInfo = bindingContext.getDataFlowInfo(expressionToBeReplaced)
+                    )
+                    if (valueTypeWithoutExpectedType == null || ErrorUtils.containsErrorType(valueTypeWithoutExpectedType)) {
+                        explicitType = valueType
+                    }
+                }
+
+                val name = suggestName { name ->
+                    resolutionScope.findLocalVariable(Name.identifier(name)) == null && !isNameUsed(name)
+                }
+
+                val declaration = psiFactory.createDeclarationByPattern<KtVariableDeclaration>("val $0 = $1", name, value)
+                statementsBefore.add(0, declaration)
+                if (explicitType != null) {
+                    addPostInsertionAction(declaration) { it.setType(explicitType!!); it }
+                }
+
+                replaceUsages(name)
+            }
+            else {
+                statementsBefore.add(0, value)
+            }
+            return
+        }
+    }
+
+    //TODO: handle mainExpression == null and statementsBefore!
+
+    val dot = if (safeCall) "?." else "."
+
+    mainExpression = if (!isNameUsed("it")) {
+        replaceUsages(Name.identifier("it"))
+        psiFactory.createExpressionByPattern("$0${dot}let { $1 }", value, mainExpression!!)
+    }
+    else {
+        val name = suggestName { !isNameUsed(it) }
+        replaceUsages(name)
+        psiFactory.createExpressionByPattern("$0${dot}let { $1 -> $2 }", value, name, mainExpression!!)
+    }
+}
+
+private fun collectNameUsages(scope: MutableReplacementCode, name: String): List<KtSimpleNameExpression> {
+    return scope.expressions.flatMap { expression ->
+        expression.collectDescendantsOfType<KtSimpleNameExpression> { it.getReceiverExpression() == null && it.getReferencedName() == name }
+    }
+}
+
@@ -33,16 +33,19 @@ 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.KtBlockExpression
 import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.types.TypeUtils

 class KotlinInlineFunctionHandler: InlineActionHandler() {
    override fun isEnabledForLanguage(language: Language) = language == KotlinLanguage.INSTANCE

+    //TODO: overrides etc
    override fun canInlineElement(element: PsiElement): Boolean {
        return element is KtNamedFunction
-               && element.hasBody() && !element.hasBlockBody() // TODO support multiline functions
+               && element.hasBody()
               && element.getUseScope() is GlobalSearchScope // TODO support local functions
    }

@@ -50,21 +53,42 @@ class KotlinInlineFunctionHandler: InlineActionHandler() {
        element as KtNamedFunction

        val descriptor = element.resolveToDescriptor() as SimpleFunctionDescriptor
-
        val bodyExpression = element.bodyExpression!!
-        val expectedType = if (element.hasDeclaredReturnType())
+        val bodyCopy = bodyExpression.copied()
+
+        val expectedType = if (!element.hasBlockBody() && element.hasDeclaredReturnType())
            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)
+        fun analyzeBodyCopy(): BindingContext {
+            return bodyCopy.analyzeInContext(bodyExpression.getResolutionScope(),
+                                             contextExpression = bodyExpression,
+                                             expectedType = expectedType)
        }

-        val replacement = ReplacementBuilder(descriptor, element.getResolutionFacade())
-                .buildReplacementCode(expression, emptyList(), ::analyzeExpression)
+        val replacementBuilder = ReplacementBuilder(descriptor, element.getResolutionFacade())
+        val replacement = if (element.hasBlockBody()) {
+            bodyCopy as KtBlockExpression
+            val statements = bodyCopy.statements
+            if (statements.isEmpty()) {
+                //TODO
+                throw UnsupportedOperationException()
+            }
+            else {
+                //TODO: check no other return's
+                val lastReturn = statements.last() as? KtReturnExpression
+                if (lastReturn == null) {
+                    //TODO
+                    throw UnsupportedOperationException()
+                }
+
+                replacementBuilder.buildReplacementCode(lastReturn.returnedExpression, statements.dropLast(1), ::analyzeBodyCopy)
+            }
+        }
+        else {
+            replacementBuilder.buildReplacementCode(bodyCopy, emptyList(), ::analyzeBodyCopy)
+        }

        val commandName = RefactoringBundle.message("inline.command", element.name)
        CallableUsageReplacementStrategy(replacement)
@@ -0,0 +1,9 @@
+fun <caret>f(p1: Int, p2: Int): Int {
+    println(p1)
+    println(p2)
+    return p1 + p2
+}
+
+fun main(args: Array<String>) {
+    println(f(3, 5))
+}
@@ -0,0 +1,5 @@
+fun main(args: Array<String>) {
+    println(3)
+    println(5)
+    println(3 + 5)
+}
@@ -0,0 +1,9 @@
+fun <caret>f(p1: Int, p2: Int): Int {
+    println(p1)
+    println(p2)
+    return p1 + p2
+}
+
+fun main(args: Array<String>) {
+    f(3, 5)
+}
@@ -0,0 +1,5 @@
+fun main(args: Array<String>) {
+    println(3)
+    println(5)
+    3 + 5
+}
@@ -0,0 +1,7 @@
+fun <caret>f(p1: Int, p2: Int): Int {
+    return p1 + p2
+}
+
+fun main(args: Array<String>) {
+    println(f(3, 5))
+}
@@ -0,0 +1,3 @@
+fun main(args: Array<String>) {
+    println(3 + 5)
+}
@@ -0,0 +1,9 @@
+fun <caret>f(p1: Int, p2: Int): Int {
+    println(p1)
+    println(p2)
+    return p1 + p2
+}
+
+fun main(args: Array<String>) {
+    val v = f(3, 5)
+}
@@ -0,0 +1,5 @@
+fun main(args: Array<String>) {
+    println(3)
+    println(5)
+    val v = 3 + 5
+}
@@ -58,7 +58,6 @@ abstract class AbstractInlineTest : KotlinLightCodeInsightFixtureTestCase() {
            try {
                runWriteAction { handler.inlineElement(myFixture.project, myFixture.editor, targetElement) }

-                TestCase.assertTrue(afterFileExists)
                UsefulTestCase.assertEmpty(expectedErrors)
                KotlinTestUtils.assertEqualsToFile(afterFile, file.text)
                for ((extraPsiFile, extraFile) in extraFilesToPsi) {
@@ -93,6 +93,39 @@ public class InlineTestGenerated extends AbstractInlineTest {
                doTest(fileName);
            }
        }
+
+        @TestMetadata("idea/testData/refactoring/inline/function/returnAtEnd")
+        @TestDataPath("$PROJECT_ROOT")
+        @RunWith(JUnit3RunnerWithInners.class)
+        public static class ReturnAtEnd extends AbstractInlineTest {
+            public void testAllFilesPresentInReturnAtEnd() throws Exception {
+                KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/refactoring/inline/function/returnAtEnd"), Pattern.compile("^(\\w+)\\.kt$"), true);
+            }
+
+            @TestMetadata("CallArgument.kt")
+            public void testCallArgument() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/inline/function/returnAtEnd/CallArgument.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("MultipleStatements.kt")
+            public void testMultipleStatements() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/inline/function/returnAtEnd/MultipleStatements.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("SingleStatement.kt")
+            public void testSingleStatement() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/inline/function/returnAtEnd/SingleStatement.kt");
+                doTest(fileName);
+            }
+
+            @TestMetadata("ValIntializer.kt")
+            public void testValIntializer() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/inline/function/returnAtEnd/ValIntializer.kt");
+                doTest(fileName);
+            }
+        }
    }

    @TestMetadata("idea/testData/refactoring/inline/inlineTypeAlias")