Move into Schema

4 files changed, 20 insertions, 15 deletions

diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/SdUtil.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/SdUtil.java
index 363f6fe88bd..1c671bda07d 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/SdUtil.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/SdUtil.java
@@ -38,8 +38,6 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
-import static org.codehaus.groovy.runtime.DefaultGroovyMethods.collect;
-
 /**
  * Util class for the plugin's code.
  *
@@ -48,16 +46,6 @@ import static org.codehaus.groovy.runtime.DefaultGroovyMethods.collect;
  */
 public class SdUtil {
     
-    public static Map<String, List<Function>> functionsIn(SdFile file) {
-        Map<String, List<Function>> functionsMap = new HashMap<>();
-        for (SdRankProfileDefinition rankProfile : PsiTreeUtil.findChildrenOfType(file, SdRankProfileDefinition.class)) {
-            for (SdFunctionDefinition function : PsiTreeUtil.findChildrenOfType(rankProfile, SdFunctionDefinition.class)) {
-                functionsMap.computeIfAbsent(function.getName(), k -> new ArrayList<>()).add(Function.from(function, null));
-            }
-        }
-        return functionsMap;
-    }
-
     public static String createFunctionDescription(SdFunctionDefinition function) {
         SdRankProfileDefinition rankProfile = PsiTreeUtil.getParentOfType(function, SdRankProfileDefinition.class);
         String rankProfileName;
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 2c7cc8ee95f..00b9590c5e9 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,6 +2,7 @@
 package ai.vespa.intellij.schema.findUsages;
 
 import ai.vespa.intellij.schema.model.Function;
+import ai.vespa.intellij.schema.model.Schema;
 import com.intellij.find.findUsages.FindUsagesHandler;
 import com.intellij.find.findUsages.FindUsagesOptions;
 import com.intellij.openapi.application.ReadAction;
@@ -33,8 +34,8 @@ public class SdFindUsagesHandler extends FindUsagesHandler {
     public SdFindUsagesHandler(PsiElement psiElement) {
         super(psiElement);
         PsiFile file = psiElement.getContainingFile();
-        functionsMap = file instanceof SdFile ? SdUtil.functionsIn((SdFile) psiElement.getContainingFile())
-                                              : new HashMap<>();
+        functionsMap = file instanceof SdFile ? new Schema((SdFile)file, null, psiElement.getProject()).definedFunctions()
+                                              : Map.of();
     }
     
     @Override
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 822c0f5eaf8..2abcb182698 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
@@ -2,6 +2,7 @@
 package ai.vespa.intellij.schema.hierarchy;
 
 import ai.vespa.intellij.schema.model.Function;
+import ai.vespa.intellij.schema.model.Schema;
 import com.intellij.ide.hierarchy.HierarchyNodeDescriptor;
 import com.intellij.ide.hierarchy.HierarchyTreeStructure;
 import com.intellij.openapi.project.Project;
@@ -36,7 +37,7 @@ public abstract class SdCallTreeStructure extends HierarchyTreeStructure {
         super(project, new SdCallHierarchyNodeDescriptor(null, element, true));
         myScopeType = currentScopeType;
         myFile = (SdFile) element.getContainingFile();
-        functionsMap = SdUtil.functionsIn(myFile);
+        functionsMap = new Schema(myFile, null, project).definedFunctions();
         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 2ba1295baf4..73d59b75e8e 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,6 +2,7 @@
 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;
@@ -11,6 +12,10 @@ import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.util.PsiTreeUtil;
 
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 /**
@@ -65,6 +70,16 @@ public class Schema {
         return definition.map(d -> new RankProfile(d, this));
     }
 
+    public Map<String, List<Function>> definedFunctions() {
+        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));
+            }
+        }
+        return functions;
+    }
+
     private Optional<SdRankProfileDefinition> findProfileElement(String name, PsiFile file) {
         return PsiTreeUtil.collectElementsOfType(file, SdRankProfileDefinition.class)
                           .stream()