Skip to content

Commit

Permalink
show mapping status in class tree (#60)
Browse files Browse the repository at this point in the history
* initial really broken impl

* temporarily remove shadowing warning

* better icons and fix names

* real time updates

* cleanup

* use StatsGenerator instead of @NebelNidas ' code

* fix StatsGenerator

* remove debug message

* fix all classes docker not properly reloading when classes are moved to new packages

* helpful comment

* fix massive lag

* fix some issues with StatsGenerator that I seem to have caused for no good reason

* automatically reload icons

* fix nodes calculating their stats repeatedly

* run icon reloading off thread to prevent lag when renaming an entry

* fix isObfuscated in EnigmaProject

* remove superfluous reload in ClassSelector

* fix all classes docker losing its expansion state on rename (I already caused this exact bug once, we stay silly :iea:)

* StatsGenerator fix for anonymous classes always showing as a mapped entry

* fix class selector exploding sometimes
  • Loading branch information
ix0rai authored Feb 27, 2023
1 parent aa05df2 commit f6ba9ee
Show file tree
Hide file tree
Showing 12 changed files with 412 additions and 95 deletions.
55 changes: 45 additions & 10 deletions enigma-swing/src/main/java/cuchaz/enigma/gui/ClassSelector.java
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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 {
Expand Down
16 changes: 12 additions & 4 deletions enigma-swing/src/main/java/cuchaz/enigma/gui/Gui.java
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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);
Expand Down
12 changes: 7 additions & 5 deletions enigma-swing/src/main/java/cuchaz/enigma/gui/GuiController.java
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand All @@ -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
Expand All @@ -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);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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++;
}
}
}
Expand All @@ -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);
}
}
}
Expand All @@ -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);
}
}
}
Expand All @@ -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;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand All @@ -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;
Expand Down
Loading

0 comments on commit f6ba9ee

Please sign in to comment.