2 files changed, 53 insertions, 15 deletions

diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
index 4e964dae..656f5a43 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -393,9 +393,8 @@ public class GuiController implements ClientPacketHandler {
         }
 
         public StructureTreeNode getClassStructure(ClassEntry entry, boolean hideDeobfuscated) {
-                Translator translator = this.project.getMapper().getDeobfuscator();
-                StructureTreeNode rootNode = new StructureTreeNode(translator, entry, entry);
-                rootNode.load(this.project.getJarIndex(), hideDeobfuscated);
+                StructureTreeNode rootNode = new StructureTreeNode(this.project, entry, entry);
+                rootNode.load(this.project, hideDeobfuscated);
                 return rootNode;
         }
 
diff --git a/enigma/src/main/java/cuchaz/enigma/analysis/StructureTreeNode.java b/enigma/src/main/java/cuchaz/enigma/analysis/StructureTreeNode.java
index 13f277c7..312c2fec 100644
--- a/enigma/src/main/java/cuchaz/enigma/analysis/StructureTreeNode.java
+++ b/enigma/src/main/java/cuchaz/enigma/analysis/StructureTreeNode.java
@@ -1,7 +1,9 @@
 package cuchaz.enigma.analysis;
 
-import cuchaz.enigma.analysis.index.JarIndex;
-import cuchaz.enigma.translation.Translator;
+import cuchaz.enigma.EnigmaProject;
+import cuchaz.enigma.api.service.NameProposalService;
+import cuchaz.enigma.api.service.ObfuscationTestService;
+import cuchaz.enigma.translation.mapping.EntryRemapper;
 import cuchaz.enigma.translation.representation.TypeDescriptor;
 import cuchaz.enigma.translation.representation.entry.*;
 
@@ -9,12 +11,14 @@ import javax.swing.tree.DefaultMutableTreeNode;
 import java.util.List;
 
 public class StructureTreeNode extends DefaultMutableTreeNode {
-    private final Translator translator;
+    private final List<NameProposalService> nameProposalServices;
+    private final EntryRemapper mapper;
     private final ClassEntry parentEntry;
     private final ParentedEntry entry;
 
-    public StructureTreeNode(Translator translator, ClassEntry parentEntry, ParentedEntry entry) {
-        this.translator = translator;
+    public StructureTreeNode(EnigmaProject project, ClassEntry parentEntry, ParentedEntry entry) {
+        this.nameProposalServices = project.getEnigma().getServices().get(NameProposalService.TYPE);
+        this.mapper = project.getMapper();
         this.parentEntry = parentEntry;
         this.entry = entry;
     }
@@ -26,19 +30,19 @@ public class StructureTreeNode extends DefaultMutableTreeNode {
         return this.entry;
     }
 
-    public void load(JarIndex jarIndex, boolean hideDeobfuscated) {
-        List<ParentedEntry> children = jarIndex.getChildrenByClass().get(this.parentEntry);
+    public void load(EnigmaProject project, boolean hideDeobfuscated) {
+        List<ParentedEntry> children = project.getJarIndex().getChildrenByClass().get(this.parentEntry);
 
         for (ParentedEntry child : children) {
-            StructureTreeNode childNode = new StructureTreeNode(this.translator, this.parentEntry, child);
+            StructureTreeNode childNode = new StructureTreeNode(project, this.parentEntry, child);
 
             if (child instanceof ClassEntry) {
-                childNode = new StructureTreeNode(this.translator, (ClassEntry) child, child);
-                childNode.load(jarIndex, hideDeobfuscated);
+                childNode = new StructureTreeNode(project, (ClassEntry) child, child);
+                childNode.load(project, hideDeobfuscated);
             }
 
             // don't add deobfuscated members if hideDeobfuscated is true, unless it's an inner class
-            if (hideDeobfuscated && this.translator.extendedTranslate(child).isDeobfuscated() && !(child instanceof ClassEntry)) {
+            if (hideDeobfuscated && this.isDeobfuscated(project, child) && !(child instanceof ClassEntry)) {
                 continue;
             }
 
@@ -51,11 +55,46 @@ public class StructureTreeNode extends DefaultMutableTreeNode {
         }
     }
 
+    private boolean isDeobfuscated(EnigmaProject project, ParentedEntry child) {
+        List<ObfuscationTestService> obfuscationTestServices = project.getEnigma().getServices().get(ObfuscationTestService.TYPE);
+
+        if (!obfuscationTestServices.isEmpty()) {
+            for (ObfuscationTestService service : obfuscationTestServices) {
+                if (service.testDeobfuscated(child)) {
+                    return true;
+                }
+            }
+        }
+
+        if (!this.nameProposalServices.isEmpty()) {
+            for (NameProposalService service : this.nameProposalServices) {
+                if (service.proposeName(child, this.mapper).isPresent()) {
+                    return true;
+                }
+            }
+        }
+
+        String mappedName = project.getMapper().deobfuscate(child).getName();
+        if (mappedName != null && !mappedName.isEmpty() && !mappedName.equals(child.getName())) {
+            return true;
+        }
+
+        return false;
+    }
+
     @Override
     public String toString() {
-        ParentedEntry translatedEntry = this.translator.extendedTranslate(this.entry).getValue();
+        ParentedEntry translatedEntry = this.mapper.deobfuscate(this.entry);
         String result = translatedEntry.getName();
 
+        if (!this.nameProposalServices.isEmpty()) {
+            for (NameProposalService service : this.nameProposalServices) {
+                if (service.proposeName(this.entry, this.mapper).isPresent()) {
+                    result = service.proposeName(this.entry, this.mapper).get();
+                }
+            }
+        }
+
         if (this.entry instanceof FieldDefEntry) {
             FieldDefEntry field = (FieldDefEntry) translatedEntry;
             String returnType = this.parseDesc(field.getDesc());