Find usages correctly

8 files changed, 126 insertions, 47 deletions

diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/findUsages/SdFindUsagesHandler.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/findUsages/SdFindUsagesHandler.java
index d95c2c045fe..4af276c26df 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/findUsages/SdFindUsagesHandler.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/findUsages/SdFindUsagesHandler.java
@@ -2,7 +2,9 @@
 package ai.vespa.intellij.schema.findUsages;
 
 import ai.vespa.intellij.schema.model.Function;
+import ai.vespa.intellij.schema.model.RankProfile;
 import ai.vespa.intellij.schema.model.Schema;
+import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
 import ai.vespa.intellij.schema.utils.Path;
 import com.intellij.find.findUsages.FindUsagesHandler;
 import com.intellij.find.findUsages.FindUsagesOptions;
@@ -13,15 +15,13 @@ import com.intellij.psi.PsiReference;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.search.SearchScope;
 import com.intellij.psi.search.searches.ReferencesSearch;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.usageView.UsageInfo;
 import com.intellij.util.Processor;
-import ai.vespa.intellij.schema.SdUtil;
-import ai.vespa.intellij.schema.psi.SdFile;
 import ai.vespa.intellij.schema.psi.SdFunctionDefinition;
 
-import java.util.HashMap;
+import java.util.Collection;
 import java.util.List;
-import java.util.Map;
 
 /**
  * This class handles creating the "Find Usages" window.
@@ -30,12 +30,8 @@ import java.util.Map;
  */
 public class SdFindUsagesHandler extends FindUsagesHandler {
     
-    private final Map<String, List<Function>> functionsMap;
-    
     public SdFindUsagesHandler(PsiElement psiElement) {
         super(psiElement);
-        PsiFile file = psiElement.getContainingFile();
-        functionsMap = file instanceof SdFile ? new Schema((SdFile)file).functions() : Map.of();
     }
     
     @Override
@@ -46,23 +42,13 @@ public class SdFindUsagesHandler extends FindUsagesHandler {
         boolean searchText = options.isSearchForTextOccurrences && scope instanceof GlobalSearchScope;
         
         if (options.isUsages) {
-            if (!(elementToSearch instanceof SdFunctionDefinition)) {
+            if (elementToSearch instanceof SdFunctionDefinition) {
+                findFunctionUsages((SdFunctionDefinition) elementToSearch, processor);
+            } else {
                 boolean success =
-                    ReferencesSearch.search(createSearchParameters(elementToSearch, scope, options))
-                                    .forEach((PsiReference ref) -> processor.process(new UsageInfo(ref)));
+                        ReferencesSearch.search(createSearchParameters(elementToSearch, scope, options))
+                                        .forEach((PsiReference ref) -> processor.process(new UsageInfo(ref)));
                 if (!success) return false;
-            } else {
-                String functionName = ReadAction.compute( ((SdFunctionDefinition) elementToSearch)::getName);
-                
-                for (Function function : functionsMap.get(functionName)) {
-                    boolean success =
-                        ReferencesSearch.search(createSearchParameters(function.definition(), scope, options))
-                                        .forEach((PsiReference ref) -> {
-                                            if (ref.getElement().getParent() == elementToSearch) return true; // Skip self ref.
-                                            return processor.process(new UsageInfo(ref));
-                                        });
-                    if (!success) return false;
-                }
             }
         }
         if (searchText) {
@@ -73,5 +59,55 @@ public class SdFindUsagesHandler extends FindUsagesHandler {
         }
         return true;
     }
-    
+
+    private void findFunctionUsages(SdFunctionDefinition functionToFind, Processor<? super UsageInfo> processor) {
+        String functionNameToFind = ReadAction.compute(functionToFind::getName);
+        PsiFile file = getPsiElement().getContainingFile();
+        var rankProfileDefinition = PsiTreeUtil.getParentOfType(functionToFind, SdRankProfileDefinition.class);
+        Schema schema;
+        if (file.getVirtualFile().getPath().endsWith(".profile")) {
+            Path schemaFile = Path.fromString(file.getVirtualFile().getParent().getPath() + ".sd");
+            schema = ReadAction.compute(() -> Schema.fromProjectFile(getProject(), schemaFile));
+        }
+        else { // schema
+            schema = ReadAction.compute(() -> Schema.fromProjectFile(getProject(), Path.fromString(file.getVirtualFile().getPath())));
+        }
+        var rankProfile = ReadAction.compute(() -> schema.rankProfiles().get(rankProfileDefinition.getName()));
+        findFunctionUsages(functionNameToFind, functionToFind, rankProfile, processor);
+    }
+
+    private void findFunctionUsages(String functionNameToFind,
+                                    SdFunctionDefinition functionToFind,
+                                    RankProfile rankProfile,
+                                    Processor<? super UsageInfo> processor) {
+        ReadAction.compute(() -> findFunctionUsagesInThis(functionNameToFind, functionToFind, rankProfile, processor));
+        Collection<RankProfile> children = ReadAction.compute(() -> rankProfile.children().values());
+        for (var child : children)
+            findFunctionUsages(functionNameToFind, functionToFind, child, processor);
+    }
+
+    private boolean findFunctionUsagesInThis(String functionNameToFind,
+                                             SdFunctionDefinition functionToFind,
+                                             RankProfile rankProfile,
+                                             Processor<? super UsageInfo> processor) {
+        Collection<List<Function>> functions = ReadAction.compute(() -> rankProfile.definedFunctions().values());
+        for (var functionList : functions) {
+            for (var function : functionList) {
+                String text = ReadAction.compute(() -> function.definition().getText());
+                int offset = 0;
+                boolean skipNext = functionToFind == function.definition(); // Skip the definition itself
+                while (offset < text.length()) {
+                    int occurrenceStart = text.indexOf(functionNameToFind, offset);
+                    if (occurrenceStart < 0) break;
+                    int occurrenceEnd = occurrenceStart + functionNameToFind.length();
+                    if ( ! skipNext)
+                        processor.process(new UsageInfo(function.definition(), occurrenceStart, occurrenceEnd));
+                    offset = occurrenceEnd;
+                    skipNext = false;
+                }
+            }
+        }
+        return true;
+    }
+
 }
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/RankProfile.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/RankProfile.java
index 53925f86984..514372f88e0 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/RankProfile.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/RankProfile.java
@@ -1,8 +1,10 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.intellij.schema.model;
 
+import ai.vespa.intellij.schema.psi.SdFirstPhaseDefinition;
 import ai.vespa.intellij.schema.psi.SdFunctionDefinition;
 import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
+import ai.vespa.intellij.schema.psi.SdSecondPhaseDefinition;
 import ai.vespa.intellij.schema.utils.AST;
 import com.intellij.psi.util.PsiTreeUtil;
 
@@ -11,6 +13,7 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
 import java.util.stream.Collectors;
 
 /**
@@ -30,6 +33,9 @@ public class RankProfile {
     /** The profiles inherited by this - lazily initialized. */
     private Map<String, RankProfile> inherited = null;
 
+    /** The children of this - lazily inherited */
+    private Map<String, RankProfile> children = null;
+
     public RankProfile(SdRankProfileDefinition definition, Schema owner) {
         this.definition = Objects.requireNonNull(definition);
         this.owner = owner;
@@ -60,7 +66,28 @@ public class RankProfile {
         functions = new HashMap<>();
         for (SdFunctionDefinition function : PsiTreeUtil.findChildrenOfType(definition, SdFunctionDefinition.class))
             functions.computeIfAbsent(function.getName(), k -> new ArrayList<>()).add(Function.from(function, this));
+        PsiTreeUtil.findChildrenOfType(definition, SdFirstPhaseDefinition.class).stream().findFirst()
+                   .map(firstPhase -> Function.from(firstPhase, this))
+                   .ifPresent(firstPhase -> functions.put(firstPhase.name(), List.of(firstPhase)));
+        PsiTreeUtil.findChildrenOfType(definition, SdSecondPhaseDefinition.class).stream().findFirst()
+                   .map(secondPhase -> Function.from(secondPhase, this))
+                   .ifPresent(secondPhase -> functions.put(secondPhase.name(), List.of(secondPhase)));
         return functions;
     }
 
+    public Map<String, RankProfile> children() {
+        if (children != null) return children;
+        children = new HashMap<>();
+        for (var profile : owner.rankProfiles().values()) {
+            if (profile.inherited().containsKey(this.name()))
+                children.put(profile.name(), profile);
+        }
+        return children;
+    }
+
+    @Override
+    public String toString() {
+        return "rank-profile " + name();
+    }
+
 }
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/Schema.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/Schema.java
index ed797fee757..180d410fe43 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/Schema.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/Schema.java
@@ -40,7 +40,7 @@ public class Schema {
 
     public SdFile definition() { return definition; }
 
-    /** The path of this schema from the project root. */
+    /** The path of the location of this schema from the project root. */
     public Path path() { return Path.fromString(definition.getContainingDirectory().getVirtualFile().getPath()); }
 
     public Optional<Schema> inherited() {
@@ -64,7 +64,7 @@ public class Schema {
         for (var profileDefinition : PsiTreeUtil.collectElementsOfType(definition, SdRankProfileDefinition.class))
             rankProfiles.put(profileDefinition.getName(), new RankProfile(profileDefinition, this));
 
-        for (var profileFile : Files.allFilesIn(path().getParentPath().append(name()), "profile", definition.getProject())) {
+        for (var profileFile : Files.allFilesIn(path().append(name()), "profile", definition.getProject())) {
             var profileDefinitions = PsiTreeUtil.collectElementsOfType(profileFile, SdRankProfileDefinition.class);
             if (profileDefinitions.size() != 1) continue; // invalid file
             var profileDefinition = profileDefinitions.stream().findAny().get();
@@ -82,6 +82,9 @@ public class Schema {
         return functions;
     }
 
+    @Override
+    public String toString() { return "schema " + name(); }
+
     /**
      * Returns the profile of the given name from the given file.
      *
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/SdDeclarationType.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/SdDeclarationType.java
index 129a32153f0..bb140c4ce75 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/SdDeclarationType.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/SdDeclarationType.java
@@ -19,8 +19,8 @@ public enum SdDeclarationType {
     IMPORTED_FIELD("Imported Field"),
     DOCUMENT_SUMMARY("Document-Summary"),
     RANK_PROFILE("Rank Profile"),
-    MACRO("Macro"),
-    MACRO_ARGUMENT("Macro's Argument"),
+    FUNCTION("Function"),
+    FIUNCTION_ARGUMENT("Function argument"),
     QUERY("Query (first use in file)"),
     ITEM_RAW_SCORE("ItemRawScore (first use in file)");
     
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/impl/SdNamedElementImpl.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/impl/SdNamedElementImpl.java
index 04a8b981fb9..8a2588ecf35 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/impl/SdNamedElementImpl.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/impl/SdNamedElementImpl.java
@@ -76,9 +76,9 @@ public abstract class SdNamedElementImpl extends ASTWrapperPsiElement implements
         } else if (this instanceof SdRankProfileDefinition) {
             return SdDeclarationType.RANK_PROFILE;
         } else if (this instanceof SdFunctionDefinition) {
-            return SdDeclarationType.MACRO;
+            return SdDeclarationType.FUNCTION;
         } else if (this instanceof SdArgumentDefinition) {
-            return SdDeclarationType.MACRO_ARGUMENT;
+            return SdDeclarationType.FIUNCTION_ARGUMENT;
         } else if (this instanceof SdDocumentDefinition) {
             return SdDeclarationType.DOCUMENT;
         } else if (this instanceof SdDocumentStructDefinition) {
diff --git a/integration/intellij/src/test/applications/rankprofilemodularity/test.sd b/integration/intellij/src/test/applications/rankprofilemodularity/test.sd
index af08815dcba..05901b4de6a 100644
--- a/integration/intellij/src/test/applications/rankprofilemodularity/test.sd
+++ b/integration/intellij/src/test/applications/rankprofilemodularity/test.sd
@@ -67,17 +67,16 @@ schema test {
 
     }
 
-#    rank-profile in_schema4 {
-#
-#        function f2() {
-#            expression: fieldMatch(title)
-#        }
-#
-#        first-phase {
-#            expression: f2
-#        }
+    rank-profile in_schema4 {
 
-    }
+        function f2() {
+            expression: fieldMatch(title)
+        }
 
+        first-phase {
+            expression: f2
+        }
+
+    }
 
 }
\ No newline at end of file
diff --git a/integration/intellij/src/test/java/ai/vespa/intellij/findUsages/FindUsagesTest.java b/integration/intellij/src/test/java/ai/vespa/intellij/findUsages/FindUsagesTest.java
index 70df63320b1..c0ddb79606b 100644
--- a/integration/intellij/src/test/java/ai/vespa/intellij/findUsages/FindUsagesTest.java
+++ b/integration/intellij/src/test/java/ai/vespa/intellij/findUsages/FindUsagesTest.java
@@ -21,17 +21,29 @@ import java.util.List;
 public class FindUsagesTest extends PluginTestBase {
 
     @Test
+    public void testTmp() {
+        useDir("src/test/applications/rankprofilemodularity");
+        var tester = new UsagesTester("test.sd", getProject());
+        tester.assertFunctionUsages("3 references in parent schema", 3, "outside_schema2", "fo2");
+    }
+
+    @Test
+    public void testTmp2() {
+        useDir("src/test/applications/rankprofilemodularity");
+        var tester = new UsagesTester("test.sd", getProject());
+    }
+
+    @Test
     public void testFindUsages() {
         useDir("src/test/applications/rankprofilemodularity");
         var tester = new UsagesTester("test.sd", getProject());
         tester.assertFunctionUsages("0 refs",       0, "in_schema1", "tensorFunction");
-        tester.assertFunctionUsages("1 local refs + 2 refs in child",  3, "in_schema2", "f2");
+        tester.assertFunctionUsages("1 local ref in first-phase",  1, "in_schema2", "f2");
         tester.assertFunctionUsages("2 local refs", 2, "in_schema2", "ff1");
-        tester.assertFunctionUsages("2 local refs + 1 ref in parent", 3, "in_schema3", "f2");
-        //tester.assertFunctionUsages("1 local ref", 1, "in_schema4", "f2");
-        //tester.assertFunctionUsages("1 local reference ", 1, "outside_schema1", "local1");
-        //tester.assertFunctionUsages("4 local references", 4, "outside_schema1", "local2");
-        //tester.assertFunctionUsages("3 references in parent schema", 3, "outside_schema2", "fo2");
+        tester.assertFunctionUsages("1 local ref", 1, "in_schema4", "f2");
+        tester.assertFunctionUsages("1 local refa", 1, "outside_schema1", "local1");
+        tester.assertFunctionUsages("4 local refs", 4, "outside_schema1", "local2");
+        tester.assertFunctionUsages("3 refs in parent schema", 3, "outside_schema2", "fo2");
     }
 
     private static class UsagesTester {
diff --git a/integration/intellij/src/test/java/ai/vespa/intellij/model/SchemaTest.java b/integration/intellij/src/test/java/ai/vespa/intellij/model/SchemaTest.java
index 420c441d302..fb71e701411 100644
--- a/integration/intellij/src/test/java/ai/vespa/intellij/model/SchemaTest.java
+++ b/integration/intellij/src/test/java/ai/vespa/intellij/model/SchemaTest.java
@@ -52,12 +52,14 @@ public class SchemaTest extends PluginTestBase {
         assertEquals("outside_schema1", profile.inherited().get("outside_schema1").name());
         assertEquals("outside_schema2", profile.inherited().get("outside_schema2").name());
         schema.functions().entrySet().stream().forEach(e -> System.out.println(e));
-        assertEquals(8, schema.functions().size());
+        assertEquals("8 proper functions + first-phase", 9, schema.functions().size());
         assertEquals(schema.rankProfiles().get("in_schema2").definedFunctions().get("ff1"),
                      schema.functions().get("ff1"));
         assertEquals(schema.rankProfiles().get("outside_schema1").definedFunctions().get("local1"),
                      schema.functions().get("local1"));
         assertEquals(4, schema.rankProfiles().get("outside_schema1").definedFunctions().size());
+        assertEquals(1, schema.rankProfiles().get("in_schema1").children().size());
+        assertEquals(4, schema.rankProfiles().get("outside_schema2").children().size());
     }
 
 }