Sequential and cancelable stats reloads

enigma-swing/src/main/java/org/quiltmc/enigma/gui/ClassSelector.java

@@ -6,6 +6,7 @@
 import org.quiltmc.enigma.gui.node.SortedMutableTreeNode;
 import org.quiltmc.enigma.gui.util.GuiUtil;
 import org.quiltmc.enigma.api.translation.representation.entry.ClassEntry;
+import org.quiltmc.enigma.util.Utils;
 
 import javax.annotation.Nullable;
 import javax.swing.JTree;
@@ -18,6 +19,8 @@
 import java.util.Collection;
 import java.util.Comparator;
 import java.util.List;
+import java.util.concurrent.Future;
+import java.util.function.Supplier;
 
 public class ClassSelector extends JTree {
 	public static final Comparator<ClassEntry> DEOBF_CLASS_COMPARATOR = Comparator.comparing(ClassEntry::getFullName);
@@ -295,12 +298,29 @@ public void reload() {
 	 * On completion, the class's stats icon will be updated.
 	 *
 	 * @param classEntry the class to reload stats for
+	 *
+	 * @return a future whose completion indicates that all asynchronous work has finished
 	 */
-	public void reloadStats(ClassEntry classEntry) {
+	public Future<?> reloadStats(ClassEntry classEntry) {
+		return this.reloadStats(classEntry, Utils.SUPPLY_FALSE);
+	}
+
+	/**
+	 * Requests an asynchronous reload of the stats for the given class.
+	 * On completion, the class's stats icon will be updated.
+	 *
+	 * @param classEntry   the class to reload stats for
+	 * @param shouldCancel a supplier that may be used to cancel asynchronous work if it returns
+	 *                     {@code true} before the work has started
+	 *
+	 * @return a future whose completion indicates that no asynchronous work remains, whether
+	 * because it was canceled using the passed {@code shouldCancel} method or because it finished normally
+	 */
+	public Future<?> reloadStats(ClassEntry classEntry, Supplier<Boolean> shouldCancel) {
 		ClassSelectorClassNode node = this.packageManager.getClassNode(classEntry);
-		if (node != null) {
-			node.reloadStats(this.controller.getGui(), this, true);
-		}
+		return node == null
+				? Utils.DUMMY_FUTURE
+				: node.reloadStats(this.controller.getGui(), this, true, shouldCancel);
 	}
 
 	public interface ClassSelectionListener {

enigma-swing/src/main/java/org/quiltmc/enigma/gui/Gui.java

@@ -68,7 +68,12 @@
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.IntFunction;
+import java.util.stream.Stream;
 
 public class Gui {
 	private final MainWindow mainWindow;
@@ -106,6 +111,20 @@ public class Gui {
 
 	private final boolean testEnvironment;
 
+	/**
+	 * Executor for {@link #reloadStats(ClassEntry, boolean) reloadStats} work.
+	 *
+	 * <p> Executes all work from one call to {@link #reloadStats(ClassEntry, boolean) reloadStats}
+	 * before starting work for the next call.
+	 * Fixes <a href="https://github.com/QuiltMC/enigma/issues/271">#271</a>.
+	 */
+	private final Executor reloadStatsExecutor = Executors.newSingleThreadExecutor();
+	/**
+	 * Setting this to true cancels unstarted work from the last call to
+	 * {@link #reloadStats(ClassEntry, boolean) reloadStats}.
+	 */
+	private AtomicBoolean priorReloadStatsCanceler = new AtomicBoolean(false);
+
 	public Gui(EnigmaProfile profile, Set<EditableType> editableTypes, boolean testEnvironment) {
 		this.dockerManager = new DockerManager(this);
 		this.mainWindow = new MainWindow(this, Enigma.NAME);
@@ -609,24 +628,42 @@ public void moveClassTree(ClassEntry classEntry, boolean updateSwingState, boole
 
 	/**
 	 * Reloads stats for the provided class in all selectors.
+	 *
 	 * @param classEntry the class to reload
 	 * @param propagate whether to also reload ancestors of the class
 	 */
 	public void reloadStats(ClassEntry classEntry, boolean propagate) {
+		this.priorReloadStatsCanceler.set(true);
+		final AtomicBoolean currentReloadCanceler = new AtomicBoolean(false);
+		this.priorReloadStatsCanceler = currentReloadCanceler;
+
 		List<ClassEntry> toUpdate = new ArrayList<>();
 		toUpdate.add(classEntry);
 		if (propagate) {
-			Collection<ClassEntry> parents = this.controller.getProject().getJarIndex().getIndex(InheritanceIndex.class).getAncestors(classEntry);
+			Collection<ClassEntry> parents = this.controller.getProject().getJarIndex().getIndex(InheritanceIndex.class)
+					.getAncestors(classEntry);
 			toUpdate.addAll(parents);
 		}
 
-		for (Docker value : this.dockerManager.getDockers()) {
-			if (value instanceof ClassesDocker docker) {
-				for (ClassEntry entry : toUpdate) {
-					docker.getClassSelector().reloadStats(entry);
-				}
-			}
-		}
+		final List<Runnable> currentReloads = this.dockerManager.getDockers().stream()
+				.flatMap(docker -> docker instanceof ClassesDocker classes ? Stream.of(classes) : Stream.empty())
+				.flatMap(docker -> toUpdate.stream().<Runnable>map(updating -> () -> {
+					try {
+						docker.getClassSelector().reloadStats(updating, currentReloadCanceler::get).get();
+					} catch (InterruptedException | ExecutionException e) {
+						throw new RuntimeException(e);
+					}
+				}))
+				.toList();
+
+		this.reloadStatsExecutor.execute(() -> CompletableFuture
+				.allOf(
+					currentReloads.stream()
+						.map(CompletableFuture::runAsync)
+						.toArray(CompletableFuture[]::new)
+				)
+				.join()
+		);
 	}
 
 	public SearchDialog getSearchDialog() {

enigma-swing/src/main/java/org/quiltmc/enigma/gui/node/ClassSelectorClassNode.java

@@ -1,18 +1,23 @@
 package org.quiltmc.enigma.gui.node;
 
 import org.quiltmc.enigma.api.ProgressListener;
+import org.quiltmc.enigma.api.stats.ProjectStatsResult;
 import org.quiltmc.enigma.gui.ClassSelector;
 import org.quiltmc.enigma.gui.Gui;
 import org.quiltmc.enigma.gui.config.Config;
 import org.quiltmc.enigma.gui.util.GuiUtil;
 import org.quiltmc.enigma.api.stats.StatsGenerator;
 import org.quiltmc.enigma.api.translation.representation.entry.ClassEntry;
+import org.quiltmc.enigma.util.Utils;
 
 import javax.swing.SwingUtilities;
 import javax.swing.SwingWorker;
 import javax.swing.tree.DefaultTreeCellRenderer;
 import javax.swing.tree.TreeNode;
 import java.util.Comparator;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.function.Supplier;
 
 public class ClassSelectorClassNode extends SortedMutableTreeNode {
 	private final ClassEntry obfEntry;
@@ -37,47 +42,82 @@ public ClassEntry getDeobfEntry() {
 	 * Reloads the stats for this class node and updates the icon in the provided class selector.
 	 * Exits if no project is open.
 	 *
-	 * @param gui the current gui instance
-	 * @param selector the class selector to reload on
+	 * @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
+	 *
+	 * @return a future whose completion indicates that all asynchronous work has finished
 	 */
-	public void reloadStats(Gui gui, ClassSelector selector, boolean updateIfPresent) {
+	public Future<?> reloadStats(Gui gui, ClassSelector selector, boolean updateIfPresent) {
+		return this.reloadStats(gui, selector, updateIfPresent, Utils.SUPPLY_FALSE);
+	}
+
+	/**
+	 * Reloads the stats for this class node and updates the icon in the provided class selector.
+	 * Exits if no project is open.
+	 *
+	 * @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
+	 * @param shouldCancel    a supplier that may be used to cancel asynchronous work if it returns
+	 *                        {@code true} before the work has started
+	 *
+	 * @return a future whose completion indicates that no asynchronous work remains, whether
+	 * because it was canceled using the passed {@code shouldCancel} method or because it finished normally
+	 */
+	public Future<?> reloadStats(Gui gui, ClassSelector selector, boolean updateIfPresent, Supplier<Boolean> shouldCancel) {
 		StatsGenerator generator = gui.getController().getStatsGenerator();
 		if (generator == null) {
-			return;
+			return Utils.DUMMY_FUTURE;
 		}
 
-		SwingWorker<ClassSelectorClassNode, Void> iconUpdateWorker = new SwingWorker<>() {
+		SwingWorker<ProjectStatsResult, Void> iconUpdateWorker = new SwingWorker<>() {
 			@Override
-			protected ClassSelectorClassNode doInBackground() {
-				var parameters = Config.stats().createIconGenParameters(gui.getEditableStatTypes());
+			protected ProjectStatsResult doInBackground() {
+				if (shouldCancel.get()) {
+					return null;
+				} else {
+					var parameters = Config.stats().createIconGenParameters(gui.getEditableStatTypes());
 
-				if (generator.getResultNullable(parameters) == null && generator.getOverallProgress() == null) {
-					generator.generate(ProgressListener.createEmpty(), parameters);
-				} else if (updateIfPresent) {
-					generator.generate(ProgressListener.createEmpty(), ClassSelectorClassNode.this.getObfEntry(), parameters);
+					if (generator.getResultNullable(parameters) == null && generator.getOverallProgress() == null) {
+						return generator.generate(ProgressListener.createEmpty(), parameters);
+					} else if (updateIfPresent) {
+						return generator.generate(ProgressListener.createEmpty(), ClassSelectorClassNode.this.getObfEntry(), parameters);
+					} else {
+						return null;
+					}
 				}
-
-				return ClassSelectorClassNode.this;
 			}
 
 			@Override
 			public void done() {
-				try {
-					var parameters = Config.stats().createIconGenParameters(gui.getEditableStatTypes());
-					((DefaultTreeCellRenderer) selector.getCellRenderer()).setIcon(GuiUtil.getDeobfuscationIcon(generator.getResultNullable(parameters), ClassSelectorClassNode.this.getObfEntry()));
-				} catch (NullPointerException ignored) {
-					// do nothing. this seems to be a race condition, likely a bug in FlatLAF caused by us suppressing the default tree icons
-					// ignoring this error should never cause issues since it only occurs at startup
+				if (!shouldCancel.get()) {
+					final ProjectStatsResult result;
+					try {
+						result = this.get();
+					} catch (ExecutionException | InterruptedException e) {
+						throw new RuntimeException(e);
+					}
+
+					if (result != null) {
+						try {
+							((DefaultTreeCellRenderer) selector.getCellRenderer()).setIcon(GuiUtil.getDeobfuscationIcon(result, ClassSelectorClassNode.this.getObfEntry()));
+						} catch (NullPointerException ignored) {
+							// do nothing. this seems to be a race condition, likely a bug in FlatLAF caused by us suppressing the default tree icons
+							// ignoring this error should never cause issues since it only occurs at startup
+						}
+
+						SwingUtilities.invokeLater(() -> selector.reload(ClassSelectorClassNode.this, false));
+					}
 				}
-
-				SwingUtilities.invokeLater(() -> selector.reload(ClassSelectorClassNode.this, false));
 			}
 		};
 
 		if (Config.main().features.enableClassTreeStatIcons.value()) {
 			SwingUtilities.invokeLater(iconUpdateWorker::execute);
 		}
+
+		return iconUpdateWorker;
 	}
 
 	@Override

enigma/src/main/java/org/quiltmc/enigma/api/stats/ProjectStatsResult.java

@@ -9,12 +9,13 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 public class ProjectStatsResult implements StatsProvider {
 	private final EnigmaProject project;
 
 	private final Map<String, List<StatsResult>> packageToClasses = new HashMap<>();
-	private final Map<ClassEntry, StatsResult> stats = new HashMap<>();
+	private final Map<ClassEntry, StatsResult> stats = new ConcurrentHashMap<>();
 	private final Map<String, StatsResult> packageStats = new HashMap<>();
 
 	private StatsResult overall;

enigma/src/main/java/org/quiltmc/enigma/util/Utils.java

@@ -2,6 +2,7 @@
 
 import com.google.common.io.CharStreams;
 
+import javax.annotation.Nonnull;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
@@ -16,12 +17,43 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Properties;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Lock;
 import java.util.function.Supplier;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipFile;
 
 public class Utils {
+	public static final Future<Void> DUMMY_FUTURE = new Future<>() {
+		@Override
+		public boolean cancel(boolean mayInterruptIfRunning) {
+			return false;
+		}
+
+		@Override
+		public boolean isCancelled() {
+			return false;
+		}
+
+		@Override
+		public boolean isDone() {
+			return true;
+		}
+
+		@Override
+		public Void get() {
+			return null;
+		}
+
+		@Override
+		public Void get(long timeout, @Nonnull TimeUnit unit) {
+			return null;
+		}
+	};
+
+	public static final Supplier<Boolean> SUPPLY_FALSE = () -> false;
+
 	public static String readStreamToString(InputStream in) throws IOException {
 		return CharStreams.toString(new InputStreamReader(in, StandardCharsets.UTF_8));
 	}