Match whole elements

3 files changed, 45 insertions, 16 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 c616800eee2..0b00b7a5e83 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
@@ -4,14 +4,18 @@ 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.SdNamedElement;
 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;
 import com.intellij.openapi.application.ReadAction;
+import com.intellij.openapi.progress.ProgressIndicatorProvider;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiElementVisitor;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.PsiReference;
+import com.intellij.psi.impl.source.tree.LeafPsiElement;
 import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.search.SearchScope;
 import com.intellij.psi.search.searches.ReferencesSearch;
@@ -94,6 +98,7 @@ public class SdFindUsagesHandler extends FindUsagesHandler {
                                     RankProfile rankProfile,
                                     SearchScope scope,
                                     Processor<? super UsageInfo> processor) {
+        ProgressIndicatorProvider.checkCanceled();
         ReadAction.compute(() -> findFunctionUsagesInThis(functionNameToFind, functionToFind, rankProfile, scope, processor));
         Collection<RankProfile> children = ReadAction.compute(() -> rankProfile.children().values());
         for (var child : children)
@@ -109,21 +114,45 @@ public class SdFindUsagesHandler extends FindUsagesHandler {
         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;
-                }
+                var matchingVisitor = new MatchingVisitor(functionNameToFind,
+                                                          functionToFind == function.definition(),
+                                                          processor);
+                ReadAction.compute(() -> { function.definition().accept(matchingVisitor); return null; } );
             }
         }
         return true;
     }
 
+    private static class MatchingVisitor extends PsiElementVisitor {
+
+        private final String textToMatch;
+        private final Processor<? super UsageInfo> processor;
+
+        private boolean skipNextMatch;
+
+        public MatchingVisitor(String textToMatch, boolean skipFirstMatch, Processor<? super UsageInfo> processor) {
+            this.textToMatch = textToMatch;
+            this.skipNextMatch = skipFirstMatch;
+            this.processor = processor;
+        }
+
+        @Override
+        public void visitElement(PsiElement element) {
+            if (element instanceof LeafPsiElement)
+                visitThis(element);
+            else
+                element.acceptChildren(this);
+        }
+
+        private void visitThis(PsiElement element) {
+            if ( ! textToMatch.equals(element.getText())) return;
+            if (skipNextMatch) {
+                skipNextMatch = false;
+                return;
+            }
+            processor.process(new UsageInfo(element));
+        }
+
+    }
+
 }
diff --git a/integration/intellij/src/test/applications/rankprofilemodularity/test/outside_schema1.profile b/integration/intellij/src/test/applications/rankprofilemodularity/test/outside_schema1.profile
index de1dc012dea..c7803b41bdd 100644
--- a/integration/intellij/src/test/applications/rankprofilemodularity/test/outside_schema1.profile
+++ b/integration/intellij/src/test/applications/rankprofilemodularity/test/outside_schema1.profile
@@ -5,15 +5,15 @@ rank-profile outside_schema1 inherits in_schema1 {
     }
 
     function local1() {
-        expression: local2 + local3 + local2
+        expression: local12 + local3 + local12
     }
 
-    function local2() {
+    function local12() {
         expression: now
     }
 
     function local3() {
-        expression: local2 + local2
+        expression: local12 + local12
     }
 
 }
\ 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 fd21e7ec4ad..bc9c7959506 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
@@ -28,7 +28,7 @@ public class FindUsagesTest extends PluginTestBase {
         tester.assertFunctionUsages("2 local refs", 2, "in_schema2", "ff1");
         tester.assertFunctionUsages("1 local ref", 1, "in_schema4", "f2");
         tester.assertFunctionUsages("1 local ref", 1, "outside_schema1", "local1");
-        tester.assertFunctionUsages("4 local refs", 4, "outside_schema1", "local2");
+        tester.assertFunctionUsages("4 local refs", 4, "outside_schema1", "local12");
         tester.assertFunctionUsages("3 refs in parent schema", 3, "outside_schema2", "fo2");
     }