path: root/src/cuchaz/enigma/gui/ClassSelector.java

/*******************************************************************************
 * Copyright (c) 2014 Jeff Martin.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the GNU Public License v3.0
 * which accompanies this distribution, and is available at
 * http://www.gnu.org/licenses/gpl.html
 * 
 * Contributors:
 *     Jeff Martin - initial API and implementation
 ******************************************************************************/
package cuchaz.enigma.gui;

import java.awt.event.MouseAdapter;
import java.awt.event.MouseEvent;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.Enumeration;
import java.util.List;
import java.util.Map;

import javax.swing.JTree;
import javax.swing.tree.DefaultMutableTreeNode;
import javax.swing.tree.DefaultTreeModel;
import javax.swing.tree.TreePath;

import com.google.common.collect.ArrayListMultimap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Multimap;

import cuchaz.enigma.mapping.ClassEntry;

public class ClassSelector extends JTree {
	
	private static final long serialVersionUID = -7632046902384775977L;
	
	public interface ClassSelectionListener {
		void onSelectClass(ClassEntry classEntry);
	}
	
	public static Comparator<ClassEntry> ObfuscatedClassEntryComparator;
	public static Comparator<ClassEntry> DeobfuscatedClassEntryComparator;
	
	static {
		ObfuscatedClassEntryComparator = new Comparator<ClassEntry>() {
			@Override
			public int compare(ClassEntry a, ClassEntry b) {
				String aname = a.getName();
				String bname = a.getName();
				if (aname.length() != bname.length()) {
					return aname.length() - bname.length();
				}
				return aname.compareTo(bname);
			}
		};
		
		DeobfuscatedClassEntryComparator = new Comparator<ClassEntry>() {
			@Override
			public int compare(ClassEntry a, ClassEntry b) {
				if (a instanceof ScoredClassEntry && b instanceof ScoredClassEntry) {
					return Float.compare(
						((ScoredClassEntry)b).getScore(),
						((ScoredClassEntry)a).getScore()
					);
				}
				return a.getName().compareTo(b.getName());
			}
		};
	}
	
	private ClassSelectionListener m_listener;
	private Comparator<ClassEntry> m_comparator;
	
	public ClassSelector(Comparator<ClassEntry> comparator) {
		m_comparator = comparator;
		
		// configure the tree control
		setRootVisible(false);
		setShowsRootHandles(false);
		setModel(null);
		
		// hook events
		addMouseListener(new MouseAdapter() {
			@Override
			public void mouseClicked(MouseEvent event) {
				if (m_listener != null && event.getClickCount() == 2) {
					// get the selected node
					TreePath path = getSelectionPath();
					if (path != null && path.getLastPathComponent() instanceof ClassSelectorClassNode) {
						ClassSelectorClassNode node = (ClassSelectorClassNode)path.getLastPathComponent();
						m_listener.onSelectClass(node.getClassEntry());
					}
				}
			}
		});
		
		// init defaults
		m_listener = null;
	}
	
	public void setListener(ClassSelectionListener val) {
		m_listener = val;
	}
	
	public void setClasses(Collection<ClassEntry> classEntries) {
		if (classEntries == null) {
			setModel(null);
			return;
		}
		
		// build the package names
		Map<String,ClassSelectorPackageNode> packages = Maps.newHashMap();
		for (ClassEntry classEntry : classEntries) {
			packages.put(classEntry.getPackageName(), null);
		}
		
		// sort the packages
		List<String> sortedPackageNames = Lists.newArrayList(packages.keySet());
		Collections.sort(sortedPackageNames, new Comparator<String>() {
			@Override
			public int compare(String a, String b) {
				// I can never keep this rule straight when writing these damn things...
				// a < b => -1, a == b => 0, a > b => +1
				
				String[] aparts = a.split("/");
				String[] bparts = b.split("/");
				for (int i = 0; true; i++) {
					if (i >= aparts.length) {
						return -1;
					} else if (i >= bparts.length) {
						return 1;
					}
					
					int result = aparts[i].compareTo(bparts[i]);
					if (result != 0) {
						return result;
					}
				}
			}
		});
		
		// create the root node and the package nodes
		DefaultMutableTreeNode root = new DefaultMutableTreeNode();
		for (String packageName : sortedPackageNames) {
			ClassSelectorPackageNode node = new ClassSelectorPackageNode(packageName);
			packages.put(packageName, node);
			root.add(node);
		}
		
		// put the classes into packages
		Multimap<String,ClassEntry> packagedClassEntries = ArrayListMultimap.create();
		for (ClassEntry classEntry : classEntries) {
			packagedClassEntries.put(classEntry.getPackageName(), classEntry);
		}
		
		// build the class nodes
		for (String packageName : packagedClassEntries.keySet()) {
			// sort the class entries
			List<ClassEntry> classEntriesInPackage = Lists.newArrayList(packagedClassEntries.get(packageName));
			Collections.sort(classEntriesInPackage, m_comparator);
			
			// create the nodes in order
			for (ClassEntry classEntry : classEntriesInPackage) {
				ClassSelectorPackageNode node = packages.get(packageName);
				node.add(new ClassSelectorClassNode(classEntry));
			}
		}
		
		// finally, update the tree control
		setModel(new DefaultTreeModel(root));
	}
	
	public ClassEntry getSelectedClass() {
		if (!isSelectionEmpty()) {
			Object selectedNode = getSelectionPath().getLastPathComponent();
			if (selectedNode instanceof ClassSelectorClassNode) {
				ClassSelectorClassNode classNode = (ClassSelectorClassNode)selectedNode;
				return classNode.getClassEntry();
			}
		}
		return null;
	}
	
	public String getSelectedPackage() {
		if (!isSelectionEmpty()) {
			Object selectedNode = getSelectionPath().getLastPathComponent();
			if (selectedNode instanceof ClassSelectorPackageNode) {
				ClassSelectorPackageNode packageNode = (ClassSelectorPackageNode)selectedNode;
				return packageNode.getPackageName();
			} else if (selectedNode instanceof ClassSelectorClassNode) {
				ClassSelectorClassNode classNode = (ClassSelectorClassNode)selectedNode;
				return classNode.getClassEntry().getPackageName();
			}
		}
		return null;
	}
	
	public Iterable<ClassSelectorPackageNode> packageNodes() {
		List<ClassSelectorPackageNode> nodes = Lists.newArrayList();
		DefaultMutableTreeNode root = (DefaultMutableTreeNode)getModel().getRoot();
		Enumeration<?> children = root.children();
		while (children.hasMoreElements()) {
			ClassSelectorPackageNode packageNode = (ClassSelectorPackageNode)children.nextElement();
			nodes.add(packageNode);
		}
		return nodes;
	}
	
	public Iterable<ClassSelectorClassNode> classNodes(ClassSelectorPackageNode packageNode) {
		List<ClassSelectorClassNode> nodes = Lists.newArrayList();
		Enumeration<?> children = packageNode.children();
		while (children.hasMoreElements()) {
			ClassSelectorClassNode classNode = (ClassSelectorClassNode)children.nextElement();
			nodes.add(classNode);
		}
		return nodes;
	}
	
	public void expandPackage(String packageName) {
		if (packageName == null) {
			return;
		}
		for (ClassSelectorPackageNode packageNode : packageNodes()) {
			if (packageNode.getPackageName().equals(packageName)) {
				expandPath(new TreePath(new Object[] {getModel().getRoot(), packageNode}));
				return;
			}
		}
	}
	
	public void expandAll() {
		for (ClassSelectorPackageNode packageNode : packageNodes()) {
			expandPath(new TreePath(new Object[] {getModel().getRoot(), packageNode}));
		}
	}
	
	public ClassEntry getFirstClass() {
		for (ClassSelectorPackageNode packageNode : packageNodes()) {
			for (ClassSelectorClassNode classNode : classNodes(packageNode)) {
				return classNode.getClassEntry();
			}
		}
		return null;
	}
	
	public ClassSelectorPackageNode getPackageNode(ClassEntry entry) {
		for (ClassSelectorPackageNode packageNode : packageNodes()) {
			if (packageNode.getPackageName().equals(entry.getPackageName())) {
				return packageNode;
			}
		}
		return null;
	}

	public ClassEntry getNextClass(ClassEntry entry) {
		boolean foundIt = false;
		for (ClassSelectorPackageNode packageNode : packageNodes()) {
			if (!foundIt) {
				// skip to the package with our target in it
				if (packageNode.getPackageName().equals(entry.getPackageName())) {
					for (ClassSelectorClassNode classNode : classNodes(packageNode)) {
						if (!foundIt) {
							if (classNode.getClassEntry().equals(entry)) {
								foundIt = true;
							}
						} else {
							// return the next class
							return classNode.getClassEntry();
						}
					}
				}
			} else {
				// return the next class
				for (ClassSelectorClassNode classNode : classNodes(packageNode)) {
					return classNode.getClassEntry();
				}
			}
		}
		return null;
	}

	public void setSelectionClass(ClassEntry classEntry) {
		expandPackage(classEntry.getPackageName());
		for (ClassSelectorPackageNode packageNode : packageNodes()) {
			for (ClassSelectorClassNode classNode : classNodes(packageNode)) {
				if (classNode.getClassEntry().equals(classEntry)) {
					setSelectionPath(new TreePath(new Object[] {getModel().getRoot(), packageNode, classNode}));
				}
			}
		}
	}
}