1 files changed, 63 insertions, 0 deletions

diff --git a/sd-plugin/src/main/java/org/intellij/sdk/language/hierarchy/SdCallerTreeStructure.java b/sd-plugin/src/main/java/org/intellij/sdk/language/hierarchy/SdCallerTreeStructure.java
new file mode 100644
index 00000000000..f6b21fc9ead
--- /dev/null
+++ b/sd-plugin/src/main/java/org/intellij/sdk/language/hierarchy/SdCallerTreeStructure.java
@@ -0,0 +1,63 @@
+package org.intellij.sdk.language.hierarchy;
+
+import com.intellij.openapi.progress.ProgressManager;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiReference;
+import com.intellij.psi.search.SearchScope;
+import com.intellij.psi.search.searches.ReferencesSearch;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.intellij.util.containers.ContainerUtil;
+import org.intellij.sdk.language.psi.SdFirstPhaseDefinition;
+import org.intellij.sdk.language.psi.SdFunctionDefinition;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.function.Consumer;
+
+public class SdCallerTreeStructure extends SdCallTreeStructure {
+    
+    private HashMap<String, HashSet<PsiElement>> macroTreeChildren;
+    
+    public SdCallerTreeStructure(Project project, PsiElement element, String currentScopeType) {
+        super(project, element, currentScopeType);
+        macroTreeChildren = new HashMap<>();
+    }
+    
+    @NotNull
+    @Override
+    protected HashSet<PsiElement> getChildren(@NotNull SdFunctionDefinition element) {
+        return getCallers(element, macrosMap);
+    }
+    
+    private HashSet<PsiElement> getCallers(@NotNull SdFunctionDefinition macro, @NotNull HashMap<String, List<PsiElement>> macrosMap) {
+        String macroName = macro.getName();
+
+        if (macroTreeChildren.containsKey(macroName)) {
+            return macroTreeChildren.get(macroName);
+        }
+        
+        HashSet<PsiElement> results = new HashSet<>();
+        
+        for (PsiElement macroImpl : macrosMap.get(macroName)) {
+            SearchScope searchScope = getSearchScope(myScopeType, macroImpl);
+            ReferencesSearch.search(macroImpl, 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().equals(macroName)) { 
+                    ContainerUtil.addIfNotNull(results, f); 
+                } else {
+                    SdFirstPhaseDefinition fp = PsiTreeUtil.getParentOfType(psiElement, SdFirstPhaseDefinition.class, false);
+                    ContainerUtil.addIfNotNull(results, fp);
+                }
+            });
+        }
+        
+        macroTreeChildren.put(macroName, results);
+        return results;
+    }
+    
+}
\ No newline at end of file