Start modelling elements

7 files changed, 102 insertions, 40 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 7c3f6dcb0ec..2bc5fa80e18 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,6 +1,7 @@
 // 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 com.intellij.lang.ASTNode;
 import com.intellij.openapi.project.Project;
 import com.intellij.openapi.vfs.VirtualFile;
@@ -36,6 +37,7 @@ 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;
 
@@ -49,14 +51,14 @@ import static org.codehaus.groovy.runtime.DefaultGroovyMethods.collect;
  */
 public class SdUtil {
     
-    public static HashMap<String, List<PsiElement>> createFunctionsMap(SdFile file) {
-        HashMap<String, List<PsiElement>> macrosMap = new HashMap<>();
+    public static Map<String, List<Function>> functionsIn(SdFile file) {
+        Map<String, List<Function>> functionsMap = new HashMap<>();
         for (SdRankProfileDefinition rankProfile : PsiTreeUtil.findChildrenOfType(file, SdRankProfileDefinition.class)) {
             for (SdFunctionDefinition function : PsiTreeUtil.findChildrenOfType(rankProfile, SdFunctionDefinition.class)) {
-                macrosMap.computeIfAbsent(function.getName(), k -> new ArrayList<>()).add(function);
+                functionsMap.computeIfAbsent(function.getName(), k -> new ArrayList<>()).add(Function.from(function));
             }
         }
-        return macrosMap;
+        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 20b3a725446..fd872b1a364 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
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.intellij.schema.findUsages;
 
+import ai.vespa.intellij.schema.model.Function;
 import com.intellij.find.findUsages.FindUsagesHandler;
 import com.intellij.find.findUsages.FindUsagesOptions;
 import com.intellij.openapi.application.ReadAction;
@@ -18,6 +19,7 @@ import ai.vespa.intellij.schema.psi.SdFunctionDefinition;
 
 import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * This class handles creating the "Find Usages" window.
@@ -26,12 +28,12 @@ import java.util.List;
  */
 public class SdFindUsagesHandler extends FindUsagesHandler {
     
-    private final HashMap<String, List<PsiElement>> functionsMap;
+    private final Map<String, List<Function>> functionsMap;
     
     protected SdFindUsagesHandler(PsiElement psiElement) {
         super(psiElement);
         PsiFile file = psiElement.getContainingFile();
-        functionsMap = file instanceof SdFile ? SdUtil.createFunctionsMap((SdFile) psiElement.getContainingFile())
+        functionsMap = file instanceof SdFile ? SdUtil.functionsIn((SdFile) psiElement.getContainingFile())
                                               : new HashMap<>();
     }
     
@@ -49,11 +51,11 @@ public class SdFindUsagesHandler extends FindUsagesHandler {
                                     .forEach((PsiReference ref) -> processor.process(new UsageInfo(ref)));
                 if (!success) return false;
             } else {
-                String macroName = ReadAction.compute( ((SdFunctionDefinition) elementToSearch)::getName);
+                String functionName = ReadAction.compute( ((SdFunctionDefinition) elementToSearch)::getName);
                 
-                for (PsiElement macroImpl : functionsMap.get(macroName)) {
+                for (Function functionImpl : functionsMap.get(functionName)) {
                     boolean success =
-                        ReferencesSearch.search(createSearchParameters(macroImpl, scope, options))
+                        ReferencesSearch.search(createSearchParameters(functionImpl.definition(), scope, options))
                                         .forEach((PsiReference ref) -> processor.process(new UsageInfo(ref)));
                     if (!success) return false;
                 }
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 6d7c90bdbc1..822c0f5eaf8 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
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.intellij.schema.hierarchy;
 
+import ai.vespa.intellij.schema.model.Function;
 import com.intellij.ide.hierarchy.HierarchyNodeDescriptor;
 import com.intellij.ide.hierarchy.HierarchyTreeStructure;
 import com.intellij.openapi.project.Project;
@@ -15,7 +16,6 @@ import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
 
 import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -29,18 +29,18 @@ public abstract class SdCallTreeStructure extends HierarchyTreeStructure {
 
     protected final String myScopeType;
     protected final SdFile myFile;
-    protected Map<String, List<PsiElement>> functionsMap;
+    protected Map<String, List<Function>> functionsMap;
     protected Map<String, Set<SdRankProfileDefinition>> ranksHeritageMap;
     
     public SdCallTreeStructure(Project project, PsiElement element, String currentScopeType) {
         super(project, new SdCallHierarchyNodeDescriptor(null, element, true));
         myScopeType = currentScopeType;
         myFile = (SdFile) element.getContainingFile();
-        functionsMap = SdUtil.createFunctionsMap(myFile);
+        functionsMap = SdUtil.functionsIn(myFile);
         ranksHeritageMap = new HashMap<>();
     }
     
-    protected abstract Set<PsiElement> getChildren(SdFunctionDefinition element);
+    protected abstract Set<Function> getChildren(SdFunctionDefinition element);
     
     @Override
     protected Object[] buildChildren(HierarchyNodeDescriptor descriptor) {
@@ -57,17 +57,17 @@ public abstract class SdCallTreeStructure extends HierarchyTreeStructure {
                 return ArrayUtilRt.EMPTY_OBJECT_ARRAY;
             }
             
-            Set<PsiElement> children = getChildren((SdFunctionDefinition) element);
+            Set<Function> children = getChildren((SdFunctionDefinition) element);
             
             Map<PsiElement, SdCallHierarchyNodeDescriptor> callerToDescriptorMap = new HashMap<>();
             PsiElement baseClass = PsiTreeUtil.getParentOfType(element, SdRankProfileDefinition.class);
             
-            for (PsiElement caller : children) {
-                if (isInScope(baseClass, caller, myScopeType)) {
+            for (Function caller : children) {
+                if (isInScope(baseClass, caller.definition(), myScopeType)) {
                     SdCallHierarchyNodeDescriptor callerDescriptor = callerToDescriptorMap.get(caller);
                     if (callerDescriptor == null) {
-                        callerDescriptor = new SdCallHierarchyNodeDescriptor(descriptor, caller, false);
-                        callerToDescriptorMap.put(caller, callerDescriptor);
+                        callerDescriptor = new SdCallHierarchyNodeDescriptor(descriptor, caller.definition(), false);
+                        callerToDescriptorMap.put(caller.definition(), callerDescriptor);
                         descriptors.add(callerDescriptor);
                     }
                 }
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 6783d47ba31..de286f6612f 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
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.intellij.schema.hierarchy;
 
+import ai.vespa.intellij.schema.model.Function;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiNamedElement;
@@ -29,12 +30,12 @@ public class SdCalleeTreeStructure extends SdCallTreeStructure {
     }
         
     @Override
-    protected Set<PsiElement> getChildren(SdFunctionDefinition element) {
+    protected Set<Function> getChildren(SdFunctionDefinition element) {
         return getCallees(element, functionsMap);
     }
     
-    private Set<PsiElement> getCallees(SdFunctionDefinition function, Map<String, List<PsiElement>> functions) {
-        Set<PsiElement> results = new HashSet<>();
+    private Set<Function> getCallees(SdFunctionDefinition function, Map<String, List<Function>> functions) {
+        Set<Function> results = new HashSet<>();
         SdExpressionDefinition expression = PsiTreeUtil.findChildOfType(function, SdExpressionDefinition.class);
         if (expression == null) {
             return results;
@@ -43,7 +44,7 @@ public class SdCalleeTreeStructure extends SdCallTreeStructure {
             if (functions.containsKey(((PsiNamedElement) identifier).getName())) {
                 PsiReference identifierRef = identifier.getReference();
                 if (identifierRef != null) {
-                    results.add(identifierRef.resolve());
+                    results.add(Function.from(identifierRef.resolve()));
                 }
             }
         }
@@ -59,8 +60,8 @@ public class SdCalleeTreeStructure extends SdCallTreeStructure {
         
         Set<SdRankProfileDefinition> inheritedRanks = ranksHeritageMap.get(rankProfileName);
         
-        for (PsiElement functionImpl : functions.get(function.getName())) {
-            if (inheritedRanks.contains(PsiTreeUtil.getParentOfType(functionImpl, SdRankProfileDefinition.class))) {
+        for (Function functionImpl : functions.get(function.getName())) {
+            if (inheritedRanks.contains(PsiTreeUtil.getParentOfType(functionImpl.definition(), SdRankProfileDefinition.class))) {
                 results.add(functionImpl);
             }
 
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 7b0a7bf7955..4422f869286 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
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.intellij.schema.hierarchy;
 
+import ai.vespa.intellij.schema.model.Function;
 import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.project.Project;
 import com.intellij.psi.PsiElement;
@@ -26,7 +27,7 @@ import java.util.function.Consumer;
  */
 public class SdCallerTreeStructure extends SdCallTreeStructure {
     
-    private Map<String, Set<PsiElement>> functionTreeChildren;
+    private Map<String, Set<Function>> functionTreeChildren;
     
     public SdCallerTreeStructure(Project project, PsiElement element, String currentScopeType) {
         super(project, element, currentScopeType);
@@ -34,35 +35,34 @@ public class SdCallerTreeStructure extends SdCallTreeStructure {
     }
     
     @Override
-    protected Set<PsiElement> getChildren(SdFunctionDefinition element) {
+    protected Set<Function> getChildren(SdFunctionDefinition element) {
         return getCallers(element, functionsMap);
     }
     
-    private Set<PsiElement> getCallers(SdFunctionDefinition macro, Map<String, List<PsiElement>> macrosMap) {
-        String macroName = macro.getName();
-
-        if (functionTreeChildren.containsKey(macroName)) {
-            return functionTreeChildren.get(macroName);
+    private Set<Function> getCallers(SdFunctionDefinition function, Map<String, List<Function>> functionsMap) {
+        String functionName = function.getName();
+        if (functionTreeChildren.containsKey(functionName)) {
+            return functionTreeChildren.get(functionName);
         }
         
-        Set<PsiElement> results = new HashSet<>();
-        
-        for (PsiElement macroImpl : macrosMap.get(macroName)) {
-            SearchScope searchScope = getSearchScope(myScopeType, macroImpl);
-            ReferencesSearch.search(macroImpl, searchScope).forEach((Consumer<? super PsiReference>) r -> {
+        Set<Function> results = new HashSet<>();
+        for (Function functionImpl : functionsMap.get(functionName)) {
+            SearchScope searchScope = getSearchScope(myScopeType, functionImpl.definition());
+            ReferencesSearch.search(functionImpl.definition(), searchScope).forEach((Consumer<? super PsiReference>) r -> {
                 ProgressManager.checkCanceled();
                 PsiElement psiElement = r.getElement();
                 SdFunctionDefinition f = PsiTreeUtil.getParentOfType(psiElement, SdFunctionDefinition.class, false);
-                if (f != null && f.getName() != null && !f.getName().equals(macroName)) { 
-                    ContainerUtil.addIfNotNull(results, f); 
+                if (f != null && f.getName() != null && !f.getName().equals(functionName)) {
+                    results.add(Function.from(f));
                 } else {
                     SdFirstPhaseDefinition fp = PsiTreeUtil.getParentOfType(psiElement, SdFirstPhaseDefinition.class, false);
-                    ContainerUtil.addIfNotNull(results, fp);
+                    if (fp != null)
+                        results.add(Function.from(fp));
                 }
             });
         }
         
-        functionTreeChildren.put(macroName, results);
+        functionTreeChildren.put(functionName, results);
         return results;
     }
     
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
new file mode 100644
index 00000000000..fc84a195b99
--- /dev/null
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/Function.java
@@ -0,0 +1,38 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package ai.vespa.intellij.schema.model;
+
+import ai.vespa.intellij.schema.psi.SdFirstPhaseDefinition;
+import ai.vespa.intellij.schema.psi.SdFunctionDefinition;
+import ai.vespa.intellij.schema.psi.SdSecondPhaseDefinition;
+import com.intellij.psi.PsiElement;
+
+/**
+ * @author bratseth
+ */
+public class Function {
+
+    private final PsiElement definition;
+
+    public Function(PsiElement definition) {
+        this.definition = definition;
+    }
+
+    public PsiElement definition() { return definition; }
+
+    public static Function from(SdFirstPhaseDefinition firstPhase) {
+        return new Function(firstPhase);
+    }
+
+    public static Function from(SdSecondPhaseDefinition secondPhase) {
+        return new Function(secondPhase);
+    }
+
+    public static Function from(SdFunctionDefinition definition) {
+        return new Function(definition);
+    }
+
+    public static Function from(PsiElement definition) {
+        return new Function(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
new file mode 100644
index 00000000000..1cc6374f812
--- /dev/null
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/model/RankProfile.java
@@ -0,0 +1,19 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package ai.vespa.intellij.schema.model;
+
+import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
+
+/**
+ * A rank profile
+ *
+ * @author bratseth
+ */
+public class RankProfile {
+
+    private final SdRankProfileDefinition definition;
+
+    public RankProfile(SdRankProfileDefinition definition) {
+        this.definition = definition;
+    }
+
+}