1 files changed, 14 insertions, 16 deletions

diff --git a/src/main/java/cuchaz/enigma/analysis/ClassReferenceTreeNode.java b/src/main/java/cuchaz/enigma/analysis/ClassReferenceTreeNode.java
index ff5f2e9..90d8a6c 100644
--- a/src/main/java/cuchaz/enigma/analysis/ClassReferenceTreeNode.java
+++ b/src/main/java/cuchaz/enigma/analysis/ClassReferenceTreeNode.java
@@ -12,12 +12,12 @@
 package cuchaz.enigma.analysis;
 
 import com.google.common.collect.Sets;
-import cuchaz.enigma.bytecode.AccessFlags;
-import cuchaz.enigma.mapping.Translator;
-import cuchaz.enigma.mapping.entry.ClassEntry;
-import cuchaz.enigma.mapping.entry.Entry;
-import cuchaz.enigma.mapping.entry.MethodDefEntry;
-import cuchaz.enigma.mapping.entry.MethodEntry;
+import cuchaz.enigma.analysis.index.JarIndex;
+import cuchaz.enigma.analysis.index.ReferenceIndex;
+import cuchaz.enigma.translation.Translator;
+import cuchaz.enigma.translation.representation.entry.ClassEntry;
+import cuchaz.enigma.translation.representation.entry.Entry;
+import cuchaz.enigma.translation.representation.entry.MethodDefEntry;
 
 import javax.swing.tree.DefaultMutableTreeNode;
 import javax.swing.tree.TreeNode;
@@ -29,7 +29,6 @@ public class ClassReferenceTreeNode extends DefaultMutableTreeNode
         private Translator deobfuscatingTranslator;
         private ClassEntry entry;
         private EntryReference<ClassEntry, MethodDefEntry> reference;
-        private AccessFlags access;
 
         public ClassReferenceTreeNode(Translator deobfuscatingTranslator, ClassEntry entry) {
                 this.deobfuscatingTranslator = deobfuscatingTranslator;
@@ -37,12 +36,10 @@ public class ClassReferenceTreeNode extends DefaultMutableTreeNode
                 this.reference = null;
         }
 
-        public ClassReferenceTreeNode(Translator deobfuscatingTranslator,
-                                  EntryReference<ClassEntry, MethodDefEntry> reference, AccessFlags access) {
+        public ClassReferenceTreeNode(Translator deobfuscatingTranslator, EntryReference<ClassEntry, MethodDefEntry> reference) {
                 this.deobfuscatingTranslator = deobfuscatingTranslator;
                 this.entry = reference.entry;
                 this.reference = reference;
-                this.access = access;
         }
 
         @Override
@@ -58,16 +55,17 @@ public class ClassReferenceTreeNode extends DefaultMutableTreeNode
         @Override
         public String toString() {
                 if (this.reference != null) {
-                        return String.format("%s (%s)", this.deobfuscatingTranslator.getTranslatedMethodDef(this.reference.context),
-                                this.access);
+                        return String.format("%s", this.deobfuscatingTranslator.translate(this.reference.context));
                 }
-                return this.deobfuscatingTranslator.getTranslatedClass(this.entry).getName();
+                return this.deobfuscatingTranslator.translate(this.entry).getFullName();
         }
 
         public void load(JarIndex index, boolean recurse) {
+                ReferenceIndex referenceIndex = index.getReferenceIndex();
+
                 // get all the child nodes
-                for (EntryReference<ClassEntry, MethodDefEntry> reference : index.getMethodsReferencing(this.entry)) {
-                        add(new ClassReferenceTreeNode(this.deobfuscatingTranslator, reference, index.getAccessFlags(this.entry)));
+                for (EntryReference<ClassEntry, MethodDefEntry> reference : referenceIndex.getReferencesToClass(this.entry)) {
+                        add(new ClassReferenceTreeNode(this.deobfuscatingTranslator, reference));
                 }
 
                 if (recurse && this.children != null) {
@@ -76,7 +74,7 @@ public class ClassReferenceTreeNode extends DefaultMutableTreeNode
                                         ClassReferenceTreeNode node = (ClassReferenceTreeNode) child;
 
                                         // don't recurse into ancestor
-                                        Set<Entry> ancestors = Sets.newHashSet();
+                                        Set<Entry<?>> ancestors = Sets.newHashSet();
                                         TreeNode n = node;
                                         while (n.getParent() != null) {
                                                 n = n.getParent();