Optimize NestedPackages by deferring sorting

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

@@ -2,23 +2,21 @@
 
 import java.util.Collection;
 import java.util.Comparator;
-import java.util.Enumeration;
 import java.util.HashMap;
 import java.util.Map;
 
-import javax.swing.tree.DefaultMutableTreeNode;
-import javax.swing.tree.MutableTreeNode;
 import javax.swing.tree.TreeNode;
 import javax.swing.tree.TreePath;
 
 import cuchaz.enigma.gui.node.ClassSelectorClassNode;
 import cuchaz.enigma.gui.node.ClassSelectorPackageNode;
+import cuchaz.enigma.gui.node.SortedMutableTreeNode;
 import cuchaz.enigma.translation.mapping.EntryRemapper;
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 
 public class NestedPackages {
-	private final DefaultMutableTreeNode root = new DefaultMutableTreeNode();
-	private final Map<String, DefaultMutableTreeNode> packageToNode = new HashMap<>();
+	private final SortedMutableTreeNode root;
+	private final Map<String, SortedMutableTreeNode> packageToNode = new HashMap<>();
 	private final Map<ClassEntry, ClassSelectorClassNode> classToNode = new HashMap<>();
 	private final EntryRemapper remapper;
 	private final Comparator<TreeNode> comparator;
@@ -42,6 +40,7 @@ public NestedPackages(Iterable<ClassEntry> entries, Comparator<ClassEntry> entry
 
 			return 0;
 		};
+		this.root = new SortedMutableTreeNode(comparator);
 
 		for (ClassEntry entry : entries) {
 			addEntry(entry);
@@ -52,31 +51,31 @@ public void addEntry(ClassEntry entry) {
 		ClassEntry translated = remapper.deobfuscate(entry);
 		var me = new ClassSelectorClassNode(entry, translated);
 		classToNode.put(entry, me);
-		insert(getPackage(translated.getPackageName()), me);
+		getPackage(translated.getPackageName()).insert(me, 0);
 	}
 
-	public DefaultMutableTreeNode getPackage(String packageName) {
-		DefaultMutableTreeNode node = packageToNode.get(packageName);
+	public SortedMutableTreeNode getPackage(String packageName) {
+		SortedMutableTreeNode node = packageToNode.get(packageName);
 
 		if (packageName == null) {
 			return root;
 		}
 
 		if (node == null) {
-			node = new ClassSelectorPackageNode(packageName);
-			insert(getPackage(ClassEntry.getParentPackage(packageName)), node);
+			node = new ClassSelectorPackageNode(this.comparator, packageName);
+			getPackage(ClassEntry.getParentPackage(packageName)).insert(node, 0);
 			packageToNode.put(packageName, node);
 		}
 
 		return node;
 	}
 
-	public DefaultMutableTreeNode getRoot() {
+	public SortedMutableTreeNode getRoot() {
 		return root;
 	}
 
 	public TreePath getPackagePath(String packageName) {
-		DefaultMutableTreeNode node = packageToNode.getOrDefault(packageName, root);
+		SortedMutableTreeNode node = packageToNode.getOrDefault(packageName, root);
 		return new TreePath(node.getPath());
 	}
 
@@ -90,11 +89,11 @@ public void removeClassNode(ClassEntry entry) {
 		if (node != null) {
 			node.removeFromParent();
 			// remove dangling packages
-			DefaultMutableTreeNode packageNode = packageToNode.get(entry.getPackageName());
+			SortedMutableTreeNode packageNode = packageToNode.get(entry.getPackageName());
 
 			while (packageNode != null && packageNode.getChildCount() == 0) {
-				DefaultMutableTreeNode theNode = packageNode;
-				packageNode = (DefaultMutableTreeNode) packageNode.getParent();
+				SortedMutableTreeNode theNode = packageNode;
+				packageNode = (SortedMutableTreeNode) packageNode.getParent();
 				theNode.removeFromParent();
 
 				if (theNode instanceof ClassSelectorPackageNode pn) {
@@ -104,22 +103,7 @@ public void removeClassNode(ClassEntry entry) {
 		}
 	}
 
-	public Collection<DefaultMutableTreeNode> getPackageNodes() {
+	public Collection<SortedMutableTreeNode> getPackageNodes() {
 		return packageToNode.values();
 	}
-
-	private void insert(DefaultMutableTreeNode parent, MutableTreeNode child) {
-		int index = 0;
-		Enumeration<TreeNode> children = parent.children();
-
-		while (children.hasMoreElements()) {
-			if (comparator.compare(children.nextElement(), child) < 0) {
-				index++;
-			} else {
-				break;
-			}
-		}
-
-		parent.insert(child, index);
-	}
 }

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

@@ -11,14 +11,18 @@
 
 package cuchaz.enigma.gui.node;
 
-import javax.swing.tree.DefaultMutableTreeNode;
+import java.util.Comparator;
+
+import javax.swing.tree.TreeNode;
 
 import cuchaz.enigma.translation.representation.entry.ClassEntry;
 
-public class ClassSelectorPackageNode extends DefaultMutableTreeNode {
+public class ClassSelectorPackageNode extends SortedMutableTreeNode {
 	private String packageName;
 
-	public ClassSelectorPackageNode(String packageName) {
+	public ClassSelectorPackageNode(Comparator<TreeNode> comparator, String packageName) {
+		super(comparator);
+
 		this.packageName = packageName != null ? packageName : "(none)";
 	}
 

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

@@ -0,0 +1,116 @@
+package cuchaz.enigma.gui.node;
+
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.Enumeration;
+import java.util.Iterator;
+import java.util.List;
+
+import javax.swing.tree.DefaultMutableTreeNode;
+import javax.swing.tree.MutableTreeNode;
+import javax.swing.tree.TreeNode;
+
+/**
+ * A MutableTreeNode whose contents are always guaranteed to be sorted with the given comparator.
+ */
+public class SortedMutableTreeNode extends DefaultMutableTreeNode {
+	private final Comparator<TreeNode> comparator;
+	private final List<TreeNode> children;
+	private boolean isSorted = true;
+
+	public SortedMutableTreeNode(Comparator<TreeNode> comparator) {
+		this.comparator = comparator;
+		this.children = new ArrayList<>();
+	}
+
+	@Override
+	public void insert(MutableTreeNode child, int index) {
+		if (child == null) {
+			throw new IllegalArgumentException("child is null");
+		}
+
+		MutableTreeNode oldParent = (MutableTreeNode) child.getParent();
+
+		if (oldParent != null) {
+			oldParent.remove(child);
+		}
+
+		child.setParent(this);
+		this.children.add(child);
+		this.isSorted = false;
+	}
+
+	private void checkSorted() {
+		if (!this.isSorted) {
+			this.isSorted = true;
+			this.children.sort(this.comparator);
+		}
+	}
+
+	@Override
+	public void remove(int index) {
+		checkSorted();
+
+		remove((MutableTreeNode) getChildAt(index));
+	}
+
+	@Override
+	public void remove(MutableTreeNode node) {
+		this.children.remove(node);
+		node.setParent(null);
+	}
+
+	@Override
+	public TreeNode getChildAt(int childIndex) {
+		checkSorted();
+
+		return this.children.get(childIndex);
+	}
+
+	@Override
+	public int getChildCount() {
+		return this.children.size();
+	}
+
+	@Override
+	public int getIndex(TreeNode node) {
+		checkSorted();
+
+		for (int i = 0; i < this.children.size(); i++) {
+			if (this.comparator.compare(node, children.get(i)) == 0) {
+				return i;
+			}
+		}
+
+		return -1;
+	}
+
+	@Override
+	public boolean getAllowsChildren() {
+		return true;
+	}
+
+	@Override
+	public boolean isLeaf() {
+		return this.children.isEmpty();
+	}
+
+	@Override
+	public Enumeration<TreeNode> children() {
+		checkSorted();
+
+		Iterator<TreeNode> iter = this.children.iterator();
+
+		return new Enumeration<>() {
+			@Override
+			public boolean hasMoreElements() {
+				return iter.hasNext();
+			}
+
+			@Override
+			public TreeNode nextElement() {
+				return iter.next();
+			}
+		};
+	}
+}