Allow filtering of packages in Mapping Stats window

NebelNidas · Sep 24, 2022 · f992016 · f992016
1 parent 98f5c0e
commit f992016
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 11 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
@@ -36,14 +36,13 @@ public static void show(Gui gui) {
 			for (StatsMember member : StatsMember.values()) {
 				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 +79,29 @@ public static void show(Gui gui, Map<StatsMember, StatsResult> results) {
 		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> _results = new HashMap<>();
+				for (StatsMember member : StatsMember.values()) {
+					_results.put(member, statsGenerator.generate(listener, Collections.singleton(member), UiConfig.getLastTopLevelPackage(), false));
+				}
+				SwingUtilities.invokeLater(() -> show(gui, _results, 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 +116,7 @@ public static void show(Gui gui, Map<StatsMember, StatsResult> results) {
 			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
@@ -55,22 +55,27 @@ public StatsResult generate(ProgressListener progress, Set<StatsMember> included
 
 		Map<String, Integer> counts = new HashMap<>();
 
-		int numDone = 0;
+		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++;
 					}
 
 					if (includedMembers.contains(StatsMember.PARAMETERS) && (!((MethodDefEntry) method).getAccess().isSynthetic() || includeSynthetic)) {
 						int index = ((MethodDefEntry) method).getAccess().isStatic() ? 0 : 1;
-
 						for (TypeDescriptor argument : method.getDesc().getArgumentDescs()) {
 							update(counts, new LocalVariableEntry(method, index, "", true, null));
 							index += argument.getSize();
@@ -81,6 +86,27 @@ public StatsResult generate(ProgressListener progress, Set<StatsMember> included
 			}
 		}
 
+		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"));

diff --git 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",