Inline refactoring: should remove redundant Unit

#KT-19443 Fixed

idea/src/org/jetbrains/kotlin/idea/codeInliner/CodeInliner.kt

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.caches.resolve.resolveImportReference
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToCall
 import org.jetbrains.kotlin.idea.core.*
+import org.jetbrains.kotlin.idea.inspections.RedundantUnitExpressionInspection
 import org.jetbrains.kotlin.idea.intentions.RemoveExplicitTypeArgumentsIntention
 import org.jetbrains.kotlin.idea.util.CommentSaver
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
@@ -407,27 +408,29 @@ class CodeInliner<TCallElement : KtElement>(
     }
 
     private fun postProcessInsertedCode(range: PsiChildRange, lexicalScope: LexicalScope?): PsiChildRange {
-        val elements = range.filterIsInstance<KtElement>().toList()
+        val pointers = range.filterIsInstance<KtElement>().map { it.createSmartPointer() }.toList()
-        if (elements.isEmpty()) return PsiChildRange.EMPTY
+        if (pointers.isEmpty()) return PsiChildRange.EMPTY
 
         lexicalScope?.let {
-            renameDuplicates(elements.dropLast(1).filterIsInstance<KtNamedDeclaration>(), it)
+            renameDuplicates(pointers.dropLast(1).mapNotNull { pointer -> pointer.element as? KtNamedDeclaration }, it)
         }
 
-        elements.forEach {
+        for (pointer in pointers) {
-            introduceNamedArguments(it)
+            val element = pointer.element ?: continue
+            introduceNamedArguments(element)
 
-            restoreFunctionLiteralArguments(it)
+            restoreFunctionLiteralArguments(element)
 
             //TODO: do this earlier
-            dropArgumentsForDefaultValues(it)
+            dropArgumentsForDefaultValues(element)
 
-            simplifySpreadArrayOfArguments(it)
+            simplifySpreadArrayOfArguments(element)
 
-            removeExplicitTypeArguments(it)
+            removeExplicitTypeArguments(element)
+
+            removeRedundantUnitExpressions(element)
         }
 
-
         val shortenFilter = { element: PsiElement ->
             if (element[USER_CODE_KEY]) {
                 ShortenReferences.FilterResult.SKIP
@@ -440,11 +443,13 @@ class CodeInliner<TCallElement : KtElement>(
             }
         }
 
-        val newElements = elements.map {
+        val newElements = pointers.mapNotNull {
-            ShortenReferences { ShortenReferences.Options(removeThis = true) }.process(it, elementFilter = shortenFilter)
+            it.element?.let { element ->
+                ShortenReferences { ShortenReferences.Options(removeThis = true) }.process(element, elementFilter = shortenFilter)
+            }
         }
 
-        newElements.forEach { element ->
+        for (element in newElements) {
             // clean up user data
             element.forEachDescendantOfType<KtExpression> {
                 it.clear(USER_CODE_KEY)
@@ -460,7 +465,15 @@ class CodeInliner<TCallElement : KtElement>(
             }
         }
 
-        return PsiChildRange(newElements.first(), newElements.last())
+        return if (newElements.isEmpty()) PsiChildRange.EMPTY else PsiChildRange(newElements.first(), newElements.last())
+    }
+
+    private fun removeRedundantUnitExpressions(result: KtElement) {
+        result.forEachDescendantOfType<KtReferenceExpression> {
+            if (RedundantUnitExpressionInspection.isRedundantUnit(it)) {
+                it.delete()
+            }
+        }
     }
 
     private fun introduceNamedArguments(result: KtElement) {

idea/testData/refactoring/inline/function/EmptyFunction.kt.after

@@ -2,9 +2,9 @@ fun <T> doIt(p: () -> T): T = TODO()
 
 fun g(p: String?) {
 
-    p?.let { Unit }
+    p?.let { }
 }
 
 fun h() = Unit
 
-fun x() = doIt { Unit }
+fun x() = doIt { }

idea/testData/refactoring/inline/function/NullableUnitReturnType.kt

@@ -0,0 +1,16 @@
+fun <T> doIt(p: () -> T): T = p()
+fun Any.doDo() = Unit
+
+abstract class A {
+    abstract fun a()
+}
+
+fun foo<caret>() {
+    null?.doDo()
+}
+
+class B : A() {
+    override fun a() = doIt {
+        foo()
+    }
+}

idea/testData/refactoring/inline/function/NullableUnitReturnType.kt.after

@@ -0,0 +1,13 @@
+fun <T> doIt(p: () -> T): T = p()
+fun Any.doDo() = Unit
+
+abstract class A {
+    abstract fun a()
+}
+
+class B : A() {
+    override fun a() = doIt {
+        null?.doDo()
+        Unit
+    }
+}

idea/testData/refactoring/inline/function/NullableUnitReturnType2.kt

@@ -0,0 +1,16 @@
+fun <T> doIt(p: () -> T): T = p()
+fun Any.doDo() = Unit
+
+abstract class A {
+    abstract fun a()
+}
+
+fun foo<caret>() {
+    1.let { it.let { it.let { it.let { null?.doDo() } } } }
+}
+
+class B : A() {
+    override fun a() = doIt {
+        foo()
+    }
+}

idea/testData/refactoring/inline/function/NullableUnitReturnType2.kt.after

@@ -0,0 +1,13 @@
+fun <T> doIt(p: () -> T): T = p()
+fun Any.doDo() = Unit
+
+abstract class A {
+    abstract fun a()
+}
+
+class B : A() {
+    override fun a() = doIt {
+        1.let { it.let { it.let { it.let { null?.doDo() } } } }
+        Unit
+    }
+}

idea/testData/refactoring/inline/function/NullableUnitReturnType3.kt

@@ -0,0 +1,14 @@
+fun <T> doIt(p: () -> T): T = p()
+fun Any.doDo() = Unit
+
+abstract class A {
+    abstract fun a()
+}
+
+fun foo<caret>(): Unit = doIt {
+    1.let { it.let { it.let { it.let { null?.doDo() } } } }
+}
+
+class B : A() {
+    override fun a() = foo()
+}

idea/testData/refactoring/inline/function/NullableUnitReturnType3.kt.after

@@ -0,0 +1,12 @@
+fun <T> doIt(p: () -> T): T = p()
+fun Any.doDo() = Unit
+
+abstract class A {
+    abstract fun a()
+}
+
+class B : A() {
+    override fun a() = doIt {
+        1.let<Int, Unit> { it.let { it.let { it.let { null?.doDo() } } } }
+    }
+}

idea/testData/refactoring/inline/function/UnitReturnType.kt.after

@@ -9,7 +9,6 @@ class C {
         p?.let {
             println(3)
             println(4)
-            Unit
         }
 
         if (other != null) {
@@ -26,6 +25,5 @@ class C {
     fun x() = doIt {
         println(9)
         println(10)
-        Unit
     }
 }

idea/testData/refactoring/inline/function/UnitReturnType3.kt

@@ -0,0 +1,12 @@
+class C {
+    fun <caret>f(p1: Int, p2: Int) {
+        println(p1)
+        nonUnit(p2)
+    }
+
+    fun nonUnit(p: Int): Int = p
+
+    fun <T> doIt(p: () -> T): T = TODO()
+
+    fun x() = doIt<Unit> { f(9, 10) }
+}

idea/testData/refactoring/inline/function/UnitReturnType3.kt.after

@@ -0,0 +1,11 @@
+class C {
+
+    fun nonUnit(p: Int): Int = p
+
+    fun <T> doIt(p: () -> T): T = TODO()
+
+    fun x() = doIt<Unit> {
+        println(9)
+        nonUnit(10)
+    }
+}

idea/testData/refactoring/inline/function/UnitReturnType4.kt

@@ -0,0 +1,16 @@
+fun <T> doIt(p: () -> T): T = p()
+fun Any.doDo() = Unit
+
+abstract class A {
+    abstract fun a()
+}
+
+fun foo<caret>() {
+    1.let { it.let { it.let { it.let { it?.doDo() } } } }
+}
+
+class B : A() {
+    override fun a() = doIt {
+        foo()
+    }
+}

idea/testData/refactoring/inline/function/UnitReturnType4.kt.after

@@ -0,0 +1,12 @@
+fun <T> doIt(p: () -> T): T = p()
+fun Any.doDo() = Unit
+
+abstract class A {
+    abstract fun a()
+}
+
+class B : A() {
+    override fun a() = doIt {
+        1.let { it.let { it.let { it.let { it?.doDo() } } } }
+    }
+}

idea/tests/org/jetbrains/kotlin/idea/refactoring/inline/InlineTestGenerated.java

@@ -75,6 +75,21 @@ public class InlineTestGenerated extends AbstractInlineTest {
             runTest("idea/testData/refactoring/inline/function/MultipleReturns.kt");
         }
 
+        @TestMetadata("NullableUnitReturnType.kt")
+        public void testNullableUnitReturnType() throws Exception {
+            runTest("idea/testData/refactoring/inline/function/NullableUnitReturnType.kt");
+        }
+
+        @TestMetadata("NullableUnitReturnType2.kt")
+        public void testNullableUnitReturnType2() throws Exception {
+            runTest("idea/testData/refactoring/inline/function/NullableUnitReturnType2.kt");
+        }
+
+        @TestMetadata("NullableUnitReturnType3.kt")
+        public void testNullableUnitReturnType3() throws Exception {
+            runTest("idea/testData/refactoring/inline/function/NullableUnitReturnType3.kt");
+        }
+
         @TestMetadata("OuterClassReceiver.kt")
         public void testOuterClassReceiver() throws Exception {
             runTest("idea/testData/refactoring/inline/function/OuterClassReceiver.kt");
@@ -135,6 +150,16 @@ public class InlineTestGenerated extends AbstractInlineTest {
             runTest("idea/testData/refactoring/inline/function/UnitReturnType2.kt");
         }
 
+        @TestMetadata("UnitReturnType3.kt")
+        public void testUnitReturnType3() throws Exception {
+            runTest("idea/testData/refactoring/inline/function/UnitReturnType3.kt");
+        }
+
+        @TestMetadata("UnitReturnType4.kt")
+        public void testUnitReturnType4() throws Exception {
+            runTest("idea/testData/refactoring/inline/function/UnitReturnType4.kt");
+        }
+
         @TestMetadata("idea/testData/refactoring/inline/function/expressionBody")
         @TestDataPath("$PROJECT_ROOT")
         @RunWith(JUnit3RunnerWithInners.class)