Add function owner

8 files changed, 39 insertions, 16 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 6c4cd81095a..f7985a173c5 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
@@ -52,7 +52,7 @@ public class SdUtil {
         Map<String, List<Function>> functionsMap = new HashMap<>();
         for (SdRankProfileDefinition rankProfile : PsiTreeUtil.findChildrenOfType(file, SdRankProfileDefinition.class)) {
             for (SdFunctionDefinition function : PsiTreeUtil.findChildrenOfType(rankProfile, SdFunctionDefinition.class)) {
-                functionsMap.computeIfAbsent(function.getName(), k -> new ArrayList<>()).add(Function.from(function));
+                functionsMap.computeIfAbsent(function.getName(), k -> new ArrayList<>()).add(Function.from(function, null));
             }
         }
         return functionsMap;
@@ -61,7 +61,7 @@ public class SdUtil {
     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));
+            functionsMap.computeIfAbsent(function.getName(), k -> new ArrayList<>()).add(Function.from(function, null));
         }
         return functionsMap;
     }
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 213955a7e9c..2c7cc8ee95f 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
@@ -53,9 +53,9 @@ public class SdFindUsagesHandler extends FindUsagesHandler {
             } else {
                 String functionName = ReadAction.compute( ((SdFunctionDefinition) elementToSearch)::getName);
                 
-                for (Function functionImpl : functionsMap.get(functionName)) {
+                for (Function function : functionsMap.get(functionName)) {
                     boolean success =
-                        ReferencesSearch.search(createSearchParameters(functionImpl.definition(), scope, options))
+                        ReferencesSearch.search(createSearchParameters(function.definition(), scope, options))
                                         .forEach((PsiReference ref) -> {
                                             if (ref.getElement().getParent() == elementToSearch) return true; // Skip self ref.
                                             return processor.process(new UsageInfo(ref));
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 cd5cfea8de2..c1d85af28a8 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()));
+                    results.add(Function.from(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 4422f869286..86dc69af004 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
@@ -53,11 +53,11 @@ public class SdCallerTreeStructure extends SdCallTreeStructure {
                 PsiElement psiElement = r.getElement();
                 SdFunctionDefinition f = PsiTreeUtil.getParentOfType(psiElement, SdFunctionDefinition.class, false);
                 if (f != null && f.getName() != null && !f.getName().equals(functionName)) {
-                    results.add(Function.from(f));
+                    results.add(Function.from(f, null));
                 } else {
                     SdFirstPhaseDefinition fp = PsiTreeUtil.getParentOfType(psiElement, SdFirstPhaseDefinition.class, false);
                     if (fp != null)
-                        results.add(Function.from(fp));
+                        results.add(Function.from(fp, 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 fc84a195b99..f4fdb48dbb1 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
@@ -13,26 +13,30 @@ public class Function {
 
     private final PsiElement definition;
 
-    public Function(PsiElement definition) {
+    private final RankProfile owner;
+
+    public Function(PsiElement definition, RankProfile owner) {
         this.definition = definition;
+        this.owner = owner;
     }
 
     public PsiElement definition() { return definition; }
+    public RankProfile owner() { return owner; }
 
-    public static Function from(SdFirstPhaseDefinition firstPhase) {
-        return new Function(firstPhase);
+    public static Function from(SdFirstPhaseDefinition firstPhase, RankProfile owner) {
+        return new Function(firstPhase, owner);
     }
 
-    public static Function from(SdSecondPhaseDefinition secondPhase) {
-        return new Function(secondPhase);
+    public static Function from(SdSecondPhaseDefinition secondPhase, RankProfile owner) {
+        return new Function(secondPhase, owner);
     }
 
-    public static Function from(SdFunctionDefinition definition) {
-        return new Function(definition);
+    public static Function from(SdFunctionDefinition definition, RankProfile owner) {
+        return new Function(definition, owner);
     }
 
-    public static Function from(PsiElement definition) {
-        return new Function(definition);
+    public static Function from(PsiElement definition, RankProfile owner) {
+        return new Function(definition, owner);
     }
 
 }
diff --git a/integration/intellij/src/test/applications/rankprofilemodularity/test.sd b/integration/intellij/src/test/applications/rankprofilemodularity/test.sd
index 20b571cbba5..af08815dcba 100644
--- a/integration/intellij/src/test/applications/rankprofilemodularity/test.sd
+++ b/integration/intellij/src/test/applications/rankprofilemodularity/test.sd
@@ -67,4 +67,17 @@ schema test {
 
     }
 
+#    rank-profile in_schema4 {
+#
+#        function f2() {
+#            expression: fieldMatch(title)
+#        }
+#
+#        first-phase {
+#            expression: f2
+#        }
+
+    }
+
+
 }
\ No newline at end of file
diff --git a/integration/intellij/src/test/java/ai/vespa/intellij/PluginTestBase.java b/integration/intellij/src/test/java/ai/vespa/intellij/PluginTestBase.java
index 32f24805c07..1ba2dcf2785 100644
--- a/integration/intellij/src/test/java/ai/vespa/intellij/PluginTestBase.java
+++ b/integration/intellij/src/test/java/ai/vespa/intellij/PluginTestBase.java
@@ -4,6 +4,7 @@ package ai.vespa.intellij;
 import ai.vespa.intellij.model.TestProjectDescriptor;
 import com.intellij.testFramework.LightProjectDescriptor;
 import com.intellij.testFramework.fixtures.LightJavaCodeInsightFixtureTestCase;
+import org.junit.Test;
 
 /**
  * Parent of unit tests. This is an antipattern, but mandated by IntelliJ.
@@ -31,4 +32,8 @@ public class PluginTestBase extends LightJavaCodeInsightFixtureTestCase {
         myFixture.copyDirectoryToProject(dir, "/");
     }
 
+    /** Avoid "no tests" warning */
+    @Test
+    public void testDummy() {}
+
 }
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 48c05c3ff00..9311bc830b4 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
@@ -28,6 +28,7 @@ public class FindUsagesTest extends PluginTestBase {
         tester.assertFunctionUsages("1 local refs + 2 refs in child",  3, "in_schema2", "f2");
         tester.assertFunctionUsages("2 local refs", 2, "in_schema2", "ff1");
         tester.assertFunctionUsages("2 local refs + 1 ref in parent", 3, "in_schema3", "f2");
+        //tester.assertFunctionUsages("1 local ref", 1, "in_schema4", "f2");
         //tester.assertFunctionUsages("1 local reference ", 1, "outside_schema1", "local1");
         //tester.assertFunctionUsages("4 local references", 4, "outside_schema1", "local2");
         //tester.assertFunctionUsages("3 references in parent schema", 3, "outside_schema2", "fo2");