Move into RankProfile

4 files changed, 22 insertions, 14 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 f7985a173c5..363f6fe88bd 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
@@ -58,14 +58,6 @@ public class SdUtil {
         return functionsMap;
     }
 
-    public static Map<String, List<Function>> functionsIn(RankProfile profile) {
-        Map<String, List<Function>> functionsMap = new HashMap<>();
-        for (SdFunctionDefinition function : PsiTreeUtil.findChildrenOfType(profile.definition(), 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/model/RankProfile.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/RankProfile.java
index 6aa43de59a9..3a1a4b2201f 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,9 +1,13 @@
 // 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.SdFunctionDefinition;
 import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
 import ai.vespa.intellij.schema.utils.AST;
+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.Objects;
@@ -20,6 +24,9 @@ public class RankProfile {
 
     private final Schema owner;
 
+    /** The functions defined in this, lazily computed */
+    private Map<String, List<Function>> functions = null;
+
     /** The profiles inherited by this - lazily initialized. */
     private Map<String, RankProfile> inherited = null;
 
@@ -47,4 +54,13 @@ public class RankProfile {
                               .collect(Collectors.toMap(p -> p.name(), p -> p));
     }
 
+    /** Returns the functions defined in this. */
+    public Map<String, List<Function>> definedFunctions() {
+        if (functions != null) return functions;
+        functions = new HashMap<>();
+        for (SdFunctionDefinition function : PsiTreeUtil.findChildrenOfType(definition, SdFunctionDefinition.class))
+            functions.computeIfAbsent(function.getName(), k -> new ArrayList<>()).add(Function.from(function, this));
+        return functions;
+    }
+
 }
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/SdFunctionDefinitionInterface.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/SdFunctionDefinitionInterface.java
index b5d1bb7d44b..b022101be2d 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/SdFunctionDefinitionInterface.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/psi/SdFunctionDefinitionInterface.java
@@ -16,17 +16,17 @@ public interface SdFunctionDefinitionInterface extends SdDeclaration {
         SdRankProfileDefinition thisRankProfile = PsiTreeUtil.getParentOfType(this, SdRankProfileDefinition.class);
         if (thisRankProfile == null) return false;
         for (var parentProfile : new RankProfile(thisRankProfile, null).inherited().values()) {
-            if (containsFunction(functionName, parentProfile.definition()))
+            if (containsFunction(functionName, parentProfile))
                 return true;
         }
         return false;
     }
 
-    default boolean containsFunction(String functionName, SdRankProfileDefinition rankProfile) {
-        if (SdUtil.functionsIn(new RankProfile(rankProfile, null)).containsKey(functionName))
+    default boolean containsFunction(String functionName, RankProfile rankProfile) {
+        if (rankProfile.definedFunctions().containsKey(functionName))
             return true;
-        for (var parentProfile : new RankProfile(rankProfile, null).inherited().values()) {
-            if (containsFunction(functionName, parentProfile.definition()))
+        for (var parentProfile : rankProfile.inherited().values()) {
+            if (containsFunction(functionName, parentProfile))
                 return true;
         }
         return false;
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 9311bc830b4..0979b861098 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
@@ -47,7 +47,7 @@ public class FindUsagesTest extends PluginTestBase {
         }
 
         void assertFunctionUsages(String explanation, int expectedUsages, String profileName, String functionName) {
-            var function = SdUtil.functionsIn(schema.rankProfile(profileName).get()).get(functionName).get(0).definition();
+            var function = schema.rankProfile(profileName).get().definedFunctions().get(functionName).get(0).definition();
             var usageProcessor = new MockUsageProcessor();
             var options = new FindUsagesOptions(project);
             options.isUsages = true;