4 files changed, 38 insertions, 35 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/IndexTreeBuilder.java b/src/main/java/cuchaz/enigma/analysis/IndexTreeBuilder.java
index 4ca7cd11..0c2dfd77 100644
--- a/src/main/java/cuchaz/enigma/analysis/IndexTreeBuilder.java
+++ b/src/main/java/cuchaz/enigma/analysis/IndexTreeBuilder.java
@@ -1,13 +1,14 @@
 package cuchaz.enigma.analysis;
 
 import com.google.common.collect.Lists;
-import cuchaz.enigma.analysis.index.EntryIndex;
 import cuchaz.enigma.analysis.index.JarIndex;
 import cuchaz.enigma.translation.Translator;
+import cuchaz.enigma.translation.mapping.EntryResolver;
 import cuchaz.enigma.translation.mapping.ResolutionStrategy;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 import cuchaz.enigma.translation.representation.entry.MethodEntry;
 
+import java.util.Collection;
 import java.util.List;
 
 public class IndexTreeBuilder {
@@ -52,34 +53,20 @@ public class IndexTreeBuilder {
                 );
 
                 // expand the full tree
-                rootNode.load(index, true);
+                rootNode.load(index);
 
                 return rootNode;
         }
 
         public List<MethodImplementationsTreeNode> buildMethodImplementations(Translator translator, MethodEntry obfMethodEntry) {
-                EntryIndex entryIndex = index.getEntryIndex();
-
-                List<MethodEntry> ancestorMethodEntries = Lists.newArrayList();
-
-                if (entryIndex.hasMethod(obfMethodEntry)) {
-                        ancestorMethodEntries.add(obfMethodEntry);
-                }
-
-                for (ClassEntry ancestorEntry : index.getInheritanceIndex().getAncestors(obfMethodEntry.getParent())) {
-                        MethodEntry ancestorMethod = obfMethodEntry.withParent(ancestorEntry);
-                        if (entryIndex.hasMethod(ancestorMethod)) {
-                                ancestorMethodEntries.add(ancestorMethod);
-                        }
-                }
+                EntryResolver resolver = index.getEntryResolver();
+                Collection<MethodEntry> resolvedEntries = resolver.resolveEntry(obfMethodEntry, ResolutionStrategy.RESOLVE_ROOT);
 
                 List<MethodImplementationsTreeNode> nodes = Lists.newArrayList();
-                if (!ancestorMethodEntries.isEmpty()) {
-                        for (MethodEntry interfaceMethodEntry : ancestorMethodEntries) {
-                                MethodImplementationsTreeNode node = new MethodImplementationsTreeNode(translator, interfaceMethodEntry);
-                                node.load(index);
-                                nodes.add(node);
-                        }
+                for (MethodEntry resolvedEntry : resolvedEntries) {
+                        MethodImplementationsTreeNode node = new MethodImplementationsTreeNode(translator, resolvedEntry);
+                        node.load(index);
+                        nodes.add(node);
                 }
 
                 return nodes;
diff --git a/src/main/java/cuchaz/enigma/analysis/MethodImplementationsTreeNode.java b/src/main/java/cuchaz/enigma/analysis/MethodImplementationsTreeNode.java
index e4b03043..b09f7ac6 100644
--- a/src/main/java/cuchaz/enigma/analysis/MethodImplementationsTreeNode.java
+++ b/src/main/java/cuchaz/enigma/analysis/MethodImplementationsTreeNode.java
@@ -71,8 +71,8 @@ public class MethodImplementationsTreeNode extends DefaultMutableTreeNode {
                 EntryIndex entryIndex = index.getEntryIndex();
                 InheritanceIndex inheritanceIndex = index.getInheritanceIndex();
 
-                Collection<ClassEntry> inheritors = inheritanceIndex.getChildren(entry.getParent());
-                for (ClassEntry inheritor : inheritors) {
+                Collection<ClassEntry> descendants = inheritanceIndex.getDescendants(entry.getParent());
+                for (ClassEntry inheritor : descendants) {
                         MethodEntry methodEntry = entry.withParent(inheritor);
                         if (entryIndex.hasMethod(methodEntry)) {
                                 nodes.add(new MethodImplementationsTreeNode(translator, methodEntry));
diff --git a/src/main/java/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java b/src/main/java/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java
index f0fd1d28..862bb92a 100644
--- a/src/main/java/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java
+++ b/src/main/java/cuchaz/enigma/analysis/MethodInheritanceTreeNode.java
@@ -11,7 +11,6 @@
 
 package cuchaz.enigma.analysis;
 
-import com.google.common.collect.Lists;
 import cuchaz.enigma.analysis.index.EntryIndex;
 import cuchaz.enigma.analysis.index.InheritanceIndex;
 import cuchaz.enigma.analysis.index.JarIndex;
@@ -20,7 +19,6 @@ import cuchaz.enigma.translation.representation.entry.ClassEntry;
 import cuchaz.enigma.translation.representation.entry.MethodEntry;
 
 import javax.swing.tree.DefaultMutableTreeNode;
-import java.util.List;
 
 public class MethodInheritanceTreeNode extends DefaultMutableTreeNode {
 
@@ -58,6 +56,10 @@ public class MethodInheritanceTreeNode extends DefaultMutableTreeNode {
                 return this.isImplemented;
         }
 
+        public boolean shouldExpand() {
+                return this.isImplemented || !this.isLeaf();
+        }
+
         @Override
         public String toString() {
                 MethodEntry translatedEntry = translator.translate(entry);
@@ -71,23 +73,19 @@ public class MethodInheritanceTreeNode extends DefaultMutableTreeNode {
                 }
         }
 
-        public void load(JarIndex index, boolean recurse) {
+        public void load(JarIndex index) {
                 // get all the child nodes
-                List<MethodInheritanceTreeNode> nodes = Lists.newArrayList();
                 EntryIndex entryIndex = index.getEntryIndex();
                 InheritanceIndex inheritanceIndex = index.getInheritanceIndex();
 
                 for (ClassEntry inheritorEntry : inheritanceIndex.getChildren(this.entry.getParent())) {
                         MethodEntry methodEntry = new MethodEntry(inheritorEntry, this.entry.getName(), this.entry.getDesc());
-                        nodes.add(new MethodInheritanceTreeNode(translator, methodEntry, entryIndex.hasMethod(methodEntry)));
-                }
 
-                // add them to this node
-                nodes.forEach(this::add);
+                        MethodInheritanceTreeNode node = new MethodInheritanceTreeNode(translator, methodEntry, entryIndex.hasMethod(methodEntry));
+                        node.load(index);
 
-                if (recurse) {
-                        for (MethodInheritanceTreeNode node : nodes) {
-                                node.load(index, true);
+                        if (node.shouldExpand()) {
+                                this.add(node);
                         }
                 }
         }
diff --git a/src/main/java/cuchaz/enigma/analysis/index/InheritanceIndex.java b/src/main/java/cuchaz/enigma/analysis/index/InheritanceIndex.java
index 1b8d9a8d..1ab2abdf 100644
--- a/src/main/java/cuchaz/enigma/analysis/index/InheritanceIndex.java
+++ b/src/main/java/cuchaz/enigma/analysis/index/InheritanceIndex.java
@@ -18,6 +18,7 @@ import cuchaz.enigma.translation.representation.entry.ClassDefEntry;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 
 import java.util.Collection;
+import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.Set;
 
@@ -60,6 +61,23 @@ public class InheritanceIndex implements JarIndexer {
                 return classChildren.get(classEntry);
         }
 
+        public Collection<ClassEntry> getDescendants(ClassEntry classEntry) {
+                Collection<ClassEntry> descendants = new HashSet<>();
+
+                LinkedList<ClassEntry> descendantQueue = new LinkedList<>();
+                descendantQueue.push(classEntry);
+
+                while (!descendantQueue.isEmpty()) {
+                        ClassEntry descendant = descendantQueue.pop();
+                        Collection<ClassEntry> children = getChildren(descendant);
+
+                        children.forEach(descendantQueue::push);
+                        descendants.addAll(children);
+                }
+
+                return descendants;
+        }
+
         public Set<ClassEntry> getAncestors(ClassEntry classEntry) {
                 Set<ClassEntry> ancestors = Sets.newHashSet();