Massively speed up inheritance graph population

recaf-core/src/main/java/software/coley/recaf/services/inheritance/ClassPathNodeProvider.java

@@ -0,0 +1,79 @@
+package software.coley.recaf.services.inheritance;
+
+import jakarta.annotation.Nonnull;
+import jakarta.annotation.Nullable;
+import software.coley.recaf.path.ClassPathNode;
+import software.coley.recaf.workspace.model.Workspace;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.stream.Stream;
+
+/**
+ * Provider of class path nodes.
+ *
+ * @author xDark
+ */
+sealed interface ClassPathNodeProvider {
+	/**
+	 * @param name
+	 * 		Class name to look up.
+	 *
+	 * @return Path node for the class with the given name, or {@code null} if no such class exists in the provider.
+	 */
+	@Nullable
+	ClassPathNode getNode(@Nonnull String name);
+
+	/**
+	 * Create a cached provider that contains all nodes from the workspace at the time of creation.
+	 *
+	 * @param workspace
+	 * 		Workspace to cache nodes from.
+	 *
+	 * @return Provider that caches all nodes from the workspace at the time of creation.
+	 */
+	static ClassPathNodeProvider.Cached cache(@Nonnull Workspace workspace) {
+		Stream<ClassPathNode> stream = workspace.classesStream();
+		Map<String, ClassPathNode> nodes = new HashMap<>(4096);
+		stream.forEach(classPathNode -> {
+			nodes.putIfAbsent(classPathNode.getValue().getName(), classPathNode);
+		});
+		return new Cached(Map.copyOf(nodes));
+	}
+
+	/**
+	 * Provider that looks up nodes directly from the workspace.
+	 * This is not recommended for repeated lookups, but it is useful for one-off lookups or when the workspace is expected to be changing frequently.
+	 *
+	 * @param workspace
+	 * 		Workspace to look up nodes from.
+	 */
+	record Live(@Nonnull Workspace workspace) implements ClassPathNodeProvider {
+		@Nullable
+		@Override
+		public ClassPathNode getNode(@Nonnull String name) {
+			return workspace.findClass(name);
+		}
+	}
+
+	/**
+	 * Provider that caches all nodes from the workspace at the time of creation.
+	 * This is recommended for repeated lookups, but it is not suitable for workspaces that are expected to be changing frequently.
+	 *
+	 * @param nodes
+	 * 		Map of class names to their corresponding path nodes. This map is expected to be immutable.
+	 *
+	 * @see #cache(Workspace)
+	 */
+	record Cached(@Nonnull Map<String, ClassPathNode> nodes) implements ClassPathNodeProvider {
+		int size() {
+			return nodes.size();
+		}
+
+		@Nullable
+		@Override
+		public ClassPathNode getNode(@Nonnull String name) {
+			return nodes.get(name);
+		}
+	}
+}

recaf-core/src/main/java/software/coley/recaf/services/inheritance/InheritanceGraph.java

@@ -49,6 +49,7 @@ public class InheritanceGraph {
 	private final Set<String> stubs = ConcurrentHashMap.newKeySet();
 	private final ListenerHost listener = new ListenerHost();
 	private final Workspace workspace;
+	private final ClassPathNodeProvider workspaceNodeProvider;
 
 	/**
 	 * Create an inheritance graph.
@@ -58,6 +59,7 @@ public class InheritanceGraph {
 	 */
 	public InheritanceGraph(@Nonnull Workspace workspace) {
 		this.workspace = workspace;
+		this.workspaceNodeProvider = new ClassPathNodeProvider.Live(workspace);
 
 		// Populate map lookups with the initial capacity of the number of classes in the workspace plus a buffer.
 		int classesInWorkspace = workspace.allResourcesStream(false /* dont count internal resource classes */)
@@ -111,10 +113,9 @@ private void refreshChildLookup() {
 		parentToChild.clear();
 
 		// Repopulate
-		workspace.findClasses(false, cls -> {
-			populateParentToChildLookup(cls);
-			return false;
-		});
+		ClassPathNodeProvider.Cached nodeProvider = ClassPathNodeProvider.cache(workspace);
+		Set<ClassInfo> visited = Collections.newSetFromMap(new IdentityHashMap<>(nodeProvider.size() + 1024 /* leeway */));
+		workspace.forEachClass(false, cls -> populateParentToChildLookup(cls, visited, nodeProvider));
 	}
 
 	/**
@@ -124,25 +125,39 @@ private void refreshChildLookup() {
 	 * 		Child class name.
 	 * @param parentName
 	 * 		Parent class name.
+	 * @param provider
+	 *      Node provider.
 	 */
-	private void populateParentToChildLookup(@Nonnull String name, @Nonnull String parentName) {
+	private void populateParentToChildLookup(@Nonnull String name, @Nonnull String parentName, @Nonnull ClassPathNodeProvider provider) {
 		parentToChild.computeIfAbsent(parentName, k -> ConcurrentHashMap.newKeySet()).add(name);
 
 		// Clear any cached relationships in the vertex and the parent vertex.
-		InheritanceVertex parentVertex = getVertex(parentName);
-		InheritanceVertex childVertex = getVertex(name);
+		InheritanceVertex parentVertex = getVertex(parentName, provider);
+		InheritanceVertex childVertex = getVertex(name, provider);
 		if (parentVertex != null) parentVertex.clearCachedVertices();
 		if (childVertex != null) childVertex.clearCachedVertices();
 	}
 
+	/**
+	 * Populate a references from the given child class to the parent class.
+	 *
+	 * @param name
+	 * 		Child class name.
+	 * @param parentName
+	 * 		Parent class name.
+	 */
+	private void populateParentToChildLookup(@Nonnull String name, @Nonnull String parentName) {
+		populateParentToChildLookup(name, parentName, workspaceNodeProvider);
+	}
+
 	/**
 	 * Populate all references from the given child class to its parents.
 	 *
 	 * @param info
 	 * 		Child class.
 	 */
 	private void populateParentToChildLookup(@Nonnull ClassInfo info) {
-		populateParentToChildLookup(info, Collections.newSetFromMap(new IdentityHashMap<>()));
+		populateParentToChildLookup(info, Collections.newSetFromMap(new IdentityHashMap<>()), workspaceNodeProvider);
 	}
 
 	/**
@@ -152,8 +167,10 @@ private void populateParentToChildLookup(@Nonnull ClassInfo info) {
 	 * 		Child class.
 	 * @param visited
 	 * 		Classes already visited in population.
+	 * @param provider
+	 *      Node provider.
 	 */
-	private void populateParentToChildLookup(@Nonnull ClassInfo info, @Nonnull Set<ClassInfo> visited) {
+	private void populateParentToChildLookup(@Nonnull ClassInfo info, @Nonnull Set<ClassInfo> visited, @Nonnull ClassPathNodeProvider provider) {
 		// Since we have observed this class to exist, we will remove the "stub" placeholder for this name.
 		stubs.remove(info.getName());
 
@@ -167,31 +184,43 @@ private void populateParentToChildLookup(@Nonnull ClassInfo info, @Nonnull Set<C
 
 		// Add direct parent
 		String name = info.getName();
-		InheritanceVertex vertex = getVertex(name);
+		InheritanceVertex vertex = getVertex(name, provider);
 		if (vertex != null)
 			vertex.clearCachedVertices();
 
 		String superName = info.getSuperName();
 		if (superName != null) {
-			populateParentToChildLookup(name, superName);
+			populateParentToChildLookup(name, superName, provider);
 
 			// Visit parent
-			InheritanceVertex superVertex = getVertex(superName);
+			InheritanceVertex superVertex = getVertex(superName, provider);
 			if (superVertex != null && !superVertex.isJavaLangObject() && !superVertex.isLoop())
-				populateParentToChildLookup(superVertex.getValue(), visited);
+				populateParentToChildLookup(superVertex.getValue(), visited, provider);
 		}
 
 		// Add direct interfaces
 		for (String itf : info.getInterfaces()) {
-			populateParentToChildLookup(name, itf);
+			populateParentToChildLookup(name, itf, provider);
 
 			// Visit interfaces
-			InheritanceVertex interfaceVertex = getVertex(itf);
+			InheritanceVertex interfaceVertex = getVertex(itf, provider);
 			if (interfaceVertex != null)
-				populateParentToChildLookup(interfaceVertex.getValue(), visited);
+				populateParentToChildLookup(interfaceVertex.getValue(), visited, provider);
 		}
 	}
 
+	/**
+	 * Populate all references from the given child class to its parents.
+	 *
+	 * @param info
+	 * 		Child class.
+	 * @param visited
+	 * 		Classes already visited in population.
+	 */
+	private void populateParentToChildLookup(@Nonnull ClassInfo info, @Nonnull Set<ClassInfo> visited) {
+		populateParentToChildLookup(info, visited, workspaceNodeProvider);
+	}
+
 	/**
 	 * Remove all references from the given child class to its parents.
 	 *
@@ -253,16 +282,18 @@ private Set<String> getDirectChildren(@Nonnull String parent) {
 	/**
 	 * @param name
 	 * 		Class name.
+	 * @param provider
+	 *      Node provider.
 	 *
 	 * @return Vertex in graph of class. {@code null} if no such class was found in the inputs.
 	 */
 	@Nullable
-	public InheritanceVertex getVertex(@Nonnull String name) {
+	private InheritanceVertex getVertex(@Nonnull String name, @Nonnull ClassPathNodeProvider provider) {
 		InheritanceVertex vertex = vertices.get(name);
 		if (vertex == null && !stubs.contains(name)) {
 			// Vertex does not exist and was not marked as a stub.
 			// We want to look up the vertex for the given class and figure out if its valid or needs to be stubbed.
-			InheritanceVertex provided = createVertex(name);
+			InheritanceVertex provided = createVertex(name, provider);
 			if (provided == STUB || provided == null) {
 				// Provider yielded either a stub OR no result. Discard it.
 				stubs.add(name);
@@ -275,6 +306,17 @@ public InheritanceVertex getVertex(@Nonnull String name) {
 		return vertex;
 	}
 
+	/**
+	 * @param name
+	 * 		Class name.
+	 *
+	 * @return Vertex in graph of class. {@code null} if no such class was found in the inputs.
+	 */
+	@Nullable
+	public InheritanceVertex getVertex(@Nonnull String name) {
+		return getVertex(name, workspaceNodeProvider);
+	}
+
 	/**
 	 * @param name
 	 * 		Class name.
@@ -441,11 +483,13 @@ public boolean isLibraryMethod(@Nonnull String name, @Nonnull String methodName,
 	 *
 	 * @param name
 	 * 		Internal class name.
+	 * @param provider
+	 *      Node provider.
 	 *
 	 * @return Vertex of class.
 	 */
 	@Nullable
-	private InheritanceVertex createVertex(@Nullable String name) {
+	private InheritanceVertex createVertex(@Nullable String name, @Nonnull ClassPathNodeProvider provider) {
 		// Edge case handling for 'java/lang/Object' doing a parent lookup.
 		// There is no parent, do not use STUB.
 		if (name == null)
@@ -456,7 +500,7 @@ private InheritanceVertex createVertex(@Nullable String name) {
 			return null;
 
 		// Find class in workspace, if not found yield stub.
-		ClassPathNode result = workspace.findClass(name);
+		ClassPathNode result = provider.getNode(name);
 		if (result == null)
 			return STUB;
 
@@ -577,7 +621,7 @@ public void onPostApply(@Nonnull Workspace workspace, @Nonnull MappingResults ma
 			mappingResults.getPreMappingPaths().forEach((name, path) -> {
 				// If we see a 'stub' from the vertex creator, we know it is no longer
 				// in the workspace and should be removed from our cache.
-				InheritanceVertex vertex = createVertex(name);
+				InheritanceVertex vertex = createVertex(name, workspaceNodeProvider);
 				if (vertex == STUB) {
 					vertices.remove(name);
 					parentToChild.remove(name);