Find usages in child schemas

10 files changed, 97 insertions, 58 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 1c671bda07d..1b45c1bf4fd 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
@@ -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.schema;
 
-import ai.vespa.intellij.schema.model.Function;
 import ai.vespa.intellij.schema.model.RankProfile;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.vfs.VirtualFile;
@@ -33,9 +32,7 @@ import ai.vespa.intellij.schema.psi.SdSummaryDefinition;
 
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 
 /**
@@ -158,7 +155,7 @@ public class SdUtil {
                                                              .filter(f -> f.getName().equals(functionName))
                                                              .findAny();
         if (function.isPresent()) return function;
-        for (var parent : profile.inherited().values()) {
+        for (var parent : profile.parents().values()) {
             function = findFunction(functionName, parent);
             if (function.isPresent()) return function;
         }
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 0b00b7a5e83..3ccc3e90e84 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
@@ -100,7 +100,7 @@ public class SdFindUsagesHandler extends FindUsagesHandler {
                                     Processor<? super UsageInfo> processor) {
         ProgressIndicatorProvider.checkCanceled();
         ReadAction.compute(() -> findFunctionUsagesInThis(functionNameToFind, functionToFind, rankProfile, scope, processor));
-        Collection<RankProfile> children = ReadAction.compute(() -> rankProfile.children().values());
+        Collection<RankProfile> children = ReadAction.compute(() -> rankProfile.children());
         for (var child : children)
             findFunctionUsages(functionNameToFind, functionToFind, child, scope, processor);
     }
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 bc4d48bf93b..e86ed1c1b4a 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
@@ -45,9 +45,9 @@ public class SdHierarchyUtil {
     private static boolean isChildOf(RankProfile targetProfile, RankProfile thisProfile) {
 
         if (Objects.equals(thisProfile.name(), targetProfile.name())) return true;
-        return thisProfile.inherited().values()
-                                           .stream()
-                                           .anyMatch(parent -> isChildOf(targetProfile, parent));
+        return thisProfile.parents().values()
+                          .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 514372f88e0..10ecde68ff9 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
@@ -13,7 +13,6 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
-import java.util.Optional;
 import java.util.stream.Collectors;
 
 /**
@@ -27,14 +26,14 @@ public class RankProfile {
 
     private final Schema owner;
 
-    /** The functions defined in this, lazily computed */
+    /** The functions defined in this. Resolved lazily. */
     private Map<String, List<Function>> functions = null;
 
-    /** The profiles inherited by this - lazily initialized. */
-    private Map<String, RankProfile> inherited = null;
+    /** The profiles inherited by this. Resolved lazily. */
+    private Map<String, RankProfile> parents = null;
 
-    /** The children of this - lazily inherited */
-    private Map<String, RankProfile> children = null;
+    /** The children of this. Resolved lazily. */
+    private List<RankProfile> children = null;
 
     public RankProfile(SdRankProfileDefinition definition, Schema owner) {
         this.definition = Objects.requireNonNull(definition);
@@ -50,14 +49,30 @@ public class RankProfile {
      *
      * @return the profiles this inherits from, empty if none
      */
-    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.rankProfiles().get(reference.getCanonicalText()))
-                              .filter(r -> r != null)
-                              .collect(Collectors.toMap(p -> p.name(), p -> p));
+    public Map<String, RankProfile> parents() {
+        if (parents != null) return parents;
+        return parents = AST.inherits(definition).stream()
+                            .map(parentIdentifierAST -> parentIdentifierAST.getPsi().getReference())
+                            .filter(reference -> reference != null)
+                            .map(reference -> owner.rankProfiles().get(reference.getCanonicalText()))
+                            .filter(r -> r != null)
+                            .collect(Collectors.toMap(p -> p.name(), p -> p));
+    }
+
+    public List<RankProfile> children() {
+        if (children != null) return children;
+        return children = children(owner);
+    }
+
+    private List<RankProfile> children(Schema schema) {
+        children = new ArrayList<>();
+        for (var profile : schema.rankProfiles().values()) {
+            if (profile.parents().containsKey(this.name()))
+                children.add(profile);
+        }
+        for (var childSchema : schema.children().values())
+            children.addAll(children(childSchema));
+        return children;
     }
 
     /** Returns the functions defined in this. */
@@ -75,16 +90,6 @@ public class RankProfile {
         return functions;
     }
 
-    public Map<String, RankProfile> children() {
-        if (children != null) return children;
-        children = new HashMap<>();
-        for (var profile : owner.rankProfiles().values()) {
-            if (profile.inherited().containsKey(this.name()))
-                children.put(profile.name(), profile);
-        }
-        return children;
-    }
-
     @Override
     public String toString() {
         return "rank-profile " + name();
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 180d410fe43..c3dbdfbdf18 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
@@ -15,6 +15,7 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 
 /**
@@ -27,7 +28,10 @@ public class Schema {
     private final SdFile definition;
 
     /** The schema this inherits, or empty if none. Resolved lazily. */
-    private Optional<Schema> inherited = null;
+    private Optional<Schema> parent = null;
+
+    /** The children of this. Resolved lazily. */
+    private Map<String, Schema> children = null;
 
     /** The profiles of this, either defined inside it or in separate .profile files. Resolved lazily. */
     private Map<String, RankProfile> rankProfiles = null;
@@ -43,14 +47,25 @@ public class Schema {
     /** The path of the location of this schema from the project root. */
     public Path path() { return Path.fromString(definition.getContainingDirectory().getVirtualFile().getPath()); }
 
-    public Optional<Schema> inherited() {
-        if (inherited != null) return inherited;
+    public Optional<Schema> parent() {
+        if (parent != null) return parent;
         Optional<SdSchemaDefinition> schemaDefinition = PsiTreeUtil.collectElementsOfType(definition, SdSchemaDefinition.class).stream().findFirst();
         if (schemaDefinition.isEmpty()) return Optional.empty(); // No valid schema definition in schema file
-        return inherited = AST.inherits(schemaDefinition.get())
-                              .stream()
-                              .findFirst() // Only one schema can be inherited; ignore any following
-                              .map(inheritedNode -> fromProjectFile(definition.getProject(), path().getParentPath().append(inheritedNode.getText() + ".sd")));
+        return parent = AST.inherits(schemaDefinition.get())
+                           .stream()
+                           .findFirst() // Only one schema can be inherited; ignore any following
+                           .map(inheritedNode -> fromProjectFile(definition.getProject(), path().getParentPath().append(inheritedNode.getText() + ".sd")));
+    }
+
+    public Map<String, Schema> children() {
+        if (children != null) return children;
+        children = new HashMap<>();
+        for (PsiFile file : Files.allFilesIn(path(),"sd", definition().getProject())) {
+            Schema schema = new Schema((SdFile)file);
+            if ( schema.parent().isPresent() && Objects.equals(schema.parent().get().name(), name()))
+                children.put(schema.name(), schema);
+        }
+        return children;
     }
 
     /**
@@ -59,7 +74,7 @@ public class Schema {
      */
     public Map<String, RankProfile> rankProfiles() {
         if (rankProfiles != null) return rankProfiles;
-        rankProfiles = inherited().isPresent() ? inherited().get().rankProfiles() : new HashMap<>();
+        rankProfiles = parent().isPresent() ? parent().get().rankProfiles() : new HashMap<>();
 
         for (var profileDefinition : PsiTreeUtil.collectElementsOfType(definition, SdRankProfileDefinition.class))
             rankProfiles.put(profileDefinition.getName(), new RankProfile(profileDefinition, this));
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 b022101be2d..e23046f1922 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
@@ -2,7 +2,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;
 
 /**
  * A function's declaration in the SD language.
@@ -15,7 +14,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).inherited().values()) {
+        for (var parentProfile : new RankProfile(thisRankProfile, null).parents().values()) {
             if (containsFunction(functionName, parentProfile))
                 return true;
         }
@@ -25,7 +24,7 @@ public interface SdFunctionDefinitionInterface extends SdDeclaration {
     default boolean containsFunction(String functionName, RankProfile rankProfile) {
         if (rankProfile.definedFunctions().containsKey(functionName))
             return true;
-        for (var parentProfile : rankProfile.inherited().values()) {
+        for (var parentProfile : rankProfile.parents().values()) {
             if (containsFunction(functionName, parentProfile))
                 return true;
         }
diff --git a/integration/intellij/src/test/applications/schemainheritance/child.sd b/integration/intellij/src/test/applications/schemainheritance/child.sd
index e45f5cb3991..ff09bad4442 100644
--- a/integration/intellij/src/test/applications/schemainheritance/child.sd
+++ b/integration/intellij/src/test/applications/schemainheritance/child.sd
@@ -23,6 +23,9 @@ schema child inherits parent {
     }
 
     rank-profile  child_profile   inherits other_child_profile, parent_profile {
+        second-phase {
+            expression: parentFunction
+        }
     }
 
     rank-profile other_child_profile {
diff --git a/integration/intellij/src/test/applications/schemainheritance/parent.sd b/integration/intellij/src/test/applications/schemainheritance/parent.sd
index 51b11dad444..ba666313af1 100644
--- a/integration/intellij/src/test/applications/schemainheritance/parent.sd
+++ b/integration/intellij/src/test/applications/schemainheritance/parent.sd
@@ -23,6 +23,9 @@ schema parent {
         indexing: input pf1 | lowercase | index | attribute | summary
     }
     rank-profile parent_profile {
+        function parentFunction() {
+            expression: random
+        }
     }
     constant parent_constant {
         file: constants/my_constant_tensor_file.json
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 bc9c7959506..41a216bb0be 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
@@ -20,7 +20,7 @@ import java.util.List;
 public class FindUsagesTest extends PluginTestBase {
 
     @Test
-    public void testFindUsages() {
+    public void testFindUsagesInRankProfileModularity() {
         useDir("src/test/applications/rankprofilemodularity");
         var tester = new UsagesTester("test.sd", getProject());
         tester.assertFunctionUsages("0 refs",       0, "in_schema1", "tensorFunction");
@@ -33,6 +33,13 @@ public class FindUsagesTest extends PluginTestBase {
     }
 
     @Test
+    public void testFindUsagesInSchemaInheritance() {
+        useDir("src/test/applications/schemainheritance");
+        var tester = new UsagesTester("parent.sd", getProject());
+        tester.assertFunctionUsages("1 ref in schild schema", 1, "parent_profile", "parentFunction");
+    }
+
+    @Test
     public void testUsageDetails() {
         useDir("src/test/applications/rankprofilemodularity");
         var tester = new UsagesTester("test.sd", getProject());
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 9cf10cdffa2..52b1ea6fd45 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
@@ -20,24 +20,34 @@ public class SchemaTest extends PluginTestBase {
         assertEquals("simple", schema.name());
         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());
-        assertEquals("parent-profile2", profile.inherited().get("parent-profile2").name());
+        assertEquals(2, profile.parents().size());
+        assertEquals("parent-profile1", profile.parents().get("parent-profile1").name());
+        assertEquals("parent-profile2", profile.parents().get("parent-profile2").name());
         assertEquals(0, schema.functions().size());
     }
 
     @Test
     public void testSchemaInheritance() {
         useDir("src/test/applications/schemaInheritance");
-        Schema schema = Schema.fromProjectFile(getProject(), Path.fromString("child.sd"));
-        assertNotNull(schema);
-        assertEquals("child", schema.name());
-        RankProfile profile = schema.rankProfiles().get("child_profile");
+        Schema child = Schema.fromProjectFile(getProject(), Path.fromString("child.sd"));
+        assertNotNull(child);
+        assertEquals("child", child.name());
+        assertEquals("parent", child.parent().get().name());
+        Schema parent = child.parent().get();
+        assertEquals("child", parent.children().get("child").name());
+
+        assertEquals(3, child.rankProfiles().size());
+        assertTrue(child.rankProfiles().containsKey("child_profile"));
+        assertTrue(child.rankProfiles().containsKey("other_child_profile"));
+        assertTrue(child.rankProfiles().containsKey("parent_profile"));
+
+        RankProfile profile = child.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());
-        assertEquals("parent_profile", profile.inherited().get("parent_profile").name());
-        assertEquals(0, schema.functions().size());
+        assertEquals(2, profile.parents().size());
+        assertEquals("other_child_profile", profile.parents().get("other_child_profile").name());
+        assertEquals("parent_profile", profile.parents().get("parent_profile").name());
+        assertEquals("child_profile", profile.parents().get("parent_profile").children().get(0).name());
+        assertEquals(2, child.functions().size());
     }
 
     @Test
@@ -48,9 +58,9 @@ public class SchemaTest extends PluginTestBase {
         assertEquals("test", schema.name());
         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());
-        assertEquals("outside_schema2", profile.inherited().get("outside_schema2").name());
+        assertEquals(2, profile.parents().size());
+        assertEquals("outside_schema1", profile.parents().get("outside_schema1").name());
+        assertEquals("outside_schema2", profile.parents().get("outside_schema2").name());
         assertEquals("8 proper functions + first-phase", 9, schema.functions().size());
         assertEquals(schema.rankProfiles().get("in_schema2").definedFunctions().get("ff1"),
                      schema.functions().get("ff1"));