show mapping status in class tree

enigma-swing/src/main/java/cuchaz/enigma/gui/ClassSelector.java

@@ -17,14 +17,25 @@
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 import cuchaz.enigma.utils.validation.ValidationContext;
 
-import javax.swing.*;
+import javax.swing.BoxLayout;
+import javax.swing.JLabel;
+import javax.swing.JPanel;
+import javax.swing.JTree;
 import javax.swing.event.CellEditorListener;
 import javax.swing.event.ChangeEvent;
-import javax.swing.tree.*;
-import java.awt.*;
+import javax.swing.tree.DefaultMutableTreeNode;
+import javax.swing.tree.DefaultTreeCellEditor;
+import javax.swing.tree.DefaultTreeCellRenderer;
+import javax.swing.tree.DefaultTreeModel;
+import javax.swing.tree.TreeNode;
+import javax.swing.tree.TreePath;
+import java.awt.Component;
 import java.awt.event.MouseEvent;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.EventObject;
 import java.util.List;
-import java.util.*;
 
 public class ClassSelector extends JTree {
 	public static final Comparator<ClassEntry> DEOBF_CLASS_COMPARATOR = Comparator.comparing(ClassEntry::getFullName);
@@ -88,7 +99,21 @@ public Component getTreeCellRendererComponent(JTree tree, Object value, boolean
 				super.getTreeCellRendererComponent(tree, value, sel, expanded, leaf, row, hasFocus);
 
 				if (gui.getController().project != null && leaf && value instanceof ClassSelectorClassNode node) {
-					this.setIcon(GuiUtil.getClassIcon(gui, node.getObfEntry()));
+					JPanel panel = new JPanel();
+					panel.setOpaque(false);
+					panel.setLayout(new BoxLayout(panel, BoxLayout.X_AXIS));
+					panel.add(new JLabel(GuiUtil.getClassIcon(gui, node.getObfEntry())));
+
+					if (node.getStats() == null) {
+						// calculate stats on a separate thread for performance reasons
+						this.setIcon(GuiUtil.PENDING_STATUS_ICON);
+						node.reloadStats(gui, ClassSelector.this, false);
+					} else {
+						this.setIcon(GuiUtil.getDeobfuscationIcon(node.getStats()));
+					}
+
+					panel.add(this);
+					return panel;
 				}
 
 				return this;
@@ -111,8 +136,7 @@ public void editingStopped(ChangeEvent e) {
 				String data = editor.getCellEditorValue().toString();
 				TreePath path = ClassSelector.this.getSelectionPath();
 
-				Object realPath = path.getLastPathComponent();
-				if (realPath instanceof DefaultMutableTreeNode node && data != null) {
+				if (path != null && path.getLastPathComponent() instanceof DefaultMutableTreeNode node && data != null) {
 					TreeNode parentNode = node.getParent();
 					if (parentNode == null)
 						return;
@@ -180,7 +204,7 @@ public void setClasses(Collection<ClassEntry> classEntries) {
 	}
 
 	public ClassEntry getSelectedClass() {
-		if (!this.isSelectionEmpty()) {
+		if (!this.isSelectionEmpty() && this.getSelectionPath() != null) {
 			Object selectedNode = this.getSelectionPath().getLastPathComponent();
 
 			if (selectedNode instanceof ClassSelectorClassNode classNode) {
@@ -266,9 +290,20 @@ public void removeEntry(ClassEntry classEntry) {
 		this.packageManager.removeClassNode(classEntry);
 	}
 
-	public void reload() {
+	public void reloadEntry(ClassEntry classEntry) {
+		this.removeEntry(classEntry);
+		this.moveClassIn(classEntry);
+		ClassSelectorClassNode node = this.packageManager.getClassNode(classEntry);
+		node.reloadStats(controller.getGui(), this, true);
+	}
+
+	public void reload(TreeNode node) {
 		DefaultTreeModel model = (DefaultTreeModel) this.getModel();
-		model.reload(this.packageManager.getRoot());
+		model.reload(node);
+	}
+
+	public void reload() {
+		this.reload(this.packageManager.getRoot());
 	}
 
 	public interface ClassSelectionListener {

enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java

@@ -571,7 +571,6 @@ public void moveClassTree(Entry<?> obfEntry, boolean isOldOb, boolean isNewOb) {
 
 		ClassSelector deobfuscatedClassSelector = Docker.getDocker(DeobfuscatedClassesDocker.class).getClassSelector();
 		ClassSelector obfuscatedClassSelector = Docker.getDocker(ObfuscatedClassesDocker.class).getClassSelector();
-		ClassSelector allClassesClassSelector = Docker.getDocker(AllClassesDocker.class).getClassSelector();
 
 		List<ClassSelector.StateEntry> deobfuscatedPanelExpansionState = deobfuscatedClassSelector.getExpansionState();
 		List<ClassSelector.StateEntry> obfuscatedPanelExpansionState = obfuscatedClassSelector.getExpansionState();
@@ -594,14 +593,23 @@ public void moveClassTree(Entry<?> obfEntry, boolean isOldOb, boolean isNewOb) {
 			deobfuscatedClassSelector.reload();
 		}
 
-		allClassesClassSelector.removeEntry(classEntry);
-		allClassesClassSelector.moveClassIn(classEntry);
-		allClassesClassSelector.reload();
+		this.reloadClassEntry(classEntry);
 
 		deobfuscatedClassSelector.restoreExpansionState(deobfuscatedPanelExpansionState);
 		obfuscatedClassSelector.restoreExpansionState(obfuscatedPanelExpansionState);
 	}
 
+	public void reloadClassEntry(ClassEntry classEntry) {
+		Docker.getDocker(DeobfuscatedClassesDocker.class).getClassSelector().reloadEntry(classEntry);
+		Docker.getDocker(ObfuscatedClassesDocker.class).getClassSelector().reloadEntry(classEntry);
+
+		ClassSelector allClassesClassSelector = Docker.getDocker(AllClassesDocker.class).getClassSelector();
+		List<ClassSelector.StateEntry> expansionState = allClassesClassSelector.getExpansionState();
+		allClassesClassSelector.reloadEntry(classEntry);
+		allClassesClassSelector.reload();
+		allClassesClassSelector.restoreExpansionState(expansionState);
+	}
+
 	public SearchDialog getSearchDialog() {
 		if (this.searchDialog == null) {
 			this.searchDialog = new SearchDialog(this);

enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java

@@ -510,10 +510,7 @@ private void applyChange0(ValidationContext vc, EntryChange<?> change) {
 		EntryMapping mapping = EntryUtil.applyChange(vc, this.project.getMapper(), change);
 
 		boolean renamed = !change.getDeobfName().isUnchanged();
-
-		if (renamed && target instanceof ClassEntry classEntry && !classEntry.isInnerClass()) {
-			this.gui.moveClassTree(target, prev.targetName() == null, mapping.targetName() == null);
-		}
+		this.gui.updateStructure(this.gui.getActiveEditor());
 
 		if (!Objects.equals(prev.targetName(), mapping.targetName())) {
 			this.chp.invalidateMapped();
@@ -523,7 +520,12 @@ private void applyChange0(ValidationContext vc, EntryChange<?> change) {
 			this.chp.invalidateJavadoc(target.getTopLevelClass());
 		}
 
-		this.gui.updateStructure(this.gui.getActiveEditor());
+		if (renamed && target instanceof ClassEntry classEntry && !classEntry.isInnerClass()) {
+			this.gui.moveClassTree(target, prev.targetName() == null, mapping.targetName() == null);
+			return;
+		}
+
+		this.gui.reloadClassEntry(change.getTarget().getTopLevelClass());
 	}
 
 	public void openStats(Set<StatsMember> includedMembers, String topLevelPackage, boolean includeSynthetic) {

enigma-swing/src/main/java/cuchaz/enigma/gui/node/ClassSelectorClassNode.java

@@ -11,17 +11,27 @@
 
 package cuchaz.enigma.gui.node;
 
+import cuchaz.enigma.ProgressListener;
+import cuchaz.enigma.gui.ClassSelector;
+import cuchaz.enigma.gui.Gui;
+import cuchaz.enigma.gui.stats.StatsGenerator;
+import cuchaz.enigma.gui.stats.StatsResult;
+import cuchaz.enigma.gui.util.GuiUtil;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 
+import javax.swing.SwingWorker;
 import javax.swing.tree.DefaultMutableTreeNode;
+import javax.swing.tree.DefaultTreeCellRenderer;
 
 public class ClassSelectorClassNode extends DefaultMutableTreeNode {
 	private final ClassEntry obfEntry;
 	private ClassEntry classEntry;
+	private StatsResult stats;
 
 	public ClassSelectorClassNode(ClassEntry obfEntry, ClassEntry classEntry) {
 		this.obfEntry = obfEntry;
 		this.classEntry = classEntry;
+		this.stats = null;
 		this.setUserObject(classEntry);
 	}
 
@@ -33,14 +43,50 @@ public ClassEntry getClassEntry() {
 		return this.classEntry;
 	}
 
+	public StatsResult getStats() {
+		return this.stats;
+	}
+
+	public void setStats(StatsResult stats) {
+		this.stats = stats;
+	}
+
+	/**
+	 * Reloads the stats for this class node and updates the icon in the provided class selector.
+	 * @param gui the current gui instance
+	 * @param selector the class selector to reload on
+	 * @param updateIfPresent whether to update the stats if they have already been generated for this node
+	 */
+	public void reloadStats(Gui gui, ClassSelector selector, boolean updateIfPresent) {
+		SwingWorker<ClassSelectorClassNode, Void> iconUpdateWorker = new SwingWorker<>() {
+			@Override
+			protected ClassSelectorClassNode doInBackground() {
+				if (ClassSelectorClassNode.this.getStats() == null || updateIfPresent) {
+					StatsResult newStats = new StatsGenerator(gui.getController().project).generateForClassTree(ProgressListener.none(), ClassSelectorClassNode.this.getObfEntry(), false);
+					ClassSelectorClassNode.this.setStats(newStats);
+				}
+
+				return ClassSelectorClassNode.this;
+			}
+
+			@Override
+			public void done() {
+				((DefaultTreeCellRenderer) selector.getCellRenderer()).setIcon(GuiUtil.getDeobfuscationIcon(ClassSelectorClassNode.this.getStats()));
+				selector.reload(ClassSelectorClassNode.this);
+			}
+		};
+
+		iconUpdateWorker.execute();
+	}
+
 	@Override
 	public String toString() {
 		return this.classEntry.getSimpleName();
 	}
 
 	@Override
 	public boolean equals(Object other) {
-		return other instanceof ClassSelectorClassNode && this.equals((ClassSelectorClassNode) other);
+		return other instanceof ClassSelectorClassNode node && this.equals(node);
 	}
 
 	@Override
@@ -60,8 +106,8 @@ public void setUserObject(Object userObject) {
 			packageName = this.classEntry.getPackageName() + "/";
 		if (userObject instanceof String)
 			this.classEntry = new ClassEntry(packageName + userObject);
-		else if (userObject instanceof ClassEntry)
-			this.classEntry = (ClassEntry) userObject;
+		else if (userObject instanceof ClassEntry entry)
+			this.classEntry = entry;
 		super.setUserObject(this.classEntry);
 	}
 

enigma-swing/src/main/java/cuchaz/enigma/gui/stats/StatsGenerator.java

@@ -34,7 +34,24 @@ public StatsGenerator(EnigmaProject project) {
 		this.entryResolver = project.getJarIndex().getEntryResolver();
 	}
 
+	public StatsResult generateForClassTree(ProgressListener progress, ClassEntry entry, boolean includeSynthetic) {
+		return generate(progress, EnumSet.allOf(StatsMember.class), entry.getFullName(), true, includeSynthetic);
+	}
+
 	public StatsResult generate(ProgressListener progress, Set<StatsMember> includedMembers, String topLevelPackage, boolean includeSynthetic) {
+		return generate(progress, includedMembers, topLevelPackage, false, includeSynthetic);
+	}
+
+	/**
+	 * Generates stats for the given package or class.
+	 * @param progress a listener to update with current progress
+	 * @param includedMembers the types of entry to include in the stats
+	 * @param topLevelPackage the package or class to generate stats for
+	 * @param forClassTree if true, the stats will be generated for the class tree - this means that non-mappable obfuscated entries will be ignored for correctness
+	 * @param includeSynthetic whether to include synthetic methods
+	 * @return the generated {@link StatsResult} for the provided class or package.
+	 */
+	public StatsResult generate(ProgressListener progress, Set<StatsMember> includedMembers, String topLevelPackage, boolean forClassTree, boolean includeSynthetic) {
 		includedMembers = EnumSet.copyOf(includedMembers);
 		int totalWork = 0;
 		int totalMappable = 0;
@@ -61,27 +78,30 @@ public StatsResult generate(ProgressListener progress, Set<StatsMember> included
 		if (includedMembers.contains(StatsMember.METHODS) || includedMembers.contains(StatsMember.PARAMETERS)) {
 			for (MethodEntry method : this.entryIndex.getMethods()) {
 				progress.step(numDone++, I18n.translate("type.methods"));
+
+				// we don't want constructors or otherwise non-mappable things to show as a mapped method!
+				if (!project.isRenamable(method)) {
+					continue;
+				}
+
 				MethodEntry root = this.entryResolver
 						.resolveEntry(method, ResolutionStrategy.RESOLVE_ROOT)
 						.stream()
 						.findFirst()
 						.orElseThrow(AssertionError::new);
 
                 ClassEntry clazz = root.getParent();
-				String deobfuscatedPackageName = this.mapper.deobfuscate(clazz).getPackageName();
 
-                if (root == method && (topLevelPackageSlash.isBlank() || (deobfuscatedPackageName != null && deobfuscatedPackageName.startsWith(topLevelPackageSlash)))) {
+                if (root == method && checkPackage(clazz, topLevelPackageSlash, forClassTree)) {
                     if (includedMembers.contains(StatsMember.METHODS) && !((MethodDefEntry) method).getAccess().isSynthetic()) {
-						this.update(counts, method);
-                        totalMappable++;
+						totalMappable += this.update(counts, method, forClassTree);
                     }
 
 					if (includedMembers.contains(StatsMember.PARAMETERS) && (!((MethodDefEntry) method).getAccess().isSynthetic() || includeSynthetic)) {
 						int index = ((MethodDefEntry) method).getAccess().isStatic() ? 0 : 1;
 						for (TypeDescriptor argument : method.getDesc().getArgumentDescs()) {
-							this.update(counts, new LocalVariableEntry(method, index, "", true, null));
+							totalMappable += this.update(counts, new LocalVariableEntry(method, index, "", true, null), forClassTree);
 							index += argument.getSize();
-							totalMappable++;
 						}
 					}
 				}
@@ -92,11 +112,9 @@ public StatsResult generate(ProgressListener progress, Set<StatsMember> included
             for (FieldEntry field : this.entryIndex.getFields()) {
                 progress.step(numDone++, I18n.translate("type.fields"));
                 ClassEntry clazz = field.getParent();
-				String deobfuscatedPackageName = this.mapper.deobfuscate(clazz).getPackageName();
 
-                if (!((FieldDefEntry) field).getAccess().isSynthetic() && (topLevelPackageSlash.isBlank() || (deobfuscatedPackageName != null && deobfuscatedPackageName.startsWith(topLevelPackageSlash)))) {
-					this.update(counts, field);
-                    totalMappable++;
+                if (!((FieldDefEntry) field).getAccess().isSynthetic() && checkPackage(clazz, topLevelPackageSlash, forClassTree)) {
+					totalMappable += this.update(counts, field, forClassTree);
                 }
             }
         }
@@ -105,11 +123,8 @@ public StatsResult generate(ProgressListener progress, Set<StatsMember> included
             for (ClassEntry clazz : this.entryIndex.getClasses()) {
                 progress.step(numDone++, I18n.translate("type.classes"));
 
-				String deobfuscatedPackageName = this.mapper.deobfuscate(clazz).getPackageName();
-
-                if (topLevelPackageSlash.isBlank() || (deobfuscatedPackageName != null && deobfuscatedPackageName.startsWith(topLevelPackageSlash))) {
-					this.update(counts, clazz);
-                    totalMappable++;
+                if (checkPackage(clazz, topLevelPackageSlash, forClassTree)) {
+					totalMappable += this.update(counts, clazz, forClassTree);
                 }
             }
         }
@@ -128,10 +143,33 @@ public StatsResult generate(ProgressListener progress, Set<StatsMember> included
 		return new StatsResult(totalMappable, counts.values().stream().mapToInt(i -> i).sum(), tree);
 	}
 
-	private void update(Map<String, Integer> counts, Entry<?> entry) {
-		if (this.project.isObfuscated(entry) && this.project.isRenamable(entry) && !this.project.isSynthetic(entry)) {
+	private boolean checkPackage(ClassEntry clazz, String topLevelPackage, boolean singleClass) {
+		String deobfuscatedName = this.mapper.deobfuscate(clazz).getPackageName();
+		if (singleClass) {
+			return (deobfuscatedName != null && deobfuscatedName.startsWith(topLevelPackage)) || clazz.getFullName().startsWith(topLevelPackage);
+		}
+
+		return topLevelPackage.isBlank() || (deobfuscatedName != null && deobfuscatedName.startsWith(topLevelPackage));
+	}
+
+	/**
+	 * @return whether to increment the total mappable entry count - 0 if no, 1 if yes
+	 */
+	private int update(Map<String, Integer> counts, Entry<?> entry, boolean forClassTree) {
+		boolean obfuscated = this.project.isObfuscated(entry);
+		boolean renamable = this.project.isRenamable(entry);
+		boolean synthetic = this.project.isSynthetic(entry);
+
+		if (forClassTree && obfuscated && !renamable) {
+			return 0;
+		}
+
+		if (obfuscated && renamable && !synthetic) {
 			String parent = this.mapper.deobfuscate(entry.getAncestry().get(0)).getName().replace('/', '.');
 			counts.put(parent, counts.getOrDefault(parent, 0) + 1);
+			return 1;
 		}
+
+		return 1;
 	}
 }

enigma-swing/src/main/java/cuchaz/enigma/gui/stats/StatsResult.java

@@ -27,10 +27,16 @@ public int getUnmapped() {
 	}
 
 	public int getMapped() {
-		return this.total - this.unmapped;
+		return this.getTotal() - this.getUnmapped();
 	}
 
 	public double getPercentage() {
+		// avoid showing "Nan%" when there are no entries to map
+		// if there are none, you've mapped them all!
+		if (this.total == 0) {
+			return 100.0f;
+		}
+
 		return (this.getMapped() * 100.0f) / this.total;
 	}
 
@@ -51,7 +57,7 @@ public static class Node<T> {
 			public String name;
 			public T value;
 			public List<Node<T>> children = new ArrayList<>();
-			private final transient Map<String, Node<T>> namedChildren = new HashMap<>();
+			private final Map<String, Node<T>> namedChildren = new HashMap<>();
 
 			public Node(String name, T value) {
 				this.name = name;