Navigation to iterator(), next(), hasNext() through 'in' in for-loops

compiler/frontend/src/org/jetbrains/jet/lang/psi/JetForExpression.java

@@ -20,6 +20,7 @@ import com.intellij.lang.ASTNode;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.jet.JetNodeTypes;
+import org.jetbrains.jet.lexer.JetTokens;
 
 public class JetForExpression extends JetLoopExpression {
     public JetForExpression(@NotNull ASTNode node) {
@@ -50,4 +51,9 @@ public class JetForExpression extends JetLoopExpression {
     public JetExpression getLoopRange() {
         return findExpressionUnder(JetNodeTypes.LOOP_RANGE);
     }
+
+    @Nullable @IfNotParsed
+    public ASTNode getInKeywordNode() {
+        return getNode().findChildByType(JetTokens.IN_KEYWORD);
+    }
 }

idea/src/org/jetbrains/jet/plugin/references/JetForLoopInReference.java

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2010-2013 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.jet.plugin.references;
+
+import com.beust.jcommander.internal.Lists;
+import com.intellij.lang.ASTNode;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiReference;
+import com.intellij.psi.ResolveResult;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.jet.lang.descriptors.FunctionDescriptor;
+import org.jetbrains.jet.lang.psi.JetExpression;
+import org.jetbrains.jet.lang.psi.JetForExpression;
+import org.jetbrains.jet.lang.resolve.BindingContext;
+import org.jetbrains.jet.lang.resolve.calls.model.ResolvedCall;
+import org.jetbrains.jet.plugin.project.AnalyzerFacadeWithCache;
+import org.jetbrains.jet.util.slicedmap.WritableSlice;
+
+import java.util.List;
+
+public class JetForLoopInReference extends AbstractPolyVariantJetReference<JetForExpression> {
+
+    public static PsiReference[] create(@NotNull JetForExpression delegate) {
+        return new PsiReference[] { new JetForLoopInReference(delegate) };
+    }
+
+    private JetForLoopInReference(@NotNull JetForExpression element) {
+        super(element);
+    }
+
+    @NotNull
+    @Override
+    public ResolveResult[] multiResolve(boolean incompleteCode) {
+        BindingContext context = AnalyzerFacadeWithCache.getContextForElement(element);
+
+        JetExpression loopRange = element.getLoopRange();
+        if (loopRange == null) return ResolveResult.EMPTY_ARRAY;
+
+        List<ResolveResult> results = Lists.newArrayList();
+
+        addResultsForKey(context, results, loopRange, BindingContext.LOOP_RANGE_ITERATOR_RESOLVED_CALL);
+        addResultsForKey(context, results, loopRange, BindingContext.LOOP_RANGE_NEXT_RESOLVED_CALL);
+        addResultsForKey(context, results, loopRange, BindingContext.LOOP_RANGE_HAS_NEXT_RESOLVED_CALL);
+
+        return results.toArray(new ResolveResult[results.size()]);
+    }
+
+    private void addResultsForKey(
+            BindingContext context,
+            List<ResolveResult> results,
+            JetExpression loopRange,
+            WritableSlice<JetExpression, ResolvedCall<FunctionDescriptor>> key
+    ) {
+        ResolvedCall<FunctionDescriptor> resolvedCall = context.get(key, loopRange);
+        if (resolvedCall == null) return;
+
+        JetReferenceUtil.findPsiElements(element.getProject(), context, results, resolvedCall.getResultingDescriptor());
+    }
+
+    @Override
+    public TextRange getRangeInElement() {
+        ASTNode inKeywordNode = element.getInKeywordNode();
+        if (inKeywordNode == null) return TextRange.EMPTY_RANGE;
+
+        int offset = inKeywordNode.getPsi().getStartOffsetInParent();
+        return new TextRange(offset, offset + inKeywordNode.getTextLength());
+    }
+
+    @NotNull
+    @Override
+    public String getCanonicalText() {
+        return "<unknown>";
+    }
+}

idea/src/org/jetbrains/jet/plugin/references/JetReferenceContributor.java

@@ -70,5 +70,14 @@ public class JetReferenceContributor extends PsiReferenceContributor {
                                                     return JetPropertyDelegationMethodsReference.create((JetPropertyDelegate) element);
                                                 }
                                             });
+
+        registrar.registerReferenceProvider(psiElement(JetForExpression.class),
+                                            new PsiReferenceProvider() {
+                                                @NotNull
+                                                @Override
+                                                public PsiReference[] getReferencesByElement(@NotNull PsiElement element, @NotNull ProcessingContext processingContext) {
+                                                    return JetForLoopInReference.create((JetForExpression) element);
+                                                }
+                                            });
     }
 }

idea/testData/resolve/references/forLoopIn/inBuiltIns/extension.kt

@@ -0,0 +1,8 @@
+fun main(it: Iterator<Any>) {
+  for (i <caret>in it.iterator()) {}
+}
+
+// MULTIRESOLVE
+// REF: (for Iterator<T> in jet).iterator()
+// REF: (in jet.Iterator).hasNext()
+// REF: (in jet.Iterator).next()

idea/testData/resolve/references/forLoopIn/inBuiltIns/member.kt

@@ -0,0 +1,8 @@
+fun main() {
+  for (i <caret>in 1..2) {}
+}
+
+// MULTIRESOLVE
+// REF: (in jet.IntIterator).next()
+// REF: (in jet.IntRange).iterator()
+// REF: (in jet.Iterator).hasNext()

idea/testData/resolve/references/forLoopIn/inLibrary/extension.kt

@@ -0,0 +1,8 @@
+fun main() {
+  for (i <caret>in "") {}
+}
+
+// MULTIRESOLVE
+// REF: (for jet.CharSequence in kotlin).iterator()
+// REF: (in jet.CharIterator).next()
+// REF: (in jet.Iterator).hasNext()

idea/testData/resolve/references/forLoopIn/inSource/allExtensions.kt

@@ -0,0 +1,16 @@
+fun test(f: Foo) {
+    for(i <caret>in f) {}
+}
+
+trait Foo
+fun Foo.iterator(): Iterator
+
+trait Iterator
+
+fun Iterator.next(): Any
+fun Iterator.hasNext(): Boolean
+
+// MULTIRESOLVE
+// REF: (for Foo in <root>).iterator()
+// REF: (for Iterator in <root>).next()
+// REF: (for Iterator in <root>).hasNext()

idea/testData/resolve/references/forLoopIn/inSource/allMembers.kt

@@ -0,0 +1,17 @@
+fun test(f: Foo) {
+    for(i <caret>in f) {}
+}
+
+trait Foo {
+  fun iterator(): Iterator
+}
+
+trait Iterator {
+  fun next(): Any
+  fun hasNext(): Boolean
+}
+
+// MULTIRESOLVE
+// REF: (in Foo).iterator()
+// REF: (in Iterator).next()
+// REF: (in Iterator).hasNext()

idea/testData/resolve/references/forLoopIn/inSource/nextMissing.kt

@@ -0,0 +1,15 @@
+fun test(f: Foo) {
+    for(i <caret>in f) {}
+}
+
+trait Foo {
+  fun iterator(): Iterator
+}
+
+trait Iterator {
+  fun hasNext(): Boolean
+}
+
+// MULTIRESOLVE
+// REF: (in Foo).iterator()
+// REF: (in Iterator).hasNext()

idea/testData/resolve/references/forLoopIn/unresolvedIterator.kt

@@ -0,0 +1,7 @@
+fun test() {
+  for (x <caret>in Y()) {}
+}
+
+class Y
+
+// REF_EMPTY

idea/tests/org/jetbrains/jet/resolve/ReferenceResolveTestGenerated.java

@@ -31,7 +31,7 @@ import org.jetbrains.jet.resolve.AbstractResolveBaseTest;
 /** This class is generated by {@link org.jetbrains.jet.generators.tests.GenerateTests}. DO NOT MODIFY MANUALLY */
 @SuppressWarnings("all")
 @TestMetadata("idea/testData/resolve/references")
-@InnerTestClasses({ReferenceResolveTestGenerated.DelegatedPropertyAccessors.class})
+@InnerTestClasses({ReferenceResolveTestGenerated.DelegatedPropertyAccessors.class, ReferenceResolveTestGenerated.ForLoopIn.class})
 public class ReferenceResolveTestGenerated extends AbstractResolveBaseTest {
     public void testAllFilesPresentInReferences() throws Exception {
         JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.GenerateTests", new File("idea/testData/resolve/references"), Pattern.compile("^(.+)\\.kt$"), true);
@@ -249,10 +249,87 @@ public class ReferenceResolveTestGenerated extends AbstractResolveBaseTest {
         }
     }
     
+    @TestMetadata("idea/testData/resolve/references/forLoopIn")
+    @InnerTestClasses({ForLoopIn.InBuiltIns.class, ForLoopIn.InLibrary.class, ForLoopIn.InSource.class})
+    public static class ForLoopIn extends AbstractResolveBaseTest {
+        public void testAllFilesPresentInForLoopIn() throws Exception {
+            JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.GenerateTests", new File("idea/testData/resolve/references/forLoopIn"), Pattern.compile("^(.+)\\.kt$"), true);
+        }
+        
+        @TestMetadata("unresolvedIterator.kt")
+        public void testUnresolvedIterator() throws Exception {
+            doTest("idea/testData/resolve/references/forLoopIn/unresolvedIterator.kt");
+        }
+        
+        @TestMetadata("idea/testData/resolve/references/forLoopIn/inBuiltIns")
+        public static class InBuiltIns extends AbstractResolveBaseTest {
+            public void testAllFilesPresentInInBuiltIns() throws Exception {
+                JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.GenerateTests", new File("idea/testData/resolve/references/forLoopIn/inBuiltIns"), Pattern.compile("^(.+)\\.kt$"), true);
+            }
+            
+            @TestMetadata("extension.kt")
+            public void testExtension() throws Exception {
+                doTest("idea/testData/resolve/references/forLoopIn/inBuiltIns/extension.kt");
+            }
+            
+            @TestMetadata("member.kt")
+            public void testMember() throws Exception {
+                doTest("idea/testData/resolve/references/forLoopIn/inBuiltIns/member.kt");
+            }
+            
+        }
+        
+        @TestMetadata("idea/testData/resolve/references/forLoopIn/inLibrary")
+        public static class InLibrary extends AbstractResolveBaseTest {
+            public void testAllFilesPresentInInLibrary() throws Exception {
+                JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.GenerateTests", new File("idea/testData/resolve/references/forLoopIn/inLibrary"), Pattern.compile("^(.+)\\.kt$"), true);
+            }
+            
+            @TestMetadata("extension.kt")
+            public void testExtension() throws Exception {
+                doTest("idea/testData/resolve/references/forLoopIn/inLibrary/extension.kt");
+            }
+            
+        }
+        
+        @TestMetadata("idea/testData/resolve/references/forLoopIn/inSource")
+        public static class InSource extends AbstractResolveBaseTest {
+            @TestMetadata("allExtensions.kt")
+            public void testAllExtensions() throws Exception {
+                doTest("idea/testData/resolve/references/forLoopIn/inSource/allExtensions.kt");
+            }
+            
+            public void testAllFilesPresentInInSource() throws Exception {
+                JetTestUtils.assertAllTestsPresentByMetadata(this.getClass(), "org.jetbrains.jet.generators.tests.GenerateTests", new File("idea/testData/resolve/references/forLoopIn/inSource"), Pattern.compile("^(.+)\\.kt$"), true);
+            }
+            
+            @TestMetadata("allMembers.kt")
+            public void testAllMembers() throws Exception {
+                doTest("idea/testData/resolve/references/forLoopIn/inSource/allMembers.kt");
+            }
+            
+            @TestMetadata("nextMissing.kt")
+            public void testNextMissing() throws Exception {
+                doTest("idea/testData/resolve/references/forLoopIn/inSource/nextMissing.kt");
+            }
+            
+        }
+        
+        public static Test innerSuite() {
+            TestSuite suite = new TestSuite("ForLoopIn");
+            suite.addTestSuite(ForLoopIn.class);
+            suite.addTestSuite(InBuiltIns.class);
+            suite.addTestSuite(InLibrary.class);
+            suite.addTestSuite(InSource.class);
+            return suite;
+        }
+    }
+    
     public static Test suite() {
         TestSuite suite = new TestSuite("ReferenceResolveTestGenerated");
         suite.addTestSuite(ReferenceResolveTestGenerated.class);
         suite.addTest(DelegatedPropertyAccessors.innerSuite());
+        suite.addTest(ForLoopIn.innerSuite());
         return suite;
     }
 }