From 5bdd52c65f49eb2c42df49b3971f7bf6b67ecf97 Mon Sep 17 00:00:00 2001
From: "Natalia.Ukhorskaya" <natalia.ukhorskaya@jetbrains.com>
Date: Thu, 14 Mar 2013 14:45:53 +0400
Subject: [PATCH] Do not remove import for iterator()

---
 .../importOptimizer/JetImportOptimizer.java   | 23 +++++++++++++++++++
 .../IteratorFunction_after.kt                 |  7 ++++++
 .../multifile/IteratorFunction/main.kt        |  7 ++++++
 .../multifile/IteratorFunction/myClass.kt     | 15 ++++++++++++
 .../OptimizeImportsMultiFileTest.java         |  4 ++++
 5 files changed, 56 insertions(+)
 create mode 100644 idea/testData/editor/optimizeImports/multifile/IteratorFunction/IteratorFunction_after.kt
 create mode 100644 idea/testData/editor/optimizeImports/multifile/IteratorFunction/main.kt
 create mode 100644 idea/testData/editor/optimizeImports/multifile/IteratorFunction/myClass.kt

diff --git a/idea/src/org/jetbrains/jet/plugin/editor/importOptimizer/JetImportOptimizer.java b/idea/src/org/jetbrains/jet/plugin/editor/importOptimizer/JetImportOptimizer.java
index ecaa14e392e..37233bdfb27 100644
--- a/idea/src/org/jetbrains/jet/plugin/editor/importOptimizer/JetImportOptimizer.java
+++ b/idea/src/org/jetbrains/jet/plugin/editor/importOptimizer/JetImportOptimizer.java
@@ -23,11 +23,19 @@ import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTreeUtil;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.jet.lang.descriptors.FunctionDescriptor;
 import org.jetbrains.jet.lang.psi.*;
+import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.resolve.DescriptorUtils;
 import org.jetbrains.jet.lang.resolve.ImportPath;
+import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCall;
 import org.jetbrains.jet.lang.resolve.java.PackageClassUtils;
+import org.jetbrains.jet.lang.resolve.lazy.ResolveSession;
+import org.jetbrains.jet.lang.resolve.lazy.ResolveSessionUtils;
 import org.jetbrains.jet.lang.resolve.name.FqName;
+import org.jetbrains.jet.lang.resolve.name.FqNameUnsafe;
 import org.jetbrains.jet.lang.resolve.name.Name;
+import org.jetbrains.jet.plugin.project.WholeProjectAnalyzerFacade;
 import org.jetbrains.jet.plugin.quickfix.ImportInsertHelper;
 import org.jetbrains.jet.util.QualifiedNamesUtil;
 
@@ -179,6 +187,21 @@ public class JetImportOptimizer implements ImportOptimizer {
 
                 super.visitReferenceExpression(expression);
             }
+
+            @Override
+            public void visitForExpression(JetForExpression expression) {
+                ResolveSession resolveSession = WholeProjectAnalyzerFacade.getLazyResolveSessionForFile((JetFile) expression.getContainingFile());
+                BindingContext context = ResolveSessionUtils.resolveToExpression(resolveSession, expression);
+                ResolvedCall<FunctionDescriptor> resolvedCall = context.get(BindingContext.LOOP_RANGE_ITERATOR_RESOLVED_CALL, expression.getLoopRange());
+                if (resolvedCall != null) {
+                    FunctionDescriptor iteratorDescriptor = resolvedCall.getResultingDescriptor();
+                    FqNameUnsafe name = DescriptorUtils.getFQName(iteratorDescriptor);
+                    assert name.isSafe();
+                    usedQualifiedNames.add(name.toSafe());
+                }
+
+                super.visitForExpression(expression);
+            }
         });
 
         return usedQualifiedNames;
diff --git a/idea/testData/editor/optimizeImports/multifile/IteratorFunction/IteratorFunction_after.kt b/idea/testData/editor/optimizeImports/multifile/IteratorFunction/IteratorFunction_after.kt
new file mode 100644
index 00000000000..c2013ae8128
--- /dev/null
+++ b/idea/testData/editor/optimizeImports/multifile/IteratorFunction/IteratorFunction_after.kt
@@ -0,0 +1,7 @@
+import test1.MyClass
+import test1.iterator
+
+fun foo() {
+    val s = MyClass()
+    for (i in s) {}
+}
diff --git a/idea/testData/editor/optimizeImports/multifile/IteratorFunction/main.kt b/idea/testData/editor/optimizeImports/multifile/IteratorFunction/main.kt
new file mode 100644
index 00000000000..c2013ae8128
--- /dev/null
+++ b/idea/testData/editor/optimizeImports/multifile/IteratorFunction/main.kt
@@ -0,0 +1,7 @@
+import test1.MyClass
+import test1.iterator
+
+fun foo() {
+    val s = MyClass()
+    for (i in s) {}
+}
diff --git a/idea/testData/editor/optimizeImports/multifile/IteratorFunction/myClass.kt b/idea/testData/editor/optimizeImports/multifile/IteratorFunction/myClass.kt
new file mode 100644
index 00000000000..d9004a82ed3
--- /dev/null
+++ b/idea/testData/editor/optimizeImports/multifile/IteratorFunction/myClass.kt
@@ -0,0 +1,15 @@
+package test1
+
+public class MyClass {
+}
+
+public fun MyClass.iterator(): Iterator<MyClass> {
+     return object: Iterator<MyClass> {
+         override fun next(): MyClass {
+             throw Exception()
+         }
+         override fun hasNext(): Boolean {
+             throw Exception()
+         }
+     }
+ }
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/jet/plugin/importOptimizer/OptimizeImportsMultiFileTest.java b/idea/tests/org/jetbrains/jet/plugin/importOptimizer/OptimizeImportsMultiFileTest.java
index bbb8d2d6d45..8f1a9090fbb 100644
--- a/idea/tests/org/jetbrains/jet/plugin/importOptimizer/OptimizeImportsMultiFileTest.java
+++ b/idea/tests/org/jetbrains/jet/plugin/importOptimizer/OptimizeImportsMultiFileTest.java
@@ -42,6 +42,10 @@ public class OptimizeImportsMultiFileTest extends CodeInsightTestCase {
         doTest(getTestName(false) + "/main.kt", getTestName(false) + "/myClass.kt");
     }
 
+    public void testIteratorFunction() throws Exception {
+        doTest(getTestName(false) + "/main.kt", getTestName(false) + "/myClass.kt");
+    }
+
     public void doTest(String... fileNames) throws Exception {
         configureByFiles(null, fileNames);
         invokeFormatFile();