Create WorkspaceResource stream methods to combine jvm+jvm-versioned classes

Col-E · Col-E · commit 88e5f26c45cb · 2025-12-06T23:28:52.000-05:00
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/callgraph/CallGraph.java b/recaf-core/src/main/java/software/coley/recaf/services/callgraph/CallGraph.java
@@ -36,7 +36,6 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
-import java.util.stream.Stream;
 
 /**
  * Represents method calls as a navigable graph.
@@ -138,11 +137,9 @@ public void initialize() {
 		// Initialize asynchronously, and mark 'isReady' if completed successfully
 		CompletableFuture.runAsync(() -> {
 			for (WorkspaceResource resource : workspace.getAllResources(false)) {
-				Stream.concat(resource.jvmClassBundleStream(),
-						resource.getVersionedJvmClassBundles().values().stream()).forEach(bundle -> {
-					for (JvmClassInfo jvmClass : bundle.values()) {
+				resource.jvmAllClassBundleStreamRecursive().forEach(bundle -> {
+					for (JvmClassInfo jvmClass : bundle.values())
 						visit(jvmClass);
-					}
 				});
 			}
 		}, threadPool).whenComplete((unused, t) -> {
@@ -377,22 +374,18 @@ public Result<Resolution<JvmClassInfo, MethodMember>> resolve(int opcode, @Nonnu
 	@Override
 	public void onAddLibrary(@Nonnull Workspace workspace, @Nonnull WorkspaceResource library) {
 		// Visit all library classes
-		Stream.concat(library.jvmClassBundleStream(),
-				library.getVersionedJvmClassBundles().values().stream()).forEach(bundle -> {
-			for (JvmClassInfo jvmClass : bundle.values()) {
+		library.jvmAllClassBundleStreamRecursive().forEach(bundle -> {
+			for (JvmClassInfo jvmClass : bundle.values())
 				onNewClass(library, bundle, jvmClass);
-			}
 		});
 	}
 
 	@Override
 	public void onRemoveLibrary(@Nonnull Workspace workspace, @Nonnull WorkspaceResource library) {
 		// Remove all vertices from library
-		Stream.concat(library.jvmClassBundleStream(),
-				library.getVersionedJvmClassBundles().values().stream()).forEach(bundle -> {
-			for (JvmClassInfo jvmClass : bundle.values()) {
+		library.jvmAllClassBundleStreamRecursive().forEach(bundle -> {
+			for (JvmClassInfo jvmClass : bundle.values())
 				onRemoveClass(library, bundle, jvmClass);
-			}
 		});
 	}
 
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/compile/VirtualFileManager.java b/recaf-core/src/main/java/software/coley/recaf/services/compile/VirtualFileManager.java
@@ -52,7 +52,7 @@ public Iterable<JavaFileObject> list(@Nonnull Location location, @Nonnull String
 						name.indexOf('/', formatted.length()) == -1;
 			}
 			return () -> new ClassPathIterator(list.iterator(), virtualClasspath.stream()
-					.flatMap(resource -> resource.jvmClassBundleStream().flatMap(b -> b.entrySet().stream()))
+					.flatMap(resource -> resource.jvmClassBundleStreamRecursive().flatMap(b -> b.entrySet().stream()))
 					.filter(entry -> check.test(entry.getKey()))
 					.<JavaFileObject>map(entry -> new ResourceVirtualJavaFileObject(entry.getKey(),
 							entry.getValue().getBytecode(), JavaFileObject.Kind.CLASS))
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/decompile/vineflower/ClassSource.java b/recaf-core/src/main/java/software/coley/recaf/services/decompile/vineflower/ClassSource.java
@@ -40,9 +40,9 @@ protected DecompiledOutputSink getSink() {
 
 	@Override
 	public Entries getEntries() {
-		// TODO: Bug in QF/VF makes it so that 'addLibrary' doesn't yield inner info for a class provided with 'addSource'
+		// TODO: Bug in Vineflower makes it so that 'addLibrary' doesn't yield inner info for a class provided with 'addSource'
 		//  So for now until this is fixed upstream we will also supply inners here.
-		//  This will make QF/VF decompile each inner class separately as well, but its the best fix for now without
+		//  This will make Vineflower decompile each inner class separately as well, but its the best fix for now without
 		//  too much of a perf hit.
 		List<Entry> entries = new ArrayList<>();
 		entries.add(new Entry(targetInfo.getName(), Entry.BASE_VERSION));
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/decompile/vineflower/LibrarySource.java b/recaf-core/src/main/java/software/coley/recaf/services/decompile/vineflower/LibrarySource.java
@@ -1,13 +1,12 @@
 package software.coley.recaf.services.decompile.vineflower;
 
 import jakarta.annotation.Nonnull;
+import org.jetbrains.java.decompiler.main.extern.IContextSource;
 import software.coley.recaf.info.JvmClassInfo;
 import software.coley.recaf.workspace.model.Workspace;
-import software.coley.recaf.workspace.model.resource.WorkspaceResource;
 
 import java.util.Collections;
 import java.util.List;
-import java.util.stream.Collectors;
 
 /**
  * Full library source for Vineflower.
@@ -16,23 +15,23 @@
  * @author therathatter
  */
 public class LibrarySource extends BaseSource {
+	private final List<Entry> entries;
+
 	/**
+	 * @param entries
+	 * 		List of context entries in the given workspace.
 	 * @param workspace
 	 * 		Workspace to pull class files from.
 	 * @param targetInfo
 	 * 		Target class to decompile.
 	 */
-	protected LibrarySource(@Nonnull Workspace workspace, @Nonnull JvmClassInfo targetInfo) {
+	protected LibrarySource(@Nonnull List<IContextSource.Entry> entries, @Nonnull Workspace workspace, @Nonnull JvmClassInfo targetInfo) {
 		super(workspace, targetInfo);
+		this.entries = entries;
 	}
 
 	@Override
 	public Entries getEntries() {
-		List<Entry> entries = workspace.getAllResources(false).stream()
-				.flatMap(WorkspaceResource::jvmClassBundleStream)
-				.flatMap(c -> c.keySet().stream())
-				.map(className -> new Entry(className, Entry.BASE_VERSION))
-				.collect(Collectors.toList());
 		return new Entries(entries, Collections.emptyList(), Collections.emptyList());
 	}
 }
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/decompile/vineflower/VineflowerDecompiler.java b/recaf-core/src/main/java/software/coley/recaf/services/decompile/vineflower/VineflowerDecompiler.java
@@ -22,6 +22,7 @@ public class VineflowerDecompiler extends AbstractJvmDecompiler {
 	private final VineflowerConfig config;
 	private final IFernflowerLogger logger;
 	private final IResultSaver dummySaver = new DummyResultSaver();
+	private final WorkspaceEntriesCache workspaceEntriesCache = new WorkspaceEntriesCache();
 
 	/**
 	 * New Vineflower decompiler instance.
@@ -45,14 +46,12 @@ public DecompileResult decompileInternal(@Nonnull Workspace workspace, @Nonnull
 		try {
 			ClassSource source = new ClassSource(workspace, info);
 			fernflower.addSource(source);
-			fernflower.addLibrary(new LibrarySource(workspace, info));
+			fernflower.addLibrary(new LibrarySource(workspaceEntriesCache.getCachedEntries(workspace), workspace, info));
 			fernflower.decompileContext();
 
 			String decompiled = source.getSink().getDecompiledOutput().get();
-
-			if (decompiled == null || decompiled.isEmpty()) {
+			if (decompiled == null || decompiled.isEmpty())
 				return new DecompileResult(new IllegalStateException("Missing decompilation output"), config.getHash());
-			}
 
 			return new DecompileResult(decompiled, config.getHash());
 		} catch (Exception e) {
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/decompile/vineflower/WorkspaceEntriesCache.java b/recaf-core/src/main/java/software/coley/recaf/services/decompile/vineflower/WorkspaceEntriesCache.java
@@ -0,0 +1,79 @@
+package software.coley.recaf.services.decompile.vineflower;
+
+import jakarta.annotation.Nonnull;
+import org.jetbrains.java.decompiler.main.extern.IContextSource;
+import software.coley.recaf.info.JvmClassInfo;
+import software.coley.recaf.workspace.model.Workspace;
+import software.coley.recaf.workspace.model.WorkspaceModificationListener;
+import software.coley.recaf.workspace.model.bundle.JvmClassBundle;
+import software.coley.recaf.workspace.model.resource.ResourceJvmClassListener;
+import software.coley.recaf.workspace.model.resource.WorkspaceResource;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * Caches the list of {@link IContextSource.Entry} items in a workspace.
+ *
+ * @author Matt Coley
+ * @see LibrarySource
+ */
+public class WorkspaceEntriesCache implements WorkspaceModificationListener, ResourceJvmClassListener {
+	private List<IContextSource.Entry> cache;
+	private int lastWorkspaceHash;
+
+	@Override
+	public void onAddLibrary(@Nonnull Workspace workspace, @Nonnull WorkspaceResource library) {
+		resetCache();
+	}
+
+	@Override
+	public void onRemoveLibrary(@Nonnull Workspace workspace, @Nonnull WorkspaceResource library) {
+		resetCache();
+	}
+
+	@Override
+	public void onNewClass(@Nonnull WorkspaceResource resource, @Nonnull JvmClassBundle bundle, @Nonnull JvmClassInfo cls) {
+		resetCache();
+	}
+
+	@Override
+	public void onUpdateClass(@Nonnull WorkspaceResource resource, @Nonnull JvmClassBundle bundle, @Nonnull JvmClassInfo oldCls, @Nonnull JvmClassInfo newCls) {
+		resetCache();
+	}
+
+	@Override
+	public void onRemoveClass(@Nonnull WorkspaceResource resource, @Nonnull JvmClassBundle bundle, @Nonnull JvmClassInfo cls) {
+		resetCache();
+	}
+
+	/**
+	 * Clears the cache.
+	 */
+	private synchronized void resetCache() {
+		cache = null;
+	}
+
+	/**
+	 * @param workspace
+	 * 		Workspace to get cached entries of.
+	 *
+	 * @return List of all distinctly named entries for classes in the workspace.
+	 */
+	@Nonnull
+	public synchronized List<IContextSource.Entry> getCachedEntries(@Nonnull Workspace workspace) {
+		List<IContextSource.Entry> local = cache;
+		int workspaceHash = System.identityHashCode(workspace);
+		if (local == null || workspaceHash != lastWorkspaceHash) {
+			local = workspace.getAllResources(false).stream()
+					.flatMap(WorkspaceResource::jvmAllClassBundleStreamRecursive)
+					.flatMap(c -> c.keySet().stream())
+					.distinct()
+					.map(className -> new IContextSource.Entry(className, IContextSource.Entry.BASE_VERSION))
+					.collect(Collectors.toList());
+			cache = local;
+			lastWorkspaceHash = workspaceHash;
+		}
+		return local;
+	}
+}
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/mapping/MappingApplier.java b/recaf-core/src/main/java/software/coley/recaf/services/mapping/MappingApplier.java
@@ -115,7 +115,7 @@ public MappingResults applyToPrimaryResource(@Nonnull Mappings mappings) {
 		Mappings finalMappings = mappings;
 		ExecutorService service = ThreadUtil.phasingService(applierThreadPool);
 		WorkspaceResource resource = workspace.getPrimaryResource();
-		Stream.concat(resource.jvmClassBundleStream(), resource.versionedJvmClassBundleStream()).forEach(bundle -> {
+		resource.jvmAllClassBundleStreamRecursive().forEach(bundle -> {
 			bundle.forEach(classInfo -> {
 				service.execute(() -> dumpIntoResults(results, workspace, resource, bundle, classInfo, finalMappings));
 			});
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/mapping/gen/MappingGenerator.java b/recaf-core/src/main/java/software/coley/recaf/services/mapping/gen/MappingGenerator.java
@@ -73,10 +73,9 @@ public Mappings generate(@Nullable Workspace workspace,
 		if (workspace != null)
 			mappings.enableClassLookup(workspace);
 		SortedMap<String, ClassInfo> classMap = new TreeMap<>();
-		resource.versionedJvmClassBundleStream()
+		resource.jvmAllClassBundleStreamRecursive()
 				.flatMap(Bundle::stream)
-				.forEach(c -> classMap.put(c.getName(), c));
-		classMap.putAll(resource.getJvmClassBundle());
+				.forEach(c -> classMap.putIfAbsent(c.getName(), c));
 
 		// Pull a class, create mappings for its inheritance family, then remove those classes from the map.
 		// When the map is empty everything has been run through the mapping generation process.
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/phantom/JPhantomGenerator.java b/recaf-core/src/main/java/software/coley/recaf/services/phantom/JPhantomGenerator.java
@@ -76,7 +76,7 @@ public JPhantomGenerator(@Nonnull JPhantomGeneratorConfig config, @Nonnull Works
 	@Override
 	public GeneratedPhantomWorkspaceResource createPhantomsForWorkspace(@Nonnull Workspace workspace) throws PhantomGenerationException {
 		// Extract all JVM classes from workspace
-		Map<String, JvmClassInfo> classMap = workspace.getPrimaryResource().jvmClassBundleStream()
+		Map<String, JvmClassInfo> classMap = workspace.getPrimaryResource().jvmAllClassBundleStreamRecursive()
 				.flatMap(Bundle::stream)
 				.collect(Collectors.toMap(Info::getName, Function.identity()));
 
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/search/SearchService.java b/recaf-core/src/main/java/software/coley/recaf/services/search/SearchService.java
@@ -42,7 +42,6 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
-import java.util.stream.Stream;
 
 /**
  * Outline for running various searches.
@@ -207,7 +206,7 @@ private static void searchResource(@Nonnull Results results,
 
 		// Visit JVM content
 		if (jvmClassVisitor != null) {
-			Stream.concat(resource.jvmClassBundleStream(), resource.versionedJvmClassBundleStream()).forEach(bundle -> {
+			resource.jvmAllClassBundleStream().forEach(bundle -> {
 				BundlePathNode bundlePath = resourcePath.child(bundle);
 				for (JvmClassInfo classInfo : bundle) {
 					if (feedback.hasRequestedCancellation())
diff --git a/recaf-core/src/main/java/software/coley/recaf/services/transform/TransformationApplier.java b/recaf-core/src/main/java/software/coley/recaf/services/transform/TransformationApplier.java
@@ -135,7 +135,7 @@ public JvmTransformResult transformJvm(@Nonnull List<Class<? extends JvmClassTra
 				throw new TransformationException(message, t);
 			}
 		}
-		resource.jvmClassBundleStreamRecursive().forEach(bundle -> {
+		resource.jvmAllClassBundleStreamRecursive().forEach(bundle -> {
 			BundlePathNode bundlePathNode = resourcePath.child(bundle);
 			for (int pass = 1; pass <= getMaxPasses(); pass++) {
 				AtomicBoolean anyWorkDone = new AtomicBoolean(false);
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
@@ -142,7 +142,7 @@ default Stream<JvmClassBundle> jvmClassBundleStream() {
 	}
 
 	/**
-	 * @return Stream of all JVM class bundles in the resource, and in any embedded resources
+	 * @return Stream of all JVM class bundles in the resource, and in any embedded resources.
 	 */
 	@Nonnull
 	default Stream<JvmClassBundle> jvmClassBundleStreamRecursive() {
@@ -159,14 +159,34 @@ default Stream<VersionedJvmClassBundle> versionedJvmClassBundleStream() {
 	}
 
 	/**
-	 * @return Stream of all versioned JVM class bundles in the resource, and in any embedded resources
+	 * @return Stream of all versioned JVM class bundles in the resource, and in any embedded resources.
 	 */
 	@Nonnull
 	default Stream<VersionedJvmClassBundle> versionedJvmClassBundleStreamRecursive() {
 		return concat(versionedJvmClassBundleStream(), getEmbeddedResources().values().stream()
 				.flatMap(WorkspaceResource::versionedJvmClassBundleStreamRecursive));
 	}
 
+	/**
+	 * Use this if you don't want to manually concat the jvm and versioned-jvm bundle streams.
+	 *
+	 * @return Stream of all immediate and versioned JVM class bundles in the resource, and in any embedded resources.
+	 */
+	@Nonnull
+	default Stream<JvmClassBundle> jvmAllClassBundleStream() {
+		return concat(jvmClassBundleStream(), versionedJvmClassBundleStream());
+	}
+
+	/**
+	 * Use this if you don't want to manually concat the jvm and versioned-jvm bundle streams.
+	 *
+	 * @return Stream of all immediate and versioned JVM class bundles in the resource, and in any embedded resources.
+	 */
+	@Nonnull
+	default Stream<JvmClassBundle> jvmAllClassBundleStreamRecursive() {
+		return concat(jvmClassBundleStreamRecursive(), versionedJvmClassBundleStreamRecursive());
+	}
+
 	/**
 	 * @return Stream of all immediate Android class bundles in the resource.
 	 */
diff --git a/recaf-core/src/test/java/software/coley/recaf/services/compile/JavacCompilerTest.java b/recaf-core/src/test/java/software/coley/recaf/services/compile/JavacCompilerTest.java
@@ -10,14 +10,20 @@
 import org.objectweb.asm.commons.SimpleRemapper;
 import software.coley.recaf.RecafConstants;
 import software.coley.recaf.info.JvmClassInfo;
+import software.coley.recaf.info.StubFileInfo;
 import software.coley.recaf.info.builder.JvmClassInfoBuilder;
 import software.coley.recaf.test.TestBase;
 import software.coley.recaf.test.TestClassUtils;
 import software.coley.recaf.test.dummy.StringConsumer;
 import software.coley.recaf.util.JavaVersion;
+import software.coley.recaf.workspace.model.BasicWorkspace;
 import software.coley.recaf.workspace.model.Workspace;
+import software.coley.recaf.workspace.model.resource.WorkspaceFileResourceBuilder;
+import software.coley.recaf.workspace.model.resource.WorkspaceResource;
+import software.coley.recaf.workspace.model.resource.WorkspaceResourceBuilder;
 
 import java.io.IOException;
+import java.util.Map;
 
 import static org.junit.jupiter.api.Assertions.*;
 
@@ -166,5 +172,16 @@ public static void main(String[] args) {
 		result = javac.compile(arguments, workspace, null);
 		assertEquals(0, result.getDiagnostics().size(), "There were unexpected diagnostic messages");
 		assertTrue(result.getCompilations().containsKey("HelloWorld"), "Class missing from compile map output");
+
+		// Put it into a workspace but as an embedded resource and try again. Should still work.
+		WorkspaceResource resource = new WorkspaceResourceBuilder()
+				.withEmbeddedResources(Map.of("embed.jar", new WorkspaceFileResourceBuilder()
+						.withFileInfo(new StubFileInfo("embed.jar"))
+						.withJvmClassBundle(TestClassUtils.fromClasses(classInfo)).build()))
+				.build();
+		workspace = new BasicWorkspace(resource);
+		result = javac.compile(arguments, workspace, null);
+		assertEquals(0, result.getDiagnostics().size(), "There were unexpected diagnostic messages");
+		assertTrue(result.getCompilations().containsKey("HelloWorld"), "Class missing from compile map output");
 	}
 }
diff --git a/recaf-ui/src/main/java/software/coley/recaf/services/info/summary/builtin/AntiDecompilationSummarizer.java b/recaf-ui/src/main/java/software/coley/recaf/services/info/summary/builtin/AntiDecompilationSummarizer.java
@@ -169,7 +169,7 @@ private JvmTransformResult computeTransformations(@Nonnull Workspace workspace)
 
 	private int computeIllegalNames(@Nonnull WorkspaceResource resource) {
 		Set<JvmClassInfo> classesWithIllegalNames = Collections.newSetFromMap(new IdentityHashMap<>());
-		resource.jvmClassBundleStream().forEach(bundle -> {
+		resource.jvmAllClassBundleStreamRecursive().forEach(bundle -> {
 			bundle.forEach(cls -> {
 				if (ILLEGAL_NAME_FILTER.shouldMapClass(cls)) {
 					classesWithIllegalNames.add(cls);
diff --git a/recaf-ui/src/main/java/software/coley/recaf/services/info/summary/builtin/EntryPointSummarizer.java b/recaf-ui/src/main/java/software/coley/recaf/services/info/summary/builtin/EntryPointSummarizer.java
@@ -62,7 +62,7 @@ public boolean summarize(@Nonnull Workspace workspace,
 
 		// Visit JVM classes
 		int[] found = {0};
-		resource.jvmClassBundleStream().forEach(bundle -> {
+		resource.jvmAllClassBundleStreamRecursive().forEach(bundle -> {
 			bundle.forEach(cls -> {
 				List<MethodMember> entryMethods = cls.getMethods().stream()
 						.filter(this::isJvmEntry)

Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,7 @@ public Iterable<JavaFileObject> list(@Nonnull Location location, @Nonnull String`
`52`	`52`	`name.indexOf('/', formatted.length()) == -1;`
`53`	`53`	`}`
`54`	`54`	`return () -> new ClassPathIterator(list.iterator(), virtualClasspath.stream()`
`55`		`- .flatMap(resource -> resource.jvmClassBundleStream().flatMap(b -> b.entrySet().stream()))`
	`55`	`+ .flatMap(resource -> resource.jvmClassBundleStreamRecursive().flatMap(b -> b.entrySet().stream()))`
`56`	`56`	`.filter(entry -> check.test(entry.getKey()))`
`57`	`57`	`.<JavaFileObject>map(entry -> new ResourceVirtualJavaFileObject(entry.getKey(),`
`58`	`58`	`entry.getValue().getBytecode(), JavaFileObject.Kind.CLASS))`