Return all profile parents

7 files changed, 111 insertions, 97 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 f1cbc026230..7c3f6dcb0ec 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
@@ -32,14 +32,20 @@ import ai.vespa.intellij.schema.psi.SdSummaryDefinition;
 import ai.vespa.intellij.schema.psi.SdTypes;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Optional;
+import java.util.stream.Collectors;
+
+import static org.codehaus.groovy.runtime.DefaultGroovyMethods.collect;
 
 /**
  * Util class for the plugin's code.
  *
  * @author Shahar Ariel
+ * @author bratseth
  */
 public class SdUtil {
     
@@ -57,24 +63,23 @@ public class SdUtil {
      * @param baseRankProfile the rank-profile node to find its parent
      * @return the rank-profile that the baseRankProfile inherits from, or null if it doesn't exist
      */
-    public static PsiElement getRankProfileParent(SdRankProfileDefinition baseRankProfile) {
-        if (baseRankProfile == null) {
-            return null;
-        }
+    public static List<SdRankProfileDefinition> getRankProfileParents(SdRankProfileDefinition baseRankProfile) {
+        if (baseRankProfile == null) return null;
         ASTNode inheritsNode = baseRankProfile.getNode().findChildByType(SdTypes.INHERITS);
-        if (inheritsNode == null) {
-            return null;
-        }
-        ASTNode ancestorAST = baseRankProfile.getNode().findChildByType(SdTypes.IDENTIFIER_VAL, inheritsNode);
-        if (ancestorAST == null) {
-            ancestorAST = baseRankProfile.getNode().findChildByType(SdTypes.IDENTIFIER_WITH_DASH_VAL, inheritsNode);
-            if (ancestorAST == null) {
-                return null;
-            }
-        }
-        SdIdentifier ancestorIdentifier = (SdIdentifier) ancestorAST.getPsi();
-        PsiReference ref = ancestorIdentifier.getReference();
-        return ref != null ? ref.resolve() : null;
+        if (inheritsNode == null) return null;
+
+        return identifiersIn(inheritsNode).stream()
+                                          .map(parentIdentifierAST -> parentIdentifierAST.getPsi().getReference())
+                                          .filter(reference -> reference != null)
+                                          .map(reference -> (SdRankProfileDefinition)reference.resolve())
+                                          .collect(Collectors.toList());
+    }
+
+    private static List<ASTNode> identifiersIn(ASTNode inheritsNode) {
+        return Arrays.stream(inheritsNode.getChildren(null))
+                     .filter(node -> node.getElementType() == SdTypes.IDENTIFIER_VAL ||
+                                     node.getElementType() == SdTypes.IDENTIFIER_WITH_DASH_VAL)
+                     .collect(Collectors.toList());
     }
     
     public static String createFunctionDescription(SdFunctionDefinition function) {
@@ -107,11 +112,10 @@ public class SdUtil {
         }
         return result;
     }
-    
-    
+
     public static List<SdDeclaration> findDeclarationsByScope(PsiElement file, PsiElement element, String name) {
         List<SdDeclaration> result = new ArrayList<>();
-        
+
         // If element is a field declared in another file (to be imported), return the declaration from the other file
         // if found, else return an empty result list
         if (element.getParent() instanceof SdImportFieldDefinition &&
@@ -159,18 +163,14 @@ public class SdUtil {
         
         // If element is a function's name, return the most specific declaration of the function
         if (((SdIdentifier) element).isFunctionName(file, name)) {
-            PsiElement curRankProfile = PsiTreeUtil.getParentOfType(element, SdRankProfileDefinition.class);
-            while (curRankProfile != null) {
-                for (SdFunctionDefinition function : PsiTreeUtil.collectElementsOfType(curRankProfile, SdFunctionDefinition.class)) {
-                    if (function.getName() != null && function.getName().equals(name)) {
-                        result.add(function);
-                        return result;
-                    }
-                }
-                curRankProfile = getRankProfileParent((SdRankProfileDefinition) curRankProfile);
+            var profile = (SdRankProfileDefinition)PsiTreeUtil.getParentOfType(element, SdRankProfileDefinition.class);
+            Optional<SdFunctionDefinition> function = findFunction(name, profile);
+            if (function.isPresent()) {
+                result.add(function.get());
+                return result;
             }
         }
-        
+
         for (PsiElement declaration : PsiTreeUtil.collectElements(file, psiElement ->
             psiElement instanceof SdDeclaration && !(psiElement instanceof SdArgumentDefinition))) {
             if (name.equals(((SdDeclaration) declaration).getName())) {
@@ -181,6 +181,25 @@ public class SdUtil {
         
         return result;
     }
+
+    /**
+     * Returns the first encountered function of the given name in the inheritance hierarchy
+     * of the given profile, or empty if it is not present.
+     *
+     * NOTE: Only profiles in the same file is considered
+     */
+    private static Optional<SdFunctionDefinition> findFunction(String functionName, SdRankProfileDefinition profile) {
+        Optional<SdFunctionDefinition> function = PsiTreeUtil.collectElementsOfType(profile, SdFunctionDefinition.class)
+                                                             .stream()
+                                                             .filter(f -> f.getName().equals(functionName))
+                                                             .findAny();
+        if (function.isPresent()) return function;
+        for (var parent : getRankProfileParents(profile)) {
+            function = findFunction(functionName, parent);
+            if (function.isPresent()) return function;
+        }
+        return Optional.empty();
+    }
     
     public static List<SdDeclaration> findDeclarations(PsiElement element) {
         return new ArrayList<>(PsiTreeUtil.collectElementsOfType(element, SdDeclaration.class));
diff --git a/integration/intellij/src/main/java/ai/vespa/intellij/schema/findUsages/SdRankProfileGroupingRule.java b/integration/intellij/src/main/java/ai/vespa/intellij/schema/findUsages/SdRankProfileGroupingRule.java
index a9e9f9337dd..140fd78be9c 100644
--- a/integration/intellij/src/main/java/ai/vespa/intellij/schema/findUsages/SdRankProfileGroupingRule.java
+++ b/integration/intellij/src/main/java/ai/vespa/intellij/schema/findUsages/SdRankProfileGroupingRule.java
@@ -25,7 +25,7 @@ public class SdRankProfileGroupingRule extends SingleParentUsageGroupingRule imp
         
         while (psiElement != null) {
             if (psiElement instanceof SdRankProfileDefinition) {
-                final SdRankProfileDefinition componentElement = (SdRankProfileDefinition) psiElement;
+                SdRankProfileDefinition componentElement = (SdRankProfileDefinition) psiElement;
                 return new SdUsageGroup(componentElement);
             }
             psiElement = psiElement.getParent();
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 524b9c31302..6d7c90bdbc1 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
@@ -17,6 +17,8 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 
 /**
  * A general tree in the "Call Hierarchy" window.
@@ -27,18 +29,18 @@ public abstract class SdCallTreeStructure extends HierarchyTreeStructure {
 
     protected final String myScopeType;
     protected final SdFile myFile;
-    protected HashMap<String, List<PsiElement>> macrosMap;
-    protected HashMap<String, HashSet<SdRankProfileDefinition>> ranksHeritageMap;
+    protected Map<String, List<PsiElement>> 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();
-        macrosMap = SdUtil.createFunctionsMap(myFile);
+        functionsMap = SdUtil.createFunctionsMap(myFile);
         ranksHeritageMap = new HashMap<>();
     }
     
-    protected abstract HashSet<PsiElement> getChildren(SdFunctionDefinition element);
+    protected abstract Set<PsiElement> getChildren(SdFunctionDefinition element);
     
     @Override
     protected Object[] buildChildren(HierarchyNodeDescriptor descriptor) {
@@ -55,9 +57,9 @@ public abstract class SdCallTreeStructure extends HierarchyTreeStructure {
                 return ArrayUtilRt.EMPTY_OBJECT_ARRAY;
             }
             
-            HashSet<PsiElement> children = getChildren((SdFunctionDefinition) element);
+            Set<PsiElement> children = getChildren((SdFunctionDefinition) element);
             
-            final HashMap<PsiElement, SdCallHierarchyNodeDescriptor> callerToDescriptorMap = new HashMap<>();
+            Map<PsiElement, SdCallHierarchyNodeDescriptor> callerToDescriptorMap = new HashMap<>();
             PsiElement baseClass = PsiTreeUtil.getParentOfType(element, SdRankProfileDefinition.class);
             
             for (PsiElement caller : children) {
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 59e6ca7add3..6783d47ba31 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
@@ -14,6 +14,8 @@ import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 
 /**
  * A Callee tree in the "Call Hierarchy" window.
@@ -27,18 +29,18 @@ public class SdCalleeTreeStructure extends SdCallTreeStructure {
     }
         
     @Override
-    protected HashSet<PsiElement> getChildren(SdFunctionDefinition element) {
-        return getCallees(element, macrosMap);
+    protected Set<PsiElement> getChildren(SdFunctionDefinition element) {
+        return getCallees(element, functionsMap);
     }
     
-    private HashSet<PsiElement> getCallees(SdFunctionDefinition macro, HashMap<String, List<PsiElement>> macrosMap) {
-        final HashSet<PsiElement> results = new HashSet<>();
-        SdExpressionDefinition expression = PsiTreeUtil.findChildOfType(macro, SdExpressionDefinition.class);
+    private Set<PsiElement> getCallees(SdFunctionDefinition function, Map<String, List<PsiElement>> functions) {
+        Set<PsiElement> results = new HashSet<>();
+        SdExpressionDefinition expression = PsiTreeUtil.findChildOfType(function, SdExpressionDefinition.class);
         if (expression == null) {
             return results;
         }
         for (SdIdentifier identifier : PsiTreeUtil.collectElementsOfType(expression, SdIdentifier.class)) {
-            if (macrosMap.containsKey(((PsiNamedElement) identifier).getName())) {
+            if (functions.containsKey(((PsiNamedElement) identifier).getName())) {
                 PsiReference identifierRef = identifier.getReference();
                 if (identifierRef != null) {
                     results.add(identifierRef.resolve());
@@ -46,7 +48,7 @@ public class SdCalleeTreeStructure extends SdCallTreeStructure {
             }
         }
         
-        SdRankProfileDefinition rankProfile = PsiTreeUtil.getParentOfType(macro, SdRankProfileDefinition.class);
+        SdRankProfileDefinition rankProfile = PsiTreeUtil.getParentOfType(function, SdRankProfileDefinition.class);
         if (rankProfile == null) {
             return results;
         }
@@ -55,11 +57,11 @@ public class SdCalleeTreeStructure extends SdCallTreeStructure {
             ranksHeritageMap.put(rankProfileName, SdHierarchyUtil.getRankProfileChildren(myFile, rankProfile));
         }
         
-        HashSet<SdRankProfileDefinition> inheritedRanks = ranksHeritageMap.get(rankProfileName);
+        Set<SdRankProfileDefinition> inheritedRanks = ranksHeritageMap.get(rankProfileName);
         
-        for (PsiElement macroImpl : macrosMap.get(macro.getName())) {
-            if (inheritedRanks.contains(PsiTreeUtil.getParentOfType(macroImpl, SdRankProfileDefinition.class))) {
-                results.add(macroImpl);
+        for (PsiElement functionImpl : functions.get(function.getName())) {
+            if (inheritedRanks.contains(PsiTreeUtil.getParentOfType(functionImpl, 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 f442c4ea5bc..7b0a7bf7955 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
@@ -15,6 +15,8 @@ import ai.vespa.intellij.schema.psi.SdFunctionDefinition;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.function.Consumer;
 
 /**
@@ -24,26 +26,26 @@ import java.util.function.Consumer;
  */
 public class SdCallerTreeStructure extends SdCallTreeStructure {
     
-    private HashMap<String, HashSet<PsiElement>> macroTreeChildren;
+    private Map<String, Set<PsiElement>> functionTreeChildren;
     
     public SdCallerTreeStructure(Project project, PsiElement element, String currentScopeType) {
         super(project, element, currentScopeType);
-        macroTreeChildren = new HashMap<>();
+        functionTreeChildren = new HashMap<>();
     }
     
     @Override
-    protected HashSet<PsiElement> getChildren(SdFunctionDefinition element) {
-        return getCallers(element, macrosMap);
+    protected Set<PsiElement> getChildren(SdFunctionDefinition element) {
+        return getCallers(element, functionsMap);
     }
     
-    private HashSet<PsiElement> getCallers(SdFunctionDefinition macro, HashMap<String, List<PsiElement>> macrosMap) {
+    private Set<PsiElement> getCallers(SdFunctionDefinition macro, Map<String, List<PsiElement>> macrosMap) {
         String macroName = macro.getName();
 
-        if (macroTreeChildren.containsKey(macroName)) {
-            return macroTreeChildren.get(macroName);
+        if (functionTreeChildren.containsKey(macroName)) {
+            return functionTreeChildren.get(macroName);
         }
         
-        HashSet<PsiElement> results = new HashSet<>();
+        Set<PsiElement> results = new HashSet<>();
         
         for (PsiElement macroImpl : macrosMap.get(macroName)) {
             SearchScope searchScope = getSearchScope(myScopeType, macroImpl);
@@ -60,7 +62,7 @@ public class SdCallerTreeStructure extends SdCallTreeStructure {
             });
         }
         
-        macroTreeChildren.put(macroName, results);
+        functionTreeChildren.put(macroName, results);
         return results;
     }
     
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 dc925f0f369..7da47607c0c 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
@@ -14,12 +14,14 @@ import ai.vespa.intellij.schema.psi.SdFunctionDefinition;
 import ai.vespa.intellij.schema.psi.SdRankProfileDefinition;
 
 import java.util.Comparator;
-import java.util.HashSet;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 /**
  * Call Hierarchy feature utilities.
  *
  * @author Shahar Ariel
+ * @author bratseth
  */
 public class SdHierarchyUtil {
     
@@ -31,35 +33,21 @@ public class SdHierarchyUtil {
     public static boolean isExecutable(PsiElement component) {
         return component instanceof SdFunctionDefinition;
     }
-    
-    public static HashSet<SdRankProfileDefinition> getRankProfileChildren(SdFile file, SdRankProfileDefinition rankProfileTarget) {
-        HashSet<SdRankProfileDefinition> result = new HashSet<>();
-        HashSet<SdRankProfileDefinition> notResult = new HashSet<>();
 
-        for (SdRankProfileDefinition rank : PsiTreeUtil.collectElementsOfType(file, SdRankProfileDefinition.class)) {
-            if (notResult.contains(rank)) {
-                continue;
-            }
-            HashSet<SdRankProfileDefinition> tempRanks = new HashSet<>();
-            SdRankProfileDefinition curRank = rank;
-            while (curRank != null) {
-                String curRankName = curRank.getName();
-                if (curRankName != null && curRankName.equals(rankProfileTarget.getName())) {
-                    result.addAll(tempRanks);
-                    break;
-                }
-                tempRanks.add(curRank);
-                PsiElement temp = SdUtil.getRankProfileParent(curRank);
-                curRank = temp != null ? (SdRankProfileDefinition) temp : null;
-            }
-            if (curRank == null) {
-                notResult.addAll(tempRanks);
-            }
-        }
-        return result;
+    public static Set<SdRankProfileDefinition> getRankProfileChildren(SdFile file, SdRankProfileDefinition targetProfile) {
+        return PsiTreeUtil.collectElementsOfType(file, SdRankProfileDefinition.class)
+                          .stream()
+                          .filter(profile -> isChildOf(targetProfile, profile))
+                          .collect(Collectors.toSet());
     }
-    
-    
+
+    private static boolean isChildOf(SdRankProfileDefinition targetProfile, SdRankProfileDefinition thisProfile) {
+        if (thisProfile.getName().equals(targetProfile.getName())) return true;
+        return SdUtil.getRankProfileParents(thisProfile)
+                     .stream()
+                     .anyMatch(parent -> isChildOf(targetProfile, parent));
+    }
+
     public static Comparator<NodeDescriptor<?>> getComparator(Project project) {
         final HierarchyBrowserManager.State state = HierarchyBrowserManager.getInstance(project).getState();
         if (state != null && state.SORT_ALPHABETICALLY) {
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 e7fa94cf81c..f60131ef8cb 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
@@ -11,19 +11,20 @@ import ai.vespa.intellij.schema.SdUtil;
 public interface SdFunctionDefinitionInterface extends SdDeclaration {
 
     default boolean isOverride() {
-        String macroName = this.getName();
-        
-        SdRankProfileDefinition curRankProfile = PsiTreeUtil.getParentOfType(this, SdRankProfileDefinition.class);
-        if (curRankProfile != null) {
-            curRankProfile = (SdRankProfileDefinition) SdUtil.getRankProfileParent(curRankProfile);
+        String functionName = this.getName();
+        SdRankProfileDefinition thisRankProfile = PsiTreeUtil.getParentOfType(this, SdRankProfileDefinition.class);
+        if (thisRankProfile == null) return false;
+        for (var parentProfile : SdUtil.getRankProfileParents(thisRankProfile)) {
+            if (containsFunction(functionName, parentProfile))
+                return true;
         }
-        while (curRankProfile != null) {
-            for (SdFunctionDefinition macro : PsiTreeUtil.collectElementsOfType(curRankProfile, SdFunctionDefinition.class)) {
-                if (macro.getName() != null && macro.getName().equals(macroName)) {
-                    return true;
-                }
-            }
-            curRankProfile = (SdRankProfileDefinition) SdUtil.getRankProfileParent(curRankProfile);
+        return false;
+    }
+
+    default boolean containsFunction(String functionName, SdRankProfileDefinition rankProfile) {
+        for (var parentProfile : SdUtil.getRankProfileParents(rankProfile)) {
+            if (containsFunction(functionName, parentProfile))
+                return true;
         }
         return false;
     }