Cache and index parent profiles

5 files changed, 27 insertions, 29 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 558c0d84e1e..6c4cd81095a 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
@@ -3,17 +3,14 @@ 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;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiManager;
 import com.intellij.psi.PsiNamedElement;
 import com.intellij.psi.PsiReference;
-import com.intellij.psi.PsiWhiteSpace;
 import com.intellij.psi.search.FileTypeIndex;
 import com.intellij.psi.search.GlobalSearchScope;
-import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.util.PsiTreeUtil;
 import ai.vespa.intellij.schema.psi.SdAnnotationFieldDefinition;
 import ai.vespa.intellij.schema.psi.SdArgumentDefinition;
@@ -35,13 +32,11 @@ import ai.vespa.intellij.schema.psi.SdSchemaFieldDefinition;
 import ai.vespa.intellij.schema.psi.SdSummaryDefinition;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.stream.Collectors;
 
 import static org.codehaus.groovy.runtime.DefaultGroovyMethods.collect;
 
@@ -183,7 +178,7 @@ public class SdUtil {
                                                              .filter(f -> f.getName().equals(functionName))
                                                              .findAny();
         if (function.isPresent()) return function;
-        for (var parent : profile.findInherited()) {
+        for (var parent : profile.inherited().values()) {
             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 9abd9511c84..3410fa76482 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
@@ -14,6 +14,7 @@ import ai.vespa.intellij.schema.psi.SdFunctionDefinition;
 import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
 
 import java.util.Comparator;
+import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -42,8 +43,8 @@ public class SdHierarchyUtil {
     }
 
     private static boolean isChildOf(SdRankProfileDefinition targetProfile, SdRankProfileDefinition thisProfile) {
-        if (thisProfile.getName().equals(targetProfile.getName())) return true;
-        return new RankProfile(thisProfile, null).findInherited()
+        if (Objects.equals(thisProfile.getName(), targetProfile.getName())) return true;
+        return new RankProfile(thisProfile, null).inherited().values()
                                            .stream()
                                            .anyMatch(parent -> isChildOf(targetProfile, parent.definition()));
     }
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 9e10905f95c..6aa43de59a9 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
@@ -5,6 +5,7 @@ import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
 import ai.vespa.intellij.schema.utils.AST;
 
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.stream.Collectors;
 
@@ -19,6 +20,9 @@ public class RankProfile {
 
     private final Schema owner;
 
+    /** The profiles inherited by this - lazily initialized. */
+    private Map<String, RankProfile> inherited = null;
+
     public RankProfile(SdRankProfileDefinition definition, Schema owner) {
         this.definition = Objects.requireNonNull(definition);
         this.owner = owner;
@@ -33,13 +37,14 @@ public class RankProfile {
      *
      * @return the profiles this inherits from, empty if none
      */
-    public List<RankProfile> findInherited() {
-        return AST.inherits(definition).stream()
-                  .map(parentIdentifierAST -> parentIdentifierAST.getPsi().getReference())
-                  .filter(reference -> reference != null)
-                  .map(reference -> owner.rankProfile(reference.getCanonicalText()))
-                  .flatMap(r -> r.stream())
-                  .collect(Collectors.toList());
+    public Map<String, RankProfile> inherited() {
+        if (inherited != null) return inherited;
+        return inherited = AST.inherits(definition).stream()
+                              .map(parentIdentifierAST -> parentIdentifierAST.getPsi().getReference())
+                              .filter(reference -> reference != null)
+                              .map(reference -> owner.rankProfile(reference.getCanonicalText()))
+                              .flatMap(r -> r.stream())
+                              .collect(Collectors.toMap(p -> p.name(), p -> p));
     }
 
 }
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 189c02774c2..b5d1bb7d44b 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
@@ -15,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 : new RankProfile(thisRankProfile, null).findInherited()) {
+        for (var parentProfile : new RankProfile(thisRankProfile, null).inherited().values()) {
             if (containsFunction(functionName, parentProfile.definition()))
                 return true;
         }
@@ -25,7 +25,7 @@ public interface SdFunctionDefinitionInterface extends SdDeclaration {
     default boolean containsFunction(String functionName, SdRankProfileDefinition rankProfile) {
         if (SdUtil.functionsIn(new RankProfile(rankProfile, null)).containsKey(functionName))
             return true;
-        for (var parentProfile : new RankProfile(rankProfile, null).findInherited()) {
+        for (var parentProfile : new RankProfile(rankProfile, null).inherited().values()) {
             if (containsFunction(functionName, parentProfile.definition()))
                 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 6ce93d3fe28..ba670d1c732 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
@@ -35,10 +35,9 @@ public class SchemaTest extends LightJavaCodeInsightFixtureTestCase {
         assertEquals("simple", schema.name());
         RankProfile profile = schema.rankProfile("simple-profile").get();
         assertEquals("simple-profile", profile.name());
-        List<RankProfile> parents = profile.findInherited();
-        assertEquals(2, parents.size());
-        assertEquals("parent-profile1", parents.get(0).name());
-        assertEquals("parent-profile2", parents.get(1).name());
+        assertEquals(2, profile.inherited().size());
+        assertEquals("parent-profile1", profile.inherited().get("parent-profile1").name());
+        assertEquals("parent-profile2", profile.inherited().get("parent-profile2").name());
     }
 
     @Test
@@ -49,10 +48,9 @@ public class SchemaTest extends LightJavaCodeInsightFixtureTestCase {
         assertEquals("child", schema.name());
         RankProfile profile = schema.rankProfile("child_profile").get();
         assertEquals("child_profile", profile.name());
-        List<RankProfile> parents = profile.findInherited();
-        assertEquals(2, parents.size());
-        assertEquals("other_child_profile", parents.get(0).name());
-        assertEquals("parent_profile", parents.get(1).name());
+        assertEquals(2, profile.inherited().size());
+        assertEquals("other_child_profile", profile.inherited().get("other_child_profile").name());
+        assertEquals("parent_profile", profile.inherited().get("parent_profile").name());
     }
 
     @Test
@@ -63,10 +61,9 @@ public class SchemaTest extends LightJavaCodeInsightFixtureTestCase {
         assertEquals("test", schema.name());
         RankProfile profile = schema.rankProfile("in_schema3").get();
         assertEquals("in_schema3", profile.name());
-        List<RankProfile> parents = profile.findInherited();
-        assertEquals(2, parents.size());
-        assertEquals("outside_schema1", parents.get(0).name());
-        assertEquals("outside_schema2", parents.get(1).name());
+        assertEquals(2, profile.inherited().size());
+        assertEquals("outside_schema1", profile.inherited().get("outside_schema1").name());
+        assertEquals("outside_schema2", profile.inherited().get("outside_schema2").name());
     }
 
 }