Test function resolution

7 files changed, 36 insertions, 22 deletions

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 c4d75206af5..d95c2c045fe 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
@@ -35,9 +35,7 @@ public class SdFindUsagesHandler extends FindUsagesHandler {
     public SdFindUsagesHandler(PsiElement psiElement) {
         super(psiElement);
         PsiFile file = psiElement.getContainingFile();
-        functionsMap = file instanceof SdFile ? new Schema((SdFile)file,
-                                                           Path.fromString(file.getContainingDirectory().toString())).functions()
-                                              : Map.of();
+        functionsMap = file instanceof SdFile ? new Schema((SdFile)file).functions() : Map.of();
     }
     
     @Override
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallTreeStructure.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallTreeStructure.java
index e65d659074d..9b94d496e67 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallTreeStructure.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallTreeStructure.java
@@ -37,7 +37,7 @@ public abstract class SdCallTreeStructure extends HierarchyTreeStructure {
         super(project, new SdCallHierarchyNodeDescriptor(null, element, true));
         myScopeType = currentScopeType;
         myFile = (SdFile) element.getContainingFile();
-        functionsMap = new Schema(myFile, null).functions();
+        functionsMap = new Schema(myFile).functions();
         ranksHeritageMap = new HashMap<>();
     }
     
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCalleeTreeStructure.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCalleeTreeStructure.java
index c1d85af28a8..1882768b3d3 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCalleeTreeStructure.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCalleeTreeStructure.java
@@ -45,7 +45,7 @@ public class SdCalleeTreeStructure extends SdCallTreeStructure {
             if (functions.containsKey(((PsiNamedElement) identifier).getName())) {
                 PsiReference identifierRef = identifier.getReference();
                 if (identifierRef != null) {
-                    results.add(Function.from(identifierRef.resolve(), null));
+                    results.add(Function.from((SdFunctionDefinition)identifierRef.resolve(), null));
                 }
             }
         }
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallerTreeStructure.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallerTreeStructure.java
index 86dc69af004..07fe9956089 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallerTreeStructure.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/hierarchy/SdCallerTreeStructure.java
@@ -2,6 +2,7 @@
 package ai.vespa.intellij.schema.hierarchy;
 
 import ai.vespa.intellij.schema.model.Function;
+import ai.vespa.intellij.schema.psi.SdSecondPhaseDefinition;
 import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
@@ -58,6 +59,11 @@ public class SdCallerTreeStructure extends SdCallTreeStructure {
                     SdFirstPhaseDefinition fp = PsiTreeUtil.getParentOfType(psiElement, SdFirstPhaseDefinition.class, false);
                     if (fp != null)
                         results.add(Function.from(fp, null));
+                    else {
+                        SdSecondPhaseDefinition sp = PsiTreeUtil.getParentOfType(psiElement, SdSecondPhaseDefinition.class, false);
+                        if (sp != null)
+                            results.add(Function.from(sp, null));
+                    }
                 }
             });
         }
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/Function.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/Function.java
index f4fdb48dbb1..f77fdbbd178 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/Function.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/Function.java
@@ -11,32 +11,39 @@ import com.intellij.psi.PsiElement;
  */
 public class Function {
 
+    private final String name;
+
     private final PsiElement definition;
 
     private final RankProfile owner;
 
-    public Function(PsiElement definition, RankProfile owner) {
+    private Function(String name, PsiElement definition, RankProfile owner) {
+        this.name = name;
         this.definition = definition;
         this.owner = owner;
     }
 
     public PsiElement definition() { return definition; }
+
+    public String name() { return name; }
+
     public RankProfile owner() { return owner; }
 
+    @Override
+    public String toString() {
+        return "function " + name();
+    }
+
     public static Function from(SdFirstPhaseDefinition firstPhase, RankProfile owner) {
-        return new Function(firstPhase, owner);
+        return new Function("first-phase", firstPhase, owner);
     }
 
     public static Function from(SdSecondPhaseDefinition secondPhase, RankProfile owner) {
-        return new Function(secondPhase, owner);
+        return new Function("second-phase", secondPhase, owner);
     }
 
     public static Function from(SdFunctionDefinition definition, RankProfile owner) {
-        return new Function(definition, owner);
-    }
-
-    public static Function from(PsiElement definition, RankProfile owner) {
-        return new Function(definition, owner);
+        return new Function(definition.getName(), definition, owner);
     }
 
 }
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 cf83864bc56..ed797fee757 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
@@ -26,24 +26,23 @@ public class Schema {
 
     private final SdFile definition;
 
-    /** The path to this schema */
-    private final Path path;
-
     /** 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. */
+    /** The profiles of this, either defined inside it or in separate .profile files. Resolved lazily. */
     private Map<String, RankProfile> rankProfiles = null;
 
-    public Schema(SdFile definition, Path path) {
+    public Schema(SdFile definition) {
         this.definition = definition;
-        this.path = path;
     }
 
     public String name() { return definition.getName().substring(0, definition.getName().length() - 3); }
 
     public SdFile definition() { return definition; }
 
+    /** The path 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;
         Optional<SdSchemaDefinition> schemaDefinition = PsiTreeUtil.collectElementsOfType(definition, SdSchemaDefinition.class).stream().findFirst();
@@ -51,7 +50,7 @@ public class Schema {
         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")));
+                              .map(inheritedNode -> fromProjectFile(definition.getProject(), path().getParentPath().append(inheritedNode.getText() + ".sd")));
     }
 
     /**
@@ -65,7 +64,7 @@ public class Schema {
         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())) {
+        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();
@@ -94,7 +93,7 @@ public class Schema {
             throw new IllegalArgumentException("Could not find file '" + file + "'");
         if ( ! (psiFile.get() instanceof SdFile))
             throw new IllegalArgumentException(file + " is not a schema file");
-        return new Schema((SdFile)psiFile.get(), file);
+        return new Schema((SdFile)psiFile.get());
     }
 
 }
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 3f58f7c2801..a3899c0e936 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
@@ -23,6 +23,7 @@ public class SchemaTest extends PluginTestBase {
         assertEquals(2, profile.inherited().size());
         assertEquals("parent-profile1", profile.inherited().get("parent-profile1").name());
         assertEquals("parent-profile2", profile.inherited().get("parent-profile2").name());
+        assertEquals(0, schema.functions().size());
     }
 
     @Test
@@ -36,6 +37,7 @@ public class SchemaTest extends PluginTestBase {
         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());
     }
 
     @Test
@@ -49,6 +51,8 @@ public class SchemaTest extends PluginTestBase {
         assertEquals(2, profile.inherited().size());
         assertEquals("outside_schema1", profile.inherited().get("outside_schema1").name());
         assertEquals("outside_schema2", profile.inherited().get("outside_schema2").name());
+        schema.functions().entrySet().stream().forEach(e -> System.out.println(e));
+        assertEquals(8, schema.functions().size());
     }
 
 }