1 files changed, 16 insertions, 19 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/ClassImplementationsTreeNode.java b/src/main/java/cuchaz/enigma/analysis/ClassImplementationsTreeNode.java
index e876bb0..0fc44ca 100644
--- a/src/main/java/cuchaz/enigma/analysis/ClassImplementationsTreeNode.java
+++ b/src/main/java/cuchaz/enigma/analysis/ClassImplementationsTreeNode.java
@@ -12,26 +12,28 @@
 package cuchaz.enigma.analysis;
 
 import com.google.common.collect.Lists;
-import cuchaz.enigma.mapping.entry.ClassEntry;
-import cuchaz.enigma.mapping.entry.MethodEntry;
-import cuchaz.enigma.mapping.Translator;
+import cuchaz.enigma.analysis.index.InheritanceIndex;
+import cuchaz.enigma.analysis.index.JarIndex;
+import cuchaz.enigma.translation.Translator;
+import cuchaz.enigma.translation.representation.entry.ClassEntry;
+import cuchaz.enigma.translation.representation.entry.MethodEntry;
 
 import javax.swing.tree.DefaultMutableTreeNode;
+import java.util.Collection;
 import java.util.List;
 
 public class ClassImplementationsTreeNode extends DefaultMutableTreeNode {
-
-        private final Translator deobfuscatingTranslator;
+        private final Translator translator;
         private final ClassEntry entry;
 
-        public ClassImplementationsTreeNode(Translator deobfuscatingTranslator, ClassEntry entry) {
-                this.deobfuscatingTranslator = deobfuscatingTranslator;
+        public ClassImplementationsTreeNode(Translator translator, ClassEntry entry) {
+                this.translator = translator;
                 this.entry = entry;
         }
 
         public static ClassImplementationsTreeNode findNode(ClassImplementationsTreeNode node, MethodEntry entry) {
                 // is this the node?
-                if (node.entry.equals(entry.getOwnerClassEntry())) {
+                if (node.entry.equals(entry.getParent())) {
                         return node;
                 }
 
@@ -49,24 +51,19 @@ public class ClassImplementationsTreeNode extends DefaultMutableTreeNode {
                 return this.entry;
         }
 
-        public String getDeobfClassName() {
-                return this.deobfuscatingTranslator.getTranslatedClass(entry).getClassName();
-        }
-
         @Override
         public String toString() {
-                String className = getDeobfClassName();
-                if (className == null) {
-                        className = this.entry.getClassName();
-                }
-                return className;
+                return translator.translate(entry).toString();
         }
 
         public void load(JarIndex index) {
                 // get all method implementations
                 List<ClassImplementationsTreeNode> nodes = Lists.newArrayList();
-                for (String implementingClassName : index.getImplementingClasses(this.entry.getClassName())) {
-                        nodes.add(new ClassImplementationsTreeNode(this.deobfuscatingTranslator, new ClassEntry(implementingClassName)));
+                InheritanceIndex inheritanceIndex = index.getInheritanceIndex();
+
+                Collection<ClassEntry> inheritors = inheritanceIndex.getChildren(entry);
+                for (ClassEntry inheritor : inheritors) {
+                        nodes.add(new ClassImplementationsTreeNode(translator, inheritor));
                 }
 
                 // add them to this node