Get all profiles

5 files changed, 39 insertions, 25 deletions

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 3a1a4b2201f..53925f86984 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
@@ -49,8 +49,8 @@ public class RankProfile {
         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())
+                              .map(reference -> owner.rankProfiles().get(reference.getCanonicalText()))
+                              .filter(r -> r != null)
                               .collect(Collectors.toMap(p -> p.name(), p -> p));
     }
 
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 94c4a85c182..ad1f1bca149 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
@@ -33,6 +33,9 @@ public class Schema {
     /** The schema this inherits, or empty if none. Resolved lazily. */
     private Optional<Schema> inherited = null;
 
+    /** The profiles of this, either defined inside it or in separatew .profile files. Resolved lazily. */
+    private Map<String, RankProfile> rankProfiles = null;
+
     public Schema(SdFile definition, Path path) {
         this.definition = definition;
         this.path = path;
@@ -52,19 +55,24 @@ public class Schema {
                               .map(inheritedNode -> fromProjectFile(definition.getProject(), path.getParentPath().append(inheritedNode.getText() + ".sd")));
     }
 
-    /** Returns a rank profile belonging to this, defined either inside it or in a separate .profile file */
-    public Optional<RankProfile> rankProfile(String name) {
-        var definition = findProfileElement(name, this.definition); // Look up in this
-        if (definition.isEmpty()) { // Look up in a separate file schema-name/profile-name.profile
-            Optional<PsiFile> file = Files.open(path.getParentPath().append(name()).append(name + ".profile"),
-                                                this.definition.getProject());
-            if (file.isPresent())
-                definition = findProfileElement(name, file.get());
-        }
-        if (definition.isEmpty() && inherited().isPresent()) { // Look up in parent
-            return inherited().get().rankProfile(name);
+    /**
+     * Returns the rank profiles in this, defined either inside it or in separate .profile files,
+     * or inherited from the parent schema.
+     */
+    public Map<String, RankProfile> rankProfiles() {
+        if (rankProfiles != null) return rankProfiles;
+        rankProfiles = inherited().isPresent() ? inherited().get().rankProfiles() : new HashMap<>();
+
+        for (var profileDefinition : PsiTreeUtil.collectElementsOfType(definition, SdRankProfileDefinition.class))
+            rankProfiles.put(profileDefinition.getName(), new RankProfile(profileDefinition, this));
+
+        for (var profileFile : Files.allFilesIn(path.getParentPath().append(name()), "profile", definition.getProject())) {
+            var profileDefinitions = PsiTreeUtil.collectElementsOfType(profileFile, SdRankProfileDefinition.class);
+            if (profileDefinitions.size() != 1) continue; // invalid file
+            var profileDefinition = profileDefinitions.stream().findAny().get();
+            rankProfiles.put(profileDefinition.getName(), new RankProfile(profileDefinition, this));
         }
-        return definition.map(d -> new RankProfile(d, this));
+        return rankProfiles;
     }
 
     public Map<String, List<Function>> definedFunctions() {
@@ -77,13 +85,6 @@ public class Schema {
         return functions;
     }
 
-    private Optional<SdRankProfileDefinition> findProfileElement(String name, PsiFile file) {
-        return PsiTreeUtil.collectElementsOfType(file, SdRankProfileDefinition.class)
-                          .stream()
-                          .filter(p -> name.equals(p.getName()))
-                          .findAny();
-    }
-
     /**
      * Returns the profile of the given name from the given file.
      *
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/utils/Files.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/utils/Files.java
index 7ac88a9650d..0f1fece1439 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/utils/Files.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/utils/Files.java
@@ -9,7 +9,9 @@ import com.intellij.psi.search.FilenameIndex;
 import com.intellij.psi.search.GlobalSearchScope;
 
 import java.util.Collection;
+import java.util.List;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 /**
  * Provide workable file operations on top of IntelliJ's API.
@@ -33,4 +35,15 @@ public class Files {
         return Optional.empty();
     }
 
+    /** Returns all the files in the given directory with the given ending. */
+    public static List<PsiFile> allFilesIn(Path dir, String extension, Project project) {
+        if (extension.startsWith("."))
+            extension = extension.substring(1);
+        return FilenameIndex.getAllFilesByExt(project, extension)
+                            .stream()
+                            .filter(file -> file.getParent().getPath().endsWith(dir.getRelative()))
+                            .map(file -> PsiManager.getInstance(project).findFile(file))
+                            .collect(Collectors.toList());
+    }
+
 }
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 0979b861098..70df63320b1 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 = schema.rankProfile(profileName).get().definedFunctions().get(functionName).get(0).definition();
+            var function = schema.rankProfiles().get(profileName).definedFunctions().get(functionName).get(0).definition();
             var usageProcessor = new MockUsageProcessor();
             var options = new FindUsagesOptions(project);
             options.isUsages = 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 f6e1ec3bdf9..3f58f7c2801 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
@@ -18,7 +18,7 @@ public class SchemaTest extends PluginTestBase {
         Schema schema = Schema.fromProjectFile(getProject(), Path.fromString("simple.sd"));
         assertNotNull(schema);
         assertEquals("simple", schema.name());
-        RankProfile profile = schema.rankProfile("simple-profile").get();
+        RankProfile profile = schema.rankProfiles().get("simple-profile");
         assertEquals("simple-profile", profile.name());
         assertEquals(2, profile.inherited().size());
         assertEquals("parent-profile1", profile.inherited().get("parent-profile1").name());
@@ -31,7 +31,7 @@ public class SchemaTest extends PluginTestBase {
         Schema schema = Schema.fromProjectFile(getProject(), Path.fromString("child.sd"));
         assertNotNull(schema);
         assertEquals("child", schema.name());
-        RankProfile profile = schema.rankProfile("child_profile").get();
+        RankProfile profile = schema.rankProfiles().get("child_profile");
         assertEquals("child_profile", profile.name());
         assertEquals(2, profile.inherited().size());
         assertEquals("other_child_profile", profile.inherited().get("other_child_profile").name());
@@ -44,7 +44,7 @@ public class SchemaTest extends PluginTestBase {
         Schema schema = Schema.fromProjectFile(getProject(), Path.fromString("test.sd"));
         assertNotNull(schema);
         assertEquals("test", schema.name());
-        RankProfile profile = schema.rankProfile("in_schema3").get();
+        RankProfile profile = schema.rankProfiles().get("in_schema3");
         assertEquals("in_schema3", profile.name());
         assertEquals(2, profile.inherited().size());
         assertEquals("outside_schema1", profile.inherited().get("outside_schema1").name());