Find all functions

3 files changed, 8 insertions, 8 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 5029d51405f..c4d75206af5 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
@@ -3,6 +3,7 @@ package ai.vespa.intellij.schema.findUsages;
 
 import ai.vespa.intellij.schema.model.Function;
 import ai.vespa.intellij.schema.model.Schema;
+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;
@@ -34,7 +35,8 @@ public class SdFindUsagesHandler extends FindUsagesHandler {
     public SdFindUsagesHandler(PsiElement psiElement) {
         super(psiElement);
         PsiFile file = psiElement.getContainingFile();
-        functionsMap = file instanceof SdFile ? new Schema((SdFile)file, null).definedFunctions()
+        functionsMap = file instanceof SdFile ? new Schema((SdFile)file,
+                                                           Path.fromString(file.getContainingDirectory().toString())).functions()
                                               : Map.of();
     }
     
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallTreeStructure.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallTreeStructure.java
index 50c65b21a81..e65d659074d 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallTreeStructure.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallTreeStructure.java
@@ -37,7 +37,7 @@ public abstract class SdCallTreeStructure extends HierarchyTreeStructure {
         super(project, new SdCallHierarchyNodeDescriptor(null, element, true));
         myScopeType = currentScopeType;
         myFile = (SdFile) element.getContainingFile();
-        functionsMap = new Schema(myFile, null).definedFunctions();
+        functionsMap = new Schema(myFile, null).functions();
         ranksHeritageMap = new HashMap<>();
     }
     
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 ad1f1bca149..cf83864bc56 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
@@ -2,7 +2,6 @@
 package ai.vespa.intellij.schema.model;
 
 import ai.vespa.intellij.schema.psi.SdFile;
-import ai.vespa.intellij.schema.psi.SdFunctionDefinition;
 import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
 import ai.vespa.intellij.schema.psi.SdSchemaDefinition;
 import ai.vespa.intellij.schema.utils.AST;
@@ -75,12 +74,11 @@ public class Schema {
         return rankProfiles;
     }
 
-    public Map<String, List<Function>> definedFunctions() {
+    public Map<String, List<Function>> functions() {
         Map<String, List<Function>> functions = new HashMap<>();
-        for (SdRankProfileDefinition rankProfile : PsiTreeUtil.findChildrenOfType(definition, SdRankProfileDefinition.class)) {
-            for (SdFunctionDefinition function : PsiTreeUtil.findChildrenOfType(rankProfile, SdFunctionDefinition.class)) {
-                functions.computeIfAbsent(function.getName(), k -> new ArrayList<>()).add(Function.from(function, null));
-            }
+        for (var profile : rankProfiles().values()) {
+            for (var entry : profile.definedFunctions().entrySet())
+                functions.computeIfAbsent(entry.getKey(), k -> new ArrayList<>()).addAll(entry.getValue());
         }
         return functions;
     }