Pull into class

5 files changed, 60 insertions, 51 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 c5f9c7fed35..3d38b071ed3 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
@@ -2,6 +2,7 @@
 package ai.vespa.intellij.schema;
 
 import ai.vespa.intellij.schema.model.Function;
+import ai.vespa.intellij.schema.model.RankProfile;
 import com.intellij.lang.ASTNode;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.vfs.VirtualFile;
@@ -32,7 +33,6 @@ import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
 import ai.vespa.intellij.schema.psi.SdSchemaAnnotationDefinition;
 import ai.vespa.intellij.schema.psi.SdSchemaFieldDefinition;
 import ai.vespa.intellij.schema.psi.SdSummaryDefinition;
-import ai.vespa.intellij.schema.psi.SdTypes;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -62,42 +62,15 @@ public class SdUtil {
         }
         return functionsMap;
     }
-    
-    /**
-     * Returns the profiles inherited by this.
-     *
-     * @param baseRankProfile the rank-profile node to find the parents of
-     * @return the profiles this inherits from, empty if none
-     */
-    public static List<SdRankProfileDefinition> getRankProfileParents(SdRankProfileDefinition baseRankProfile) {
-        if (baseRankProfile == null) return List.of();
-        ASTNode inheritsNode = baseRankProfile.getNode().findChildByType(SdTypes.INHERITS);
-        if (inheritsNode == null) return List.of();
-        return inherits(baseRankProfile).stream()
-                                        .map(parentIdentifierAST -> parentIdentifierAST.getPsi().getReference())
-                                        .filter(reference -> reference != null)
-                                        .map(reference -> (SdRankProfileDefinition)reference.resolve())
-                                        .collect(Collectors.toList());
-    }
 
-    private static List<ASTNode> inherits(SdRankProfileDefinition baseRankProfile) {
-        Tokens tokens = Tokens.of(baseRankProfile);
-        tokens.require(SdTypes.RANK_PROFILE);
-        tokens.requireWhitespace();
-        tokens.require(SdTypes.IDENTIFIER_VAL, SdTypes.IDENTIFIER_WITH_DASH_VAL);
-        if ( ! tokens.skipWhitespace()) return List.of();
-        if ( ! tokens.skip(SdTypes.INHERITS)) return List.of();
-        tokens.requireWhitespace();
-        List<ASTNode> inherited = new ArrayList<>();
-        do {
-            inherited.add(tokens.require(SdTypes.IDENTIFIER_VAL, SdTypes.IDENTIFIER_WITH_DASH_VAL));
-            tokens.skipWhitespace();
-            if ( ! tokens.skip(SdTypes.COMMA)) break;
-            tokens.skipWhitespace();
-        } while (true);
-        return inherited;
+    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));
+        }
+        return functionsMap;
     }
-    
+
     public static String createFunctionDescription(SdFunctionDefinition function) {
         SdRankProfileDefinition rankProfile = PsiTreeUtil.getParentOfType(function, SdRankProfileDefinition.class);
         String rankProfileName;
@@ -210,7 +183,7 @@ public class SdUtil {
                                                              .filter(f -> f.getName().equals(functionName))
                                                              .findAny();
         if (function.isPresent()) return function;
-        for (var parent : getRankProfileParents(profile)) {
+        for (var parent : new RankProfile(profile).findInherited()) {
             function = findFunction(functionName, parent);
             if (function.isPresent()) return function;
         }
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdHierarchyUtil.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdHierarchyUtil.java
index 7da47607c0c..da8be16cca2 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdHierarchyUtil.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdHierarchyUtil.java
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.intellij.schema.hierarchy;
 
+import ai.vespa.intellij.schema.model.RankProfile;
 import com.intellij.ide.hierarchy.HierarchyBrowserManager;
 import com.intellij.ide.util.treeView.AlphaComparator;
 import com.intellij.ide.util.treeView.NodeDescriptor;
@@ -8,7 +9,6 @@ import com.intellij.openapi.project.Project;
 import com.intellij.psi.*;
 import com.intellij.psi.util.PsiTreeUtil;
 
-import ai.vespa.intellij.schema.SdUtil;
 import ai.vespa.intellij.schema.psi.SdFile;
 import ai.vespa.intellij.schema.psi.SdFunctionDefinition;
 import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
@@ -43,9 +43,9 @@ public class SdHierarchyUtil {
 
     private static boolean isChildOf(SdRankProfileDefinition targetProfile, SdRankProfileDefinition thisProfile) {
         if (thisProfile.getName().equals(targetProfile.getName())) return true;
-        return SdUtil.getRankProfileParents(thisProfile)
-                     .stream()
-                     .anyMatch(parent -> isChildOf(targetProfile, parent));
+        return new RankProfile(thisProfile).findInherited()
+                                           .stream()
+                                           .anyMatch(parent -> isChildOf(targetProfile, parent));
     }
 
     public static Comparator<NodeDescriptor<?>> getComparator(Project project) {
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 5e927d669dc..27430fcfbcb 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,18 +1,18 @@
 // 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.SdFile;
+import ai.vespa.intellij.schema.SdUtil;
 import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
+import ai.vespa.intellij.schema.psi.SdTypes;
+import com.intellij.lang.ASTNode;
 import com.intellij.openapi.project.Project;
-import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiFile;
-import com.intellij.psi.PsiManager;
-import com.intellij.psi.search.FilenameIndex;
-import com.intellij.psi.search.GlobalSearchScope;
 import com.intellij.psi.util.PsiTreeUtil;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 /**
  * A rank profile
@@ -30,6 +30,40 @@ public class RankProfile {
     public SdRankProfileDefinition definition() { return definition; }
 
     /**
+     * Returns the profiles inherited by this.
+     *
+     * @return the profiles this inherits from, empty if none
+     */
+    public List<SdRankProfileDefinition> findInherited() {
+        ASTNode inheritsNode = definition.getNode().findChildByType(SdTypes.INHERITS);
+        if (inheritsNode == null) return List.of();
+        return inherits().stream()
+                         .map(parentIdentifierAST -> parentIdentifierAST.getPsi().getReference())
+                         .filter(reference -> reference != null)
+                         .map(reference -> (SdRankProfileDefinition)reference.resolve())
+                         .collect(Collectors.toList());
+    }
+
+    /** Returns the nodes following "inherits" in the definition of this */
+    private List<ASTNode> inherits() {
+        SdUtil.Tokens tokens = SdUtil.Tokens.of(definition);
+        tokens.require(SdTypes.RANK_PROFILE);
+        tokens.requireWhitespace();
+        tokens.require(SdTypes.IDENTIFIER_VAL, SdTypes.IDENTIFIER_WITH_DASH_VAL);
+        if ( ! tokens.skipWhitespace()) return List.of();
+        if ( ! tokens.skip(SdTypes.INHERITS)) return List.of();
+        tokens.requireWhitespace();
+        List<ASTNode> inherited = new ArrayList<>();
+        do {
+            inherited.add(tokens.require(SdTypes.IDENTIFIER_VAL, SdTypes.IDENTIFIER_WITH_DASH_VAL));
+            tokens.skipWhitespace();
+            if ( ! tokens.skip(SdTypes.COMMA)) break;
+            tokens.skipWhitespace();
+        } while (true);
+        return inherited;
+    }
+
+    /**
      * Returns the profile of the given name from the given file.
      *
      * @throws IllegalArgumentException if not found
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 f60131ef8cb..06928a46cc7 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
@@ -1,5 +1,6 @@
 package ai.vespa.intellij.schema.psi;
 
+import ai.vespa.intellij.schema.model.RankProfile;
 import com.intellij.psi.util.PsiTreeUtil;
 import ai.vespa.intellij.schema.SdUtil;
 
@@ -14,7 +15,7 @@ public interface SdFunctionDefinitionInterface extends SdDeclaration {
         String functionName = this.getName();
         SdRankProfileDefinition thisRankProfile = PsiTreeUtil.getParentOfType(this, SdRankProfileDefinition.class);
         if (thisRankProfile == null) return false;
-        for (var parentProfile : SdUtil.getRankProfileParents(thisRankProfile)) {
+        for (var parentProfile : new RankProfile(thisRankProfile).findInherited()) {
             if (containsFunction(functionName, parentProfile))
                 return true;
         }
@@ -22,7 +23,9 @@ public interface SdFunctionDefinitionInterface extends SdDeclaration {
     }
 
     default boolean containsFunction(String functionName, SdRankProfileDefinition rankProfile) {
-        for (var parentProfile : SdUtil.getRankProfileParents(rankProfile)) {
+        if (SdUtil.functionsIn(new RankProfile(rankProfile)).containsKey(functionName))
+            return true;
+        for (var parentProfile : new RankProfile(rankProfile).findInherited()) {
             if (containsFunction(functionName, parentProfile))
                 return true;
         }
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 38fd9a97a5b..432509a101c 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
@@ -1,7 +1,6 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.intellij.model;
 
-import ai.vespa.intellij.schema.SdUtil;
 import ai.vespa.intellij.schema.model.RankProfile;
 import ai.vespa.intellij.schema.model.Schema;
 import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
@@ -36,7 +35,7 @@ public class SchemaTest extends LightJavaCodeInsightFixtureTestCase {
         assertEquals("simple.sd", schema.definition().getName());
         RankProfile profile = RankProfile.fromProjectFile(getProject(), "simple.sd", "simple-profile");
         assertEquals("simple-profile", profile.definition().getName());
-        List<SdRankProfileDefinition> parents = SdUtil.getRankProfileParents(profile.definition());
+        List<SdRankProfileDefinition> parents = profile.findInherited();
         assertEquals(2, parents.size());
         assertEquals("parent-profile1", parents.get(0).getName());
         assertEquals("parent-profile2", parents.get(1).getName());
@@ -50,7 +49,7 @@ public class SchemaTest extends LightJavaCodeInsightFixtureTestCase {
         assertEquals("child.sd", schema.definition().getName());
         RankProfile profile = RankProfile.fromProjectFile(getProject(), "child.sd", "child_profile");
         assertEquals("child_profile", profile.definition().getName());
-        List<SdRankProfileDefinition> parents = SdUtil.getRankProfileParents(profile.definition());
+        List<SdRankProfileDefinition> parents = profile.findInherited();
         assertEquals(2, parents.size());
         assertEquals("other_child_profile", parents.get(0).getName());
         // assertEquals("parent-profile", parents.get(1).getName()); TODO