Merge pull request #379 from YanisBft/entry-and-search

The Search update

5 files changed, 110 insertions, 36 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 ce4823e2..aaa04460 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java
@@ -47,10 +47,7 @@ import cuchaz.enigma.gui.renderer.CallsTreeCellRenderer;
 import cuchaz.enigma.gui.renderer.ImplementationsTreeCellRenderer;
 import cuchaz.enigma.gui.renderer.InheritanceTreeCellRenderer;
 import cuchaz.enigma.gui.renderer.MessageListCellRenderer;
-import cuchaz.enigma.gui.util.History;
-import cuchaz.enigma.gui.util.LanguageChangeListener;
-import cuchaz.enigma.gui.util.LanguageUtil;
-import cuchaz.enigma.gui.util.ScaleUtil;
+import cuchaz.enigma.gui.util.*;
 import cuchaz.enigma.network.Message;
 import cuchaz.enigma.network.packet.MarkDeobfuscatedC2SPacket;
 import cuchaz.enigma.network.packet.MessageC2SPacket;
@@ -172,6 +169,7 @@ public class Gui implements LanguageChangeListener {
                 inheritanceTree = new JTree();
                 inheritanceTree.setModel(null);
                 inheritanceTree.setCellRenderer(new InheritanceTreeCellRenderer(this));
+                inheritanceTree.setSelectionModel(new SingleTreeSelectionModel());
                 inheritanceTree.setShowsRootHandles(true);
                 inheritanceTree.addMouseListener(new MouseAdapter() {
                         @Override
@@ -205,6 +203,7 @@ public class Gui implements LanguageChangeListener {
                 implementationsTree = new JTree();
                 implementationsTree.setModel(null);
                 implementationsTree.setCellRenderer(new ImplementationsTreeCellRenderer(this));
+                implementationsTree.setSelectionModel(new SingleTreeSelectionModel());
                 implementationsTree.setShowsRootHandles(true);
                 implementationsTree.addMouseListener(new MouseAdapter() {
                         @Override
@@ -235,6 +234,7 @@ public class Gui implements LanguageChangeListener {
                 callsTree = new JTree();
                 callsTree.setModel(null);
                 callsTree.setCellRenderer(new CallsTreeCellRenderer(this));
+                callsTree.setSelectionModel(new SingleTreeSelectionModel());
                 callsTree.setShowsRootHandles(true);
                 callsTree.addMouseListener(new MouseAdapter() {
                         @SuppressWarnings("unchecked")
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/SearchDialog.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/SearchDialog.java
index 2d396c36..9ed7339c 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/SearchDialog.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/SearchDialog.java
@@ -16,19 +16,22 @@ import java.awt.Color;
 import java.awt.FlowLayout;
 import java.awt.Font;
 import java.awt.event.*;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.List;
+import java.util.*;
 
 import javax.swing.*;
 import javax.swing.event.DocumentEvent;
 import javax.swing.event.DocumentListener;
 
+import cuchaz.enigma.analysis.index.EntryIndex;
 import cuchaz.enigma.gui.Gui;
 import cuchaz.enigma.gui.GuiController;
 import cuchaz.enigma.gui.util.AbstractListCellRenderer;
+import cuchaz.enigma.gui.util.GuiUtil;
 import cuchaz.enigma.gui.util.ScaleUtil;
 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 cuchaz.enigma.gui.search.SearchEntry;
 import cuchaz.enigma.gui.search.SearchUtil;
@@ -49,7 +52,7 @@ public class SearchDialog {
 
                 su = new SearchUtil<>();
 
-                dialog = new JDialog(parent.getFrame(), I18n.translate("menu.view.search"), true);
+                dialog = new JDialog(parent.getFrame(), I18n.translate("menu.search"), true);
                 JPanel contentPane = new JPanel();
                 contentPane.setBorder(ScaleUtil.createEmptyBorder(4, 4, 4, 4));
                 contentPane.setLayout(new BorderLayout(ScaleUtil.scale(4), ScaleUtil.scale(4)));
@@ -95,7 +98,7 @@ public class SearchDialog {
                 classListModel = new DefaultListModel<>();
                 classList = new JList<>();
                 classList.setModel(classListModel);
-                classList.setCellRenderer(new ListCellRendererImpl());
+                classList.setCellRenderer(new ListCellRendererImpl(parent));
                 classList.setSelectionMode(ListSelectionModel.SINGLE_SELECTION);
                 classList.addMouseListener(new MouseAdapter() {
                         @Override
@@ -133,14 +136,37 @@ public class SearchDialog {
                 dialog.setLocationRelativeTo(parent.getFrame());
         }
 
-        public void show() {
+        public void show(Type type) {
                 su.clear();
-                parent.getController().project.getJarIndex().getEntryIndex().getClasses().parallelStream()
-                                .filter(e -> !e.isInnerClass())
-                                .map(e -> SearchEntryImpl.from(e, parent.getController()))
-                                .map(SearchUtil.Entry::from)
-                                .sequential()
-                                .forEach(su::add);
+
+                final EntryIndex entryIndex = parent.getController().project.getJarIndex().getEntryIndex();
+
+                switch (type) {
+                        default:
+                        case CLASS:
+                                entryIndex.getClasses().parallelStream()
+                                                .filter(e -> !e.isInnerClass())
+                                                .map(e -> SearchEntryImpl.from(e, parent.getController()))
+                                                .map(SearchUtil.Entry::from)
+                                                .sequential()
+                                                .forEach(su::add);
+                                break;
+                        case METHOD:
+                                entryIndex.getMethods().parallelStream()
+                                                .filter(e -> !e.isConstructor() && !entryIndex.getMethodAccess(e).isSynthetic())
+                                                .map(e -> SearchEntryImpl.from(e, parent.getController()))
+                                                .map(SearchUtil.Entry::from)
+                                                .sequential()
+                                                .forEach(su::add);
+                                break;
+                        case FIELD:
+                                entryIndex.getFields().parallelStream()
+                                                .map(e -> SearchEntryImpl.from(e, parent.getController()))
+                                                .map(SearchUtil.Entry::from)
+                                                .sequential()
+                                                .forEach(su::add);
+                                break;
+                }
 
                 updateList();
 
@@ -161,10 +187,18 @@ public class SearchDialog {
                 close();
                 su.hit(e);
                 parent.getController().navigateTo(e.obf);
-                if (e.deobf != null) {
-                        parent.getDeobfPanel().deobfClasses.setSelectionClass(e.deobf);
+                if (e.obf instanceof ClassEntry) {
+                        if (e.deobf != null) {
+                                parent.getDeobfPanel().deobfClasses.setSelectionClass((ClassEntry) e.deobf);
+                        } else {
+                                parent.getObfPanel().obfClasses.setSelectionClass((ClassEntry) e.obf);
+                        }
                 } else {
-                        parent.getObfPanel().obfClasses.setSelectionClass(e.obf);
+                        if (e.deobf != null) {
+                                parent.getDeobfPanel().deobfClasses.setSelectionClass((ClassEntry) e.deobf.getParent());
+                        } else {
+                                parent.getObfPanel().obfClasses.setSelectionClass((ClassEntry) e.obf.getParent());
+                        }
                 }
         }
 
@@ -189,10 +223,10 @@ public class SearchDialog {
 
         private static final class SearchEntryImpl implements SearchEntry {
 
-                public final ClassEntry obf;
-                public final ClassEntry deobf;
+                public final ParentedEntry<?> obf;
+                public final ParentedEntry<?> deobf;
 
-                private SearchEntryImpl(ClassEntry obf, ClassEntry deobf) {
+                private SearchEntryImpl(ParentedEntry<?> obf, ParentedEntry<?> deobf) {
                         this.obf = obf;
                         this.deobf = deobf;
                 }
@@ -216,8 +250,8 @@ public class SearchDialog {
                         return String.format("SearchEntryImpl { obf: %s, deobf: %s }", obf, deobf);
                 }
 
-                public static SearchEntryImpl from(ClassEntry e, GuiController controller) {
-                        ClassEntry deobf = controller.project.getMapper().deobfuscate(e);
+                public static SearchEntryImpl from(ParentedEntry<?> e, GuiController controller) {
+                        ParentedEntry<?> deobf = controller.project.getMapper().deobfuscate(e);
                         if (deobf.equals(e)) deobf = null;
                         return new SearchEntryImpl(e, deobf);
                 }
@@ -225,12 +259,13 @@ public class SearchDialog {
         }
 
         private static final class ListCellRendererImpl extends AbstractListCellRenderer<SearchEntryImpl> {
-
+                private final Gui gui;
                 private final JLabel mainName;
                 private final JLabel secondaryName;
 
-                public ListCellRendererImpl() {
+                public ListCellRendererImpl(Gui gui) {
                         this.setLayout(new BorderLayout());
+                        this.gui = gui;
 
                         mainName = new JLabel();
                         this.add(mainName, BorderLayout.WEST);
@@ -244,18 +279,31 @@ public class SearchDialog {
                 @Override
                 public void updateUiForEntry(JList<? extends SearchEntryImpl> list, SearchEntryImpl value, int index, boolean isSelected, boolean cellHasFocus) {
                         if (value.deobf == null) {
-                                mainName.setText(value.obf.getSimpleName());
+                                mainName.setText(value.obf.getContextualName());
                                 mainName.setToolTipText(value.obf.getFullName());
                                 secondaryName.setText("");
                                 secondaryName.setToolTipText("");
                         } else {
-                                mainName.setText(value.deobf.getSimpleName());
+                                mainName.setText(value.deobf.getContextualName());
                                 mainName.setToolTipText(value.deobf.getFullName());
                                 secondaryName.setText(value.obf.getSimpleName());
                                 secondaryName.setToolTipText(value.obf.getFullName());
                         }
+
+                        if (value.obf instanceof ClassEntry) {
+                                mainName.setIcon(GuiUtil.getClassIcon(gui, (ClassEntry) value.obf));
+                        } else if (value.obf instanceof MethodEntry) {
+                                mainName.setIcon(GuiUtil.getMethodIcon((MethodEntry) value.obf));
+                        } else if (value.obf instanceof FieldEntry) {
+                                mainName.setIcon(GuiUtil.FIELD_ICON);
+                        }
                 }
 
         }
 
+        public enum Type {
+                CLASS,
+                METHOD,
+                FIELD
+        }
 }
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
index f7f42a58..39aaebb5 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/elements/MenuBar.java
@@ -57,7 +57,11 @@ public class MenuBar {
         private final JMenu scaleMenu = new JMenu();
         private final JMenuItem fontItem = new JMenuItem();
         private final JMenuItem customScaleItem = new JMenuItem();
-        private final JMenuItem searchItem = new JMenuItem();
+
+        private final JMenu searchMenu = new JMenu();
+        private final JMenuItem searchClassItem = new JMenuItem();
+        private final JMenuItem searchMethodItem = new JMenuItem();
+        private final JMenuItem searchFieldItem = new JMenuItem();
 
         private final JMenu collabMenu = new JMenu();
         private final JMenuItem connectItem = new JMenuItem();
@@ -108,10 +112,13 @@ public class MenuBar {
                 this.scaleMenu.add(this.customScaleItem);
                 this.viewMenu.add(this.scaleMenu);
                 this.viewMenu.add(this.fontItem);
-                this.viewMenu.addSeparator();
-                this.viewMenu.add(this.searchItem);
                 this.ui.add(this.viewMenu);
 
+                this.searchMenu.add(this.searchClassItem);
+                this.searchMenu.add(this.searchMethodItem);
+                this.searchMenu.add(this.searchFieldItem);
+                this.ui.add(this.searchMenu);
+
                 this.collabMenu.add(this.connectItem);
                 this.collabMenu.add(this.startServerItem);
                 this.ui.add(this.collabMenu);
@@ -121,7 +128,7 @@ public class MenuBar {
                 this.ui.add(this.helpMenu);
 
                 this.saveMappingsItem.setAccelerator(KeyStroke.getKeyStroke(KeyEvent.VK_S, InputEvent.CTRL_DOWN_MASK));
-                this.searchItem.setAccelerator(KeyStroke.getKeyStroke(KeyEvent.VK_SPACE, InputEvent.SHIFT_DOWN_MASK));
+                this.searchClassItem.setAccelerator(KeyStroke.getKeyStroke(KeyEvent.VK_SPACE, InputEvent.SHIFT_DOWN_MASK));
 
                 this.jarOpenItem.addActionListener(_e -> this.onOpenJarClicked());
                 this.jarCloseItem.addActionListener(_e -> this.gui.getController().closeJar());
@@ -136,7 +143,9 @@ public class MenuBar {
                 this.exitItem.addActionListener(_e -> this.gui.close());
                 this.customScaleItem.addActionListener(_e -> this.onCustomScaleClicked());
                 this.fontItem.addActionListener(_e -> this.onFontClicked(this.gui));
-                this.searchItem.addActionListener(_e -> this.onSearchClicked());
+                this.searchClassItem.addActionListener(_e -> this.onSearchClicked(SearchDialog.Type.CLASS));
+                this.searchMethodItem.addActionListener(_e -> this.onSearchClicked(SearchDialog.Type.METHOD));
+                this.searchFieldItem.addActionListener(_e -> this.onSearchClicked(SearchDialog.Type.FIELD));
                 this.connectItem.addActionListener(_e -> this.onConnectClicked());
                 this.startServerItem.addActionListener(_e -> this.onStartServerClicked());
                 this.aboutItem.addActionListener(_e -> AboutDialog.show(this.gui.getFrame()));
@@ -188,7 +197,11 @@ public class MenuBar {
                 this.scaleMenu.setText(I18n.translate("menu.view.scale"));
                 this.fontItem.setText(I18n.translate("menu.view.font"));
                 this.customScaleItem.setText(I18n.translate("menu.view.scale.custom"));
-                this.searchItem.setText(I18n.translate("menu.view.search"));
+
+                this.searchMenu.setText(I18n.translate("menu.search"));
+                this.searchClassItem.setText(I18n.translate("menu.search.class"));
+                this.searchMethodItem.setText(I18n.translate("menu.search.method"));
+                this.searchFieldItem.setText(I18n.translate("menu.search.field"));
 
                 this.collabMenu.setText(I18n.translate("menu.collab"));
                 this.connectItem.setText(I18n.translate("menu.collab.connect"));
@@ -295,9 +308,9 @@ public class MenuBar {
                 FontDialog.display(gui.getFrame());
         }
 
-        private void onSearchClicked() {
+        private void onSearchClicked(SearchDialog.Type type) {
                 if (this.gui.getController().project != null) {
-                        this.gui.getSearchDialog().show();
+                        this.gui.getSearchDialog().show(type);
                 }
         }
 
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
index 6078145f..d8c46614 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/panels/StructurePanel.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/panels/StructurePanel.java
@@ -3,6 +3,7 @@ package cuchaz.enigma.gui.panels;
 import cuchaz.enigma.analysis.StructureTreeNode;
 import cuchaz.enigma.gui.Gui;
 import cuchaz.enigma.gui.util.GuiUtil;
+import cuchaz.enigma.gui.util.SingleTreeSelectionModel;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 import cuchaz.enigma.translation.representation.entry.FieldEntry;
 import cuchaz.enigma.translation.representation.entry.MethodEntry;
@@ -32,6 +33,7 @@ public class StructurePanel extends JPanel {
         this.structureTree = new JTree();
         this.structureTree.setModel(null);
         this.structureTree.setCellRenderer(new StructureTreeCellRenderer(gui));
+        this.structureTree.setSelectionModel(new SingleTreeSelectionModel());
         this.structureTree.setShowsRootHandles(true);
         this.structureTree.addMouseListener(new MouseAdapter() {
             @Override
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/util/SingleTreeSelectionModel.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/util/SingleTreeSelectionModel.java
new file mode 100644
index 00000000..8915264b
--- /dev/null
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/util/SingleTreeSelectionModel.java
@@ -0,0 +1,11 @@
+package cuchaz.enigma.gui.util;
+
+import javax.swing.tree.DefaultTreeSelectionModel;
+import javax.swing.tree.TreeSelectionModel;
+
+public class SingleTreeSelectionModel extends DefaultTreeSelectionModel {
+
+    public SingleTreeSelectionModel() {
+        this.setSelectionMode(TreeSelectionModel.SINGLE_TREE_SELECTION);
+    }
+}