Extract Function: Support extraction of "it" parameter from lambdas

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/ExtractionData.kt

@@ -46,6 +46,11 @@ import org.jetbrains.jet.lang.psi.JetDeclaration
 import org.jetbrains.jet.lang.psi.JetDeclarationWithBody
 import org.jetbrains.jet.lang.psi.JetUserType
 import org.jetbrains.jet.lang.resolve.calls.model.VariableAsFunctionResolvedCall
+import org.jetbrains.jet.lang.psi.JetParameter
+import org.jetbrains.jet.lang.descriptors.ValueParameterDescriptor
+import org.jetbrains.jet.lang.psi.JetPsiFactory
+import org.jetbrains.jet.lang.resolve.BindingContextUtils
+import org.jetbrains.jet.lang.psi.JetFunctionLiteral
 
 data class ExtractionOptions(val inferUnitTypeForUnusedValues: Boolean) {
     class object {
@@ -96,9 +101,18 @@ class ExtractionData(
 
     val originalStartOffset = originalElements.first?.let { e -> e.getTextRange()!!.getStartOffset() }
 
+    private val itFakeDeclaration by Delegates.lazy { JetPsiFactory.createParameter(project, "it", null) }
+
     val refOffsetToDeclaration by Delegates.lazy {
+        fun isExtractableIt(descriptor: DeclarationDescriptor, context: BindingContext): Boolean {
+            if (!(descriptor is ValueParameterDescriptor && (context[BindingContext.AUTO_CREATED_IT, descriptor] ?: false))) return false
+            val function = BindingContextUtils.descriptorToDeclaration(context, descriptor.getContainingDeclaration()) as? JetFunctionLiteral
+            return function == null || !function.isInsideOf(originalElements)
+        }
+
         if (originalStartOffset != null) {
             val resultMap = HashMap<Int, ResolveResult>()
+
             for ((ref, context) in JetFileReferencesResolver.resolve(originalFile, getExpressions())) {
                 if (ref !is JetSimpleNameExpression) continue
 
@@ -111,7 +125,7 @@ class ExtractionData(
                 if (descriptor == null) continue
 
                 val declaration = DescriptorToDeclarationUtil.getDeclaration(project, descriptor, context) as? PsiNamedElement
-                if (declaration == null) continue
+                        ?: if (isExtractableIt(descriptor, context)) itFakeDeclaration else continue
 
                 val offset = ref.getTextRange()!!.getStartOffset() - originalStartOffset
                 resultMap[offset] = ResolveResult(ref, declaration, descriptor, resolvedCall)

idea/src/org/jetbrains/jet/plugin/refactoring/extractFunction/extractFunctionUtils.kt

@@ -133,8 +133,8 @@ private fun List<Instruction>.analyzeControlFlow(
 ): Pair<ControlFlow, ErrorMessage?> {
     fun isCurrentFunctionReturn(expression: JetReturnExpression): Boolean {
         val functionDescriptor = expression.getTargetFunctionDescriptor(bindingContext)
-        val function = functionDescriptor?.let { BindingContextUtils.descriptorToDeclaration(bindingContext, it) }
+        val currentDescriptor = bindingContext[BindingContext.DECLARATION_TO_DESCRIPTOR, pseudocode.getCorrespondingElement()]
-        return function == pseudocode.getCorrespondingElement()
+        return currentDescriptor == functionDescriptor
     }
 
     val exitPoints = getExitPoints()
@@ -461,13 +461,14 @@ private fun ExtractionData.inferParametersInfo(commonParent: PsiElement,
 }
 
 private fun ExtractionData.checkLocalDeclarationsWithNonLocalUsages(
+        pseudocode: Pseudocode,
         localInstructions: List<Instruction>,
         bindingContext: BindingContext
 ): ErrorMessage? {
     // todo: non-locally used declaration can be turned into the output value
 
     val declarations = ArrayList<JetNamedDeclaration>()
-    localInstructions.firstOrNull()?.owner?.traverse(TraversalOrder.FORWARD) { instruction ->
+    pseudocode.traverse(TraversalOrder.FORWARD) { instruction ->
         if (instruction !in localInstructions) {
             PseudocodeUtil.extractVariableDescriptorIfAny(instruction, true, bindingContext)?.let { descriptor ->
                 val declaration = DescriptorToDeclarationUtil.getDeclaration(project, descriptor, bindingContext)
@@ -562,7 +563,7 @@ fun ExtractionData.performAnalysis(): AnalysisResult {
         )
     }
 
-    checkLocalDeclarationsWithNonLocalUsages(localInstructions, bindingContext)?.let { messages.add(it) }
+    checkLocalDeclarationsWithNonLocalUsages(pseudocode, localInstructions, bindingContext)?.let { messages.add(it) }
     checkDeclarationsMovingOutOfScope(controlFlow)?.let { messages.add(it) }
 
     val functionNameValidator =

idea/testData/refactoring/extractFunction/parameters/it/innerIt.kt

@@ -0,0 +1,9 @@
+// PARAM_TYPES: kotlin.Int
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Array<Int>>) {
+    if (t.check { it.check{ <selection>it + 1</selection> > 1 } }) {
+        println("OK")
+    }
+}

idea/testData/refactoring/extractFunction/parameters/it/innerIt.kt.after

@@ -0,0 +1,13 @@
+// PARAM_TYPES: kotlin.Int
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Array<Int>>) {
+    if (t.check { it.check{ i(it) > 1 } }) {
+        println("OK")
+    }
+}
+
+fun i(it: Int): Int {
+    return it + 1
+}

idea/testData/refactoring/extractFunction/parameters/it/lambdaWithIt.kt

@@ -0,0 +1,9 @@
+// PARAM_TYPES: kotlin.Array<kotlin.Int>
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Int>) {
+    if (<selection>t.check { it + 1 > 1 }</selection>) {
+        println("OK")
+    }
+}

idea/testData/refactoring/extractFunction/parameters/it/lambdaWithIt.kt.after

@@ -0,0 +1,13 @@
+// PARAM_TYPES: kotlin.Array<kotlin.Int>
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Int>) {
+    if (b(t)) {
+        println("OK")
+    }
+}
+
+fun b(t: Array<Int>): Boolean {
+    return t.check { it + 1 > 1 }
+}

idea/testData/refactoring/extractFunction/parameters/it/outerIt.kt

@@ -0,0 +1,9 @@
+// PARAM_TYPES: kotlin.Array<kotlin.Int>
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Array<Int>>) {
+    if (t.check { <selection>it.check{ it + 1 > 1 }</selection> }) {
+        println("OK")
+    }
+}

idea/testData/refactoring/extractFunction/parameters/it/outerIt.kt.after

@@ -0,0 +1,13 @@
+// PARAM_TYPES: kotlin.Array<kotlin.Int>
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Array<Int>>) {
+    if (t.check { b(it) }) {
+        println("OK")
+    }
+}
+
+fun b(it: Array<Int>): Boolean {
+    return it.check { it + 1 > 1 }
+}

idea/testData/refactoring/extractFunction/parameters/it/simpleIt.kt

@@ -0,0 +1,9 @@
+// PARAM_TYPES: kotlin.Int
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Int>) {
+    if (t.check { <selection>it + 1</selection> > 1 }) {
+        println("OK")
+    }
+}

idea/testData/refactoring/extractFunction/parameters/it/simpleIt.kt.after

@@ -0,0 +1,13 @@
+// PARAM_TYPES: kotlin.Int
+fun <T> Array<T>.check(f: (T) -> Boolean): Boolean = false
+
+// SIBLING:
+fun foo(t: Array<Int>) {
+    if (t.check { i(it) > 1 }) {
+        println("OK")
+    }
+}
+
+fun i(it: Int): Int {
+    return it + 1
+}

idea/tests/org/jetbrains/jet/plugin/refactoring/introduce/introduceVariable/JetExtractionTestGenerated.java

@@ -651,7 +651,7 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
         }
         
         @TestMetadata("idea/testData/refactoring/extractFunction/parameters")
-        @InnerTestClasses({Parameters.ExtractSuper.class, Parameters.ExtractThis.class, Parameters.Misc.class, Parameters.NonDenotableTypes.class})
+        @InnerTestClasses({Parameters.ExtractSuper.class, Parameters.ExtractThis.class, Parameters.It.class, Parameters.Misc.class, Parameters.NonDenotableTypes.class})
         public static class Parameters extends AbstractJetExtractionTest {
             public void testAllFilesPresentInParameters() throws Exception {
                 JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.TestsPackage", new File("idea/testData/refactoring/extractFunction/parameters"), Pattern.compile("^(.+)\\.kt$"), true);
@@ -728,6 +728,34 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                 
             }
             
+            @TestMetadata("idea/testData/refactoring/extractFunction/parameters/it")
+            public static class It extends AbstractJetExtractionTest {
+                public void testAllFilesPresentInIt() throws Exception {
+                    JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.TestsPackage", new File("idea/testData/refactoring/extractFunction/parameters/it"), Pattern.compile("^(.+)\\.kt$"), true);
+                }
+                
+                @TestMetadata("innerIt.kt")
+                public void testInnerIt() throws Exception {
+                    doExtractFunctionTest("idea/testData/refactoring/extractFunction/parameters/it/innerIt.kt");
+                }
+                
+                @TestMetadata("lambdaWithIt.kt")
+                public void testLambdaWithIt() throws Exception {
+                    doExtractFunctionTest("idea/testData/refactoring/extractFunction/parameters/it/lambdaWithIt.kt");
+                }
+                
+                @TestMetadata("outerIt.kt")
+                public void testOuterIt() throws Exception {
+                    doExtractFunctionTest("idea/testData/refactoring/extractFunction/parameters/it/outerIt.kt");
+                }
+                
+                @TestMetadata("simpleIt.kt")
+                public void testSimpleIt() throws Exception {
+                    doExtractFunctionTest("idea/testData/refactoring/extractFunction/parameters/it/simpleIt.kt");
+                }
+                
+            }
+            
             @TestMetadata("idea/testData/refactoring/extractFunction/parameters/misc")
             public static class Misc extends AbstractJetExtractionTest {
                 public void testAllFilesPresentInMisc() throws Exception {
@@ -839,6 +867,7 @@ public class JetExtractionTestGenerated extends AbstractJetExtractionTest {
                 suite.addTestSuite(Parameters.class);
                 suite.addTestSuite(ExtractSuper.class);
                 suite.addTestSuite(ExtractThis.class);
+                suite.addTestSuite(It.class);
                 suite.addTestSuite(Misc.class);
                 suite.addTestSuite(NonDenotableTypes.class);
                 return suite;