Structure panel!

4 files changed, 177 insertions, 24 deletions

diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java
index c56731dc..a3a438ee 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java
@@ -51,10 +51,7 @@ import cuchaz.enigma.network.packet.RemoveMappingC2SPacket;
 import cuchaz.enigma.network.packet.RenameC2SPacket;
 import cuchaz.enigma.source.Token;
 import cuchaz.enigma.translation.mapping.EntryRemapper;
-import cuchaz.enigma.translation.representation.entry.ClassEntry;
-import cuchaz.enigma.translation.representation.entry.Entry;
-import cuchaz.enigma.translation.representation.entry.FieldEntry;
-import cuchaz.enigma.translation.representation.entry.MethodEntry;
+import cuchaz.enigma.translation.representation.entry.*;
 import cuchaz.enigma.utils.I18n;
 import cuchaz.enigma.utils.validation.ParameterizedMessage;
 import cuchaz.enigma.utils.validation.ValidationContext;
@@ -82,6 +79,7 @@ public class Gui implements LanguageChangeListener {
         private JPanel classesPanel;
         private JSplitPane splitClasses;
         private IdentifierPanel infoPanel;
+        private StructurePanel structurePanel;
         private JTree inheritanceTree;
         private JTree implementationsTree;
         private JTree callsTree;
@@ -163,6 +161,9 @@ public class Gui implements LanguageChangeListener {
                 // init info panel
                 infoPanel = new IdentifierPanel(this);
 
+                // init structure panel
+                this.structurePanel = new StructurePanel(this);
+
                 // init inheritance panel
                 inheritanceTree = new JTree();
                 inheritanceTree.setModel(null);
@@ -287,6 +288,8 @@ public class Gui implements LanguageChangeListener {
                                                 editorTabPopupMenu.show(openFiles, e.getX(), e.getY(), EditorPanel.byUi(openFiles.getComponentAt(i)));
                                         }
                                 }
+
+                                showStructure(getActiveEditor());
                         }
                 });
 
@@ -311,6 +314,7 @@ public class Gui implements LanguageChangeListener {
                 centerPanel.add(openFiles, BorderLayout.CENTER);
                 tabs = new JTabbedPane();
                 tabs.setPreferredSize(ScaleUtil.getDimension(250, 0));
+                tabs.addTab(I18n.translate("info_panel.tree.structure"), structurePanel);
                 tabs.addTab(I18n.translate("info_panel.tree.inheritance"), inheritancePanel);
                 tabs.addTab(I18n.translate("info_panel.tree.implementations"), implementationsPanel);
                 tabs.addTab(I18n.translate("info_panel.tree.calls"), callPanel);
@@ -482,11 +486,14 @@ public class Gui implements LanguageChangeListener {
                                 }
                         });
 
+                        showStructure(ed);
+
                         return ed;
                 });
                 if (editorPanel != null) {
                         openFiles.setSelectedComponent(editors.get(entry).getUi());
                 }
+
                 return editorPanel;
         }
 
@@ -506,6 +513,7 @@ public class Gui implements LanguageChangeListener {
         public void closeEditor(EditorPanel ed) {
                 openFiles.remove(ed.getUi());
                 editors.inverse().remove(ed);
+                showStructure(getActiveEditor());
                 ed.destroy();
         }
 
@@ -595,6 +603,32 @@ public class Gui implements LanguageChangeListener {
                 infoPanel.startRenaming();
         }
 
+        public void showStructure(EditorPanel editor) {
+                JTree structureTree = this.structurePanel.getStructureTree();
+                structureTree.setModel(null);
+
+                if (editor == null) {
+                        this.structurePanel.getSortingPanel().setVisible(false);
+                        return;
+                }
+
+                ClassEntry classEntry = editor.getClassHandle().getRef();
+                if (classEntry == null) return;
+
+                this.structurePanel.getSortingPanel().setVisible(true);
+
+                // get the class structure
+                StructureTreeNode node = this.controller.getClassStructure(classEntry, this.structurePanel.shouldHideDeobfuscated());
+
+                // show the tree at the root
+                TreePath path = getPathToRoot(node);
+                structureTree.setModel(new DefaultTreeModel((TreeNode) path.getPathComponent(0)));
+                structureTree.expandPath(path);
+                structureTree.setSelectionRow(structureTree.getRowForPath(path));
+
+                redraw();
+        }
+
         public void showInheritance(EditorPanel editor) {
                 EntryReference<Entry<?>, Entry<?>> cursorReference = editor.getCursorReference();
                 if (cursorReference == null) return;
@@ -621,7 +655,7 @@ public class Gui implements LanguageChangeListener {
                         inheritanceTree.setSelectionRow(inheritanceTree.getRowForPath(path));
                 }
 
-                tabs.setSelectedIndex(0);
+                tabs.setSelectedIndex(1);
 
                 redraw();
         }
@@ -649,7 +683,7 @@ public class Gui implements LanguageChangeListener {
                         implementationsTree.setSelectionRow(implementationsTree.getRowForPath(path));
                 }
 
-                tabs.setSelectedIndex(1);
+                tabs.setSelectedIndex(2);
 
                 redraw();
         }
@@ -669,7 +703,7 @@ public class Gui implements LanguageChangeListener {
                         callsTree.setModel(new DefaultTreeModel(node));
                 }
 
-                tabs.setSelectedIndex(2);
+                tabs.setSelectedIndex(3);
 
                 redraw();
         }
@@ -893,9 +927,10 @@ public class Gui implements LanguageChangeListener {
         public void retranslateUi() {
                 this.jarFileChooser.setTitle(I18n.translate("menu.file.jar.open"));
                 this.exportJarFileChooser.setTitle(I18n.translate("menu.file.export.jar"));
-                this.tabs.setTitleAt(0, I18n.translate("info_panel.tree.inheritance"));
-                this.tabs.setTitleAt(1, I18n.translate("info_panel.tree.implementations"));
-                this.tabs.setTitleAt(2, I18n.translate("info_panel.tree.calls"));
+                this.tabs.setTitleAt(0, I18n.translate("info_panel.tree.structure"));
+                this.tabs.setTitleAt(1, I18n.translate("info_panel.tree.inheritance"));
+                this.tabs.setTitleAt(2, I18n.translate("info_panel.tree.implementations"));
+                this.tabs.setTitleAt(3, I18n.translate("info_panel.tree.calls"));
                 this.logTabs.setTitleAt(0, I18n.translate("log_panel.users"));
                 this.logTabs.setTitleAt(1, I18n.translate("log_panel.messages"));
                 this.connectionStatusLabel.setText(I18n.translate(connectionState == ConnectionState.NOT_CONNECTED ? "status.disconnected" : "status.connected"));
@@ -907,6 +942,7 @@ public class Gui implements LanguageChangeListener {
                 this.deobfPanel.retranslateUi();
                 this.deobfPanelPopupMenu.retranslateUi();
                 this.infoPanel.retranslateUi();
+                this.structurePanel.retranslateUi();
                 this.editors.values().forEach(EditorPanel::retranslateUi);
         }
 
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 4f7819e4..4e964dae 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
@@ -57,10 +57,7 @@ import cuchaz.enigma.translation.mapping.serde.MappingParseException;
 import cuchaz.enigma.translation.mapping.serde.MappingSaveParameters;
 import cuchaz.enigma.translation.mapping.tree.EntryTree;
 import cuchaz.enigma.translation.mapping.tree.HashEntryTree;
-import cuchaz.enigma.translation.representation.entry.ClassEntry;
-import cuchaz.enigma.translation.representation.entry.Entry;
-import cuchaz.enigma.translation.representation.entry.FieldEntry;
-import cuchaz.enigma.translation.representation.entry.MethodEntry;
+import cuchaz.enigma.translation.representation.entry.*;
 import cuchaz.enigma.utils.I18n;
 import cuchaz.enigma.utils.Utils;
 import cuchaz.enigma.utils.validation.ValidationContext;
@@ -395,6 +392,13 @@ public class GuiController implements ClientPacketHandler {
                 chp.invalidateMapped();
         }
 
+        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);
+                return rootNode;
+        }
+
         public ClassInheritanceTreeNode getClassInheritance(ClassEntry entry) {
                 Translator translator = project.getMapper().getDeobfuscator();
                 ClassInheritanceTreeNode rootNode = indexTreeBuilder.buildClassInheritance(translator, entry);
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/panels/StructurePanel.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/panels/StructurePanel.java
new file mode 100644
index 00000000..32f803f8
--- /dev/null
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/panels/StructurePanel.java
@@ -0,0 +1,102 @@
+package cuchaz.enigma.gui.panels;
+
+import cuchaz.enigma.analysis.StructureTreeNode;
+import cuchaz.enigma.gui.Gui;
+import cuchaz.enigma.gui.util.GuiUtil;
+import cuchaz.enigma.translation.representation.entry.ClassEntry;
+import cuchaz.enigma.translation.representation.entry.FieldEntry;
+import cuchaz.enigma.translation.representation.entry.MethodEntry;
+import cuchaz.enigma.translation.representation.entry.ParentedEntry;
+import cuchaz.enigma.utils.I18n;
+
+import javax.swing.*;
+import javax.swing.tree.TreeCellRenderer;
+import javax.swing.tree.TreePath;
+import java.awt.*;
+import java.awt.event.MouseAdapter;
+import java.awt.event.MouseEvent;
+
+public class StructurePanel extends JPanel {
+    private final Gui gui;
+
+    private JPanel sortingPanel;
+    private JCheckBox hideDeobfuscated;
+
+    private JTree structureTree;
+
+    public StructurePanel(Gui gui) {
+        this.gui = gui;
+
+        this.sortingPanel = new JPanel();
+        this.hideDeobfuscated = new JCheckBox(I18n.translate("info_panel.tree.structure.hide_deobfuscated"));
+        this.hideDeobfuscated.addActionListener(event -> gui.showStructure(gui.getActiveEditor()));
+        this.sortingPanel.add(this.hideDeobfuscated);
+        this.sortingPanel.setVisible(false);
+
+        this.structureTree = new JTree();
+        this.structureTree.setModel(null);
+        this.structureTree.setCellRenderer(new StructureTreeCellRenderer());
+        this.structureTree.addMouseListener(new MouseAdapter() {
+            @Override
+            public void mouseClicked(MouseEvent event) {
+                if (event.getClickCount() >= 2) {
+                    // get the selected node
+                    TreePath path = structureTree.getSelectionPath();
+                    if (path == null) {
+                        return;
+                    }
+
+                    Object node = path.getLastPathComponent();
+                    if (node instanceof StructureTreeNode) {
+                        gui.getController().navigateTo(((StructureTreeNode) node).getEntry());
+                    }
+                }
+            }
+        });
+
+        this.setLayout(new BorderLayout());
+        this.add(this.sortingPanel, BorderLayout.NORTH);
+        this.add(new JScrollPane(this.structureTree));
+    }
+
+    public JPanel getSortingPanel() {
+        return this.sortingPanel;
+    }
+
+    public boolean shouldHideDeobfuscated() {
+        return this.hideDeobfuscated.isSelected();
+    }
+
+    public JTree getStructureTree() {
+        return this.structureTree;
+    }
+
+    public void retranslateUi() {
+        this.hideDeobfuscated.setText(I18n.translate("info_panel.tree.structure.hide_deobfuscated"));
+    }
+
+    class StructureTreeCellRenderer implements TreeCellRenderer {
+        private JLabel label;
+
+        public StructureTreeCellRenderer() {
+            this.label = new JLabel();
+        }
+
+        @Override
+        public Component getTreeCellRendererComponent(JTree tree, Object value, boolean selected, boolean expanded, boolean leaf, int row, boolean hasFocus) {
+            ParentedEntry entry = ((StructureTreeNode) value).getEntry();
+
+            if (entry instanceof ClassEntry) {
+                this.label.setIcon(GuiUtil.CLASS_ICON);
+            } else if (entry instanceof MethodEntry) {
+                this.label.setIcon(GuiUtil.METHOD_ICON);
+            } else if (entry instanceof FieldEntry) {
+                this.label.setIcon(GuiUtil.FIELD_ICON);
+            }
+
+            this.label.setText(value.toString());
+
+            return this.label;
+        }
+    }
+}
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/util/GuiUtil.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/util/GuiUtil.java
index 7fe942d0..666fc0aa 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/util/GuiUtil.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/util/GuiUtil.java
@@ -1,24 +1,24 @@
 package cuchaz.enigma.gui.util;
 
-import java.awt.Color;
-import java.awt.Cursor;
-import java.awt.Desktop;
-import java.awt.Font;
+import cuchaz.enigma.utils.Os;
+
+import javax.imageio.ImageIO;
+import javax.swing.*;
+import java.awt.*;
 import java.awt.event.MouseAdapter;
 import java.awt.event.MouseEvent;
 import java.awt.font.TextAttribute;
 import java.io.IOException;
+import java.io.InputStream;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.Map;
 
-import javax.swing.JComponent;
-import javax.swing.JLabel;
-import javax.swing.ToolTipManager;
-
-import cuchaz.enigma.utils.Os;
-
 public class GuiUtil {
+    public static final Icon CLASS_ICON = loadIcon("class");
+    public static final Icon METHOD_ICON = loadIcon("method");
+    public static final Icon FIELD_ICON = loadIcon("field");
+
     public static void openUrl(String url) {
         try {
             switch (Os.getOs()) {
@@ -70,4 +70,15 @@ public class GuiUtil {
         return link;
     }
 
+    public static Icon loadIcon(String name) {
+        try {
+            InputStream inputStream = GuiUtil.class.getResourceAsStream("/icons/" + name + ".png");
+            Image image = ImageIO.read(inputStream).getScaledInstance(ScaleUtil.scale(16), ScaleUtil.scale(16), Image.SCALE_DEFAULT);
+            return new ImageIcon(image);
+        } catch (IOException e) {
+            e.printStackTrace();
+        }
+
+        return null;
+    }
 }