diff --git a/recaf-core/src/main/java/software/coley/recaf/services/inheritance/ClassPathNodeProvider.java b/recaf-core/src/main/java/software/coley/recaf/services/inheritance/ClassPathNodeProvider.java new file mode 100644 index 000000000..4bde6a0b9 --- /dev/null +++ b/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 stream = workspace.classesStream(); + Map 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 nodes) implements ClassPathNodeProvider { + int size() { + return nodes.size(); + } + + @Nullable + @Override + public ClassPathNode getNode(@Nonnull String name) { + return nodes.get(name); + } + } +} diff --git a/recaf-core/src/main/java/software/coley/recaf/services/inheritance/InheritanceGraph.java b/recaf-core/src/main/java/software/coley/recaf/services/inheritance/InheritanceGraph.java index b0e87dd7c..8720bfe2f 100644 --- a/recaf-core/src/main/java/software/coley/recaf/services/inheritance/InheritanceGraph.java +++ b/recaf-core/src/main/java/software/coley/recaf/services/inheritance/InheritanceGraph.java @@ -49,6 +49,7 @@ public class InheritanceGraph { private final Set 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 visited = Collections.newSetFromMap(new IdentityHashMap<>(nodeProvider.size() + 1024 /* leeway */)); + workspace.forEachClass(false, cls -> populateParentToChildLookup(cls, visited, nodeProvider)); } /** @@ -124,17 +125,31 @@ 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. * @@ -142,7 +157,7 @@ private void populateParentToChildLookup(@Nonnull String name, @Nonnull String p * 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 visited) { + private void populateParentToChildLookup(@Nonnull ClassInfo info, @Nonnull Set 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 visited) { + populateParentToChildLookup(info, visited, workspaceNodeProvider); + } + /** * Remove all references from the given child class to its parents. * @@ -253,16 +282,18 @@ private Set 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); diff --git a/recaf-core/src/main/java/software/coley/recaf/workspace/model/Workspace.java b/recaf-core/src/main/java/software/coley/recaf/workspace/model/Workspace.java index 530b6d89e..6e7381958 100644 --- a/recaf-core/src/main/java/software/coley/recaf/workspace/model/Workspace.java +++ b/recaf-core/src/main/java/software/coley/recaf/workspace/model/Workspace.java @@ -23,12 +23,15 @@ import java.util.ArrayDeque; import java.util.ArrayList; +import java.util.Collection; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.NavigableMap; import java.util.Queue; import java.util.SortedSet; import java.util.TreeSet; +import java.util.function.Consumer; import java.util.function.Function; import java.util.function.Predicate; import java.util.stream.Collectors; @@ -187,26 +190,26 @@ default ClassPathNode findJvmClass(@Nonnull String name) { */ @Nullable default ClassPathNode findJvmClass(boolean includeInternal, @Nonnull String name) { - Queue resourceQueue = new ArrayDeque<>(getAllResources(includeInternal)); - while (!resourceQueue.isEmpty()) { - WorkspaceResource resource = resourceQueue.remove(); - - // Check JVM bundles for class by the given name - JvmClassInfo classInfo; - for (JvmClassBundle bundle : resource.jvmClassBundleStream().toList()) { - classInfo = bundle.get(name); - if (classInfo != null) - return PathNodes.classPath(this, resource, bundle, classInfo); - } - for (VersionedJvmClassBundle versionedBundle : resource.versionedJvmClassBundleStream().toList()) { - classInfo = versionedBundle.get(name); - if (classInfo != null) - return PathNodes.classPath(this, resource, versionedBundle, classInfo); + Queue> resourceQueue = new ArrayDeque<>(); + Collection resources = getAllResources(includeInternal); + do { + for (WorkspaceResource resource : resources) { + // Check JVM bundles for class by the given name + JvmClassInfo classInfo; + for (JvmClassBundle bundle : resource.jvmClassBundles()) { + classInfo = bundle.get(name); + if (classInfo != null) + return PathNodes.classPath(this, resource, bundle, classInfo); + } + for (VersionedJvmClassBundle versionedBundle : resource.getVersionedJvmClassBundles().values()) { + classInfo = versionedBundle.get(name); + if (classInfo != null) + return PathNodes.classPath(this, resource, versionedBundle, classInfo); + } + // Queue up embedded resources + resourceQueue.add(resource.getEmbeddedResources().values()); } - - // Queue up embedded resources - resourceQueue.addAll(resource.getEmbeddedResources().values()); - } + } while ((resources = resourceQueue.poll()) != null); return null; } @@ -241,24 +244,25 @@ default ClassPathNode findLatestVersionedJvmClass(@Nonnull String name) { @Nullable default ClassPathNode findVersionedJvmClass(@Nonnull String name, int version) { // Internal resources don't have versioned classes, so we won't iterate over those. - Queue resourceQueue = new ArrayDeque<>(getAllResources(false)); - while (!resourceQueue.isEmpty()) { - WorkspaceResource resource = resourceQueue.remove(); + Queue> resourceQueue = new ArrayDeque<>(); + Collection resources = getAllResources(false); + do { + for (WorkspaceResource resource : resources) { + // Check versioned bundles for class by the given name, in descending order from the given version. + NavigableMap versionedBundleMap = resource.getVersionedJvmClassBundles(); + Map.Entry entry = versionedBundleMap.floorEntry(version); + while (entry != null) { + VersionedJvmClassBundle versionedBundle = entry.getValue(); + JvmClassInfo classInfo = versionedBundle.get(name); + if (classInfo != null) + return PathNodes.classPath(this, resource, versionedBundle, classInfo); + entry = versionedBundleMap.floorEntry(entry.getKey() - 1); + } - // Check versioned bundles for class by the given name, in descending order from the given version. - NavigableMap versionedBundleMap = resource.getVersionedJvmClassBundles(); - Map.Entry entry = versionedBundleMap.floorEntry(version); - while (entry != null) { - VersionedJvmClassBundle versionedBundle = entry.getValue(); - JvmClassInfo classInfo = versionedBundle.get(name); - if (classInfo != null) - return PathNodes.classPath(this, resource, versionedBundle, classInfo); - entry = versionedBundleMap.floorEntry(entry.getKey() - 1); + // Queue up embedded resources. + resourceQueue.add(resource.getEmbeddedResources().values()); } - - // Queue up embedded resources. - resourceQueue.addAll(resource.getEmbeddedResources().values()); - } + } while ((resources = resourceQueue.poll()) != null); return null; } @@ -359,6 +363,11 @@ default SortedSet findClasses(boolean includeInternal, @Nonnull P return result; } + default void forEachClass(boolean includeInternal, @Nonnull Consumer consumer) { + forEachJvmClass(includeInternal, consumer); + forEachAndroidClass(consumer); + } + /** * @return Stream of all classes. */ @@ -514,6 +523,12 @@ default SortedSet findJvmClasses(boolean includeInternal, @Nonnul .collect(Collectors.toCollection(TreeSet::new)); } + default void forEachJvmClass(boolean includeInternal, @Nonnull Consumer consumer) { + jvmClassesStream(includeInternal) + .map(node -> (JvmClassInfo) node.getValue()) + .forEach(consumer); + } + /** * @param filter * Android class filter. @@ -527,6 +542,12 @@ default SortedSet findAndroidClasses(@Nonnull Predicate consumer) { + androidClassesStream() + .map(node -> (AndroidClassInfo) node.getValue()) + .forEach(consumer); + } + /** * @param name * File name. diff --git a/recaf-core/src/main/java/software/coley/recaf/workspace/model/resource/AgentServerRemoteVmResource.java b/recaf-core/src/main/java/software/coley/recaf/workspace/model/resource/AgentServerRemoteVmResource.java index d6f7fb35a..be433db5e 100644 --- a/recaf-core/src/main/java/software/coley/recaf/workspace/model/resource/AgentServerRemoteVmResource.java +++ b/recaf-core/src/main/java/software/coley/recaf/workspace/model/resource/AgentServerRemoteVmResource.java @@ -1,5 +1,6 @@ package software.coley.recaf.workspace.model.resource; +import com.google.common.collect.Iterables; import com.sun.tools.attach.VirtualMachine; import jakarta.annotation.Nonnull; import jakarta.annotation.Nullable; @@ -83,6 +84,12 @@ public Map getJvmClassloaderBundles() { return (Map) (Object) remoteBundleMap; } + @Nonnull + @Override + public Iterable jvmClassBundles() { + return Iterables.concat(super.jvmClassBundles(), new ArrayList<>(remoteBundleMap.values())); + } + @Nonnull @Override public Stream jvmClassBundleStream() { diff --git a/recaf-core/src/main/java/software/coley/recaf/workspace/model/resource/WorkspaceResource.java b/recaf-core/src/main/java/software/coley/recaf/workspace/model/resource/WorkspaceResource.java index 0efd7f4f9..e46eadf20 100644 --- a/recaf-core/src/main/java/software/coley/recaf/workspace/model/resource/WorkspaceResource.java +++ b/recaf-core/src/main/java/software/coley/recaf/workspace/model/resource/WorkspaceResource.java @@ -15,6 +15,7 @@ import software.coley.recaf.workspace.model.bundle.JvmClassBundle; import software.coley.recaf.workspace.model.bundle.VersionedJvmClassBundle; +import java.util.List; import java.util.Map; import java.util.NavigableMap; import java.util.stream.Stream; @@ -133,6 +134,14 @@ default boolean isEmbeddedResource() { return getContainingResource() != null; } + /** + * @return Iterable of all immediate JVM class bundles in the resource. + */ + @Nonnull + default Iterable jvmClassBundles() { + return List.of(getJvmClassBundle()); + } + /** * @return Stream of all immediate JVM class bundles in the resource. */