Allow filtering of packages in Mapping Stats window (#481)

Co-authored-by: Eli Orona <eliorona@live.com>

3 files changed, 59 insertions, 9 deletions

diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/StatsDialog.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/StatsDialog.java
index 1ab66ef2..5f030382 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/StatsDialog.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/dialog/StatsDialog.java
@@ -37,13 +37,13 @@ public class StatsDialog {
                                 results.put(member, statsGenerator.generate(listener, Collections.singleton(member), "", false));
                         }
 
-                        SwingUtilities.invokeLater(() -> show(gui, results));
+                        SwingUtilities.invokeLater(() -> show(gui, results, ""));
                 });
         }
 
-        public static void show(Gui gui, Map<StatsMember, StatsResult> results) {
+        public static void show(Gui gui, Map<StatsMember, StatsResult> results, String packageName) {
                 // init frame
-                JDialog dialog = new JDialog(gui.getFrame(), I18n.translate("menu.file.stats.title"), true);
+                JDialog dialog = new JDialog(gui.getFrame(), packageName.isEmpty() ? I18n.translate("menu.file.stats.title") : I18n.translateFormatted("menu.file.stats.title_filtered", packageName), true);
                 Container contentPane = dialog.getContentPane();
                 contentPane.setLayout(new GridBagLayout());
 
@@ -80,10 +80,30 @@ public class StatsDialog {
                 topLevelPackage.setText(UiConfig.getLastTopLevelPackage());
                 contentPane.add(topLevelPackage, cb1.pos(0, results.size() + 2).fill(GridBagConstraints.HORIZONTAL).build());
 
+                // Show filter button
+                JButton filterButton = new JButton(I18n.translate("menu.file.stats.filter"));
+                filterButton.addActionListener(action -> {
+                        dialog.dispose();
+                        ProgressDialog.runOffThread(gui.getFrame(), listener -> {
+                                UiConfig.setLastTopLevelPackage(topLevelPackage.getText());
+                                UiConfig.save();
+
+                                final StatsGenerator statsGenerator = new StatsGenerator(gui.getController().project);
+                                final Map<StatsMember, StatsResult> filteredResults = new HashMap<>();
+
+                                for (StatsMember member : StatsMember.values()) {
+                                        filteredResults.put(member, statsGenerator.generate(listener, Collections.singleton(member), UiConfig.getLastTopLevelPackage(), false));
+                                }
+
+                                SwingUtilities.invokeLater(() -> show(gui, filteredResults, UiConfig.getLastTopLevelPackage()));
+                        });
+                });
+                contentPane.add(filterButton, cb1.pos(0, results.size() + 3).anchor(GridBagConstraints.EAST).build());
+
                 // show synthetic members option
                 JCheckBox syntheticParametersOption = new JCheckBox(I18n.translate("menu.file.stats.synthetic_parameters"));
                 syntheticParametersOption.setSelected(UiConfig.shouldIncludeSyntheticParameters());
-                contentPane.add(syntheticParametersOption, cb1.pos(0, results.size() + 3).build());
+                contentPane.add(syntheticParametersOption, cb1.pos(0, results.size() + 4).build());
 
                 // show generate button
                 JButton button = new JButton(I18n.translate("menu.file.stats.generate"));
@@ -98,7 +118,7 @@ public class StatsDialog {
                         generateStats(gui, checkboxes, topLevelPackage.getText(), syntheticParametersOption.isSelected());
                 });
 
-                contentPane.add(button, cb1.pos(0, results.size() + 4).weightY(1.0).anchor(GridBagConstraints.SOUTHEAST).build());
+                contentPane.add(button, cb1.pos(0, results.size() + 5).weightY(1.0).anchor(GridBagConstraints.SOUTHWEST).build());
 
                 // add action listener to each checkbox
                 checkboxes.forEach((key, value) -> value.addActionListener(action -> {
diff --git a/enigma-swing/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java b/enigma-swing/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java
index 99b5572b..5d241665 100644
--- a/enigma-swing/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java
+++ b/enigma-swing/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java
@@ -54,15 +54,20 @@ public class StatsGenerator {
                 progress.init(totalWork, I18n.translate("progress.stats"));
 
                 Map<String, Integer> counts = new HashMap<>();
-
+                String topLevelPackageSlash = topLevelPackage.replace(".", "/");
                 int numDone = 0;
 
                 if (includedMembers.contains(StatsMember.METHODS) || includedMembers.contains(StatsMember.PARAMETERS)) {
                         for (MethodEntry method : entryIndex.getMethods()) {
                                 progress.step(numDone++, I18n.translate("type.methods"));
-                                MethodEntry root = entryResolver.resolveEntry(method, ResolutionStrategy.RESOLVE_ROOT).stream().findFirst().orElseThrow(AssertionError::new);
-
-                                if (root == method) {
+                                MethodEntry root = entryResolver
+                                                .resolveEntry(method, ResolutionStrategy.RESOLVE_ROOT)
+                                                .stream()
+                                                .findFirst()
+                                                .orElseThrow(AssertionError::new);
+                                ClassEntry clazz = root.getParent();
+
+                                if (root == method && this.mapper.deobfuscate(clazz).getPackageName().startsWith(topLevelPackageSlash)) {
                                         if (includedMembers.contains(StatsMember.METHODS) && !((MethodDefEntry) method).getAccess().isSynthetic()) {
                                                 update(counts, method);
                                                 totalMappable++;
@@ -84,6 +89,29 @@ public class StatsGenerator {
                 if (includedMembers.contains(StatsMember.FIELDS)) {
                         for (FieldEntry field : entryIndex.getFields()) {
                                 progress.step(numDone++, I18n.translate("type.fields"));
+                                ClassEntry clazz = field.getParent();
+
+                                if (!((FieldDefEntry) field).getAccess().isSynthetic() && this.mapper.deobfuscate(clazz).getPackageName().startsWith(topLevelPackageSlash)) {
+                                        update(counts, field);
+                                        totalMappable++;
+                                }
+                        }
+                }
+
+                if (includedMembers.contains(StatsMember.CLASSES)) {
+                        for (ClassEntry clazz : entryIndex.getClasses()) {
+                                progress.step(numDone++, I18n.translate("type.classes"));
+
+                                if (this.mapper.deobfuscate(clazz).getPackageName().startsWith(topLevelPackageSlash)) {
+                                        update(counts, clazz);
+                                        totalMappable++;
+                                }
+                        }
+                }
+
+                if (includedMembers.contains(StatsMember.FIELDS)) {
+                        for (FieldEntry field : entryIndex.getFields()) {
+                                progress.step(numDone++, I18n.translate("type.fields"));
 
                                 if (!((FieldDefEntry) field).getAccess().isSynthetic()) {
                                         update(counts, field);
diff --git a/enigma/src/main/resources/lang/en_us.json b/enigma/src/main/resources/lang/en_us.json
index 4ec2f087..76afebc2 100644
--- a/enigma/src/main/resources/lang/en_us.json
+++ b/enigma/src/main/resources/lang/en_us.json
@@ -28,6 +28,8 @@
         "menu.file.export.jar": "Export Jar...",
         "menu.file.stats": "Mapping Stats...",
         "menu.file.stats.title": "Mapping Stats",
+        "menu.file.stats.title_filtered": "Mapping Stats for %s",
+        "menu.file.stats.filter": "Filter",
         "menu.file.stats.top_level_package": "Top-Level Package:",
         "menu.file.stats.synthetic_parameters": "Include Synthetic Parameters",
         "menu.file.stats.generate": "Generate Diagram",