Add support for compiling multiple sources at once using JavacCompiler service.

Author: Col-E

recaf-core/src/main/java/software/coley/recaf/services/compile/JavacArguments.java

@@ -4,6 +4,9 @@
 import jakarta.annotation.Nullable;
 import software.coley.recaf.workspace.model.Workspace;
 
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -15,7 +18,7 @@
 public class JavacArguments {
 	// Primary inputs
 	private final String className;
-	private final String classSource;
+	private final Map<String, String> classSources;
 	// Options
 	private final String classPath;
 	private final int versionTarget;
@@ -45,8 +48,38 @@ public class JavacArguments {
 	public JavacArguments(@Nonnull String className, @Nonnull String classSource,
 						  @Nullable String classPath, int versionTarget, int downsampleTarget,
 						  boolean debugVariables, boolean debugLineNumbers, boolean debugSourceName) {
+		this(className, Collections.singletonMap(className, classSource),
+				classPath, versionTarget, downsampleTarget,
+				debugVariables, debugLineNumbers, debugSourceName);
+	}
+
+	/**
+	 * @param className
+	 * 		Internal name of the primary class being compiled.
+	 * @param classSources
+	 * 		Sources of all classes to compile, keyed by internal name.
+	 * @param classPath
+	 * 		Classpath to use with compiler.
+	 * @param versionTarget
+	 * 		Java version to target.
+	 * @param downsampleTarget
+	 * 		Java version to target via down sampling. Negative to disable downs sampling.
+	 * @param debugVariables
+	 * 		Debug flag to include variable info.
+	 * @param debugLineNumbers
+	 * 		Debug flag to include line number info.
+	 * @param debugSourceName
+	 * 		Debug flag to include source file name.
+	 */
+	public JavacArguments(@Nonnull String className, @Nonnull Map<String, String> classSources,
+						  @Nullable String classPath, int versionTarget, int downsampleTarget,
+						  boolean debugVariables, boolean debugLineNumbers, boolean debugSourceName) {
 		this.className = className;
-		this.classSource = classSource;
+		if (classSources.isEmpty())
+			throw new IllegalArgumentException("Class sources must not be empty");
+		if (!classSources.containsKey(className))
+			throw new IllegalArgumentException("Class sources must contain the primary class: " + className);
+		this.classSources = classSources;
 		this.classPath = classPath;
 		this.versionTarget = versionTarget;
 		this.downsampleTarget = downsampleTarget;
@@ -88,11 +121,19 @@ public String getClassName() {
 	}
 
 	/**
-	 * @return Source of the class.
+	 * @return Source of the primary class.
 	 */
 	@Nonnull
 	public String getClassSource() {
-		return classSource;
+		return Objects.requireNonNull(classSources.get(className));
+	}
+
+	/**
+	 * @return Sources of all classes to compile, keyed by internal name.
+	 */
+	@Nonnull
+	public Map<String, String> getClassSources() {
+		return classSources;
 	}
 
 	/**
@@ -150,14 +191,14 @@ public boolean equals(Object o) {
 		if (debugLineNumbers != other.debugLineNumbers) return false;
 		if (debugSourceName != other.debugSourceName) return false;
 		if (!className.equals(other.className)) return false;
-		if (!classSource.equals(other.classSource)) return false;
+		if (!classSources.equals(other.classSources)) return false;
 		return Objects.equals(classPath, other.classPath);
 	}
 
 	@Override
 	public int hashCode() {
 		int result = className.hashCode();
-		result = 31 * result + classSource.hashCode();
+		result = 31 * result + classSources.hashCode();
 		result = 31 * result + (classPath != null ? classPath.hashCode() : 0);
 		result = 31 * result + versionTarget;
 		result = 31 * result + (debugVariables ? 1 : 0);

recaf-core/src/main/java/software/coley/recaf/services/compile/JavacArgumentsBuilder.java

@@ -4,14 +4,22 @@
 import jakarta.annotation.Nullable;
 import software.coley.recaf.util.JavaVersion;
 
+import java.util.HashMap;
+import java.util.Map;
+
 /**
  * Builder for {@link JavacArguments}.
  *
  * @author Matt Coley
  */
 public final class JavacArgumentsBuilder {
+	// Primary class to compile.
 	private String className;
 	private String classSource;
+	// All classes to compile, keyed by internal name.
+	// Does not need to contain the primary class, that will be added automatically if provided.
+	private Map<String, String> classSources;
+	// Additional compiler arguments.
 	private String classPath = System.getProperty("java.class.path");
 	private int versionTarget = JavaVersion.get();
 	private int downsampleTarget = -1;
@@ -21,7 +29,7 @@ public final class JavacArgumentsBuilder {
 
 	/**
 	 * @param className
-	 * 		Internal name of the class being compiled.
+	 * 		Internal name of the primary class being compiled.
 	 *
 	 * @return Builder.
 	 */
@@ -33,7 +41,7 @@ public JavacArgumentsBuilder withClassName(@Nonnull String className) {
 
 	/**
 	 * @param classSource
-	 * 		Source of the class.
+	 * 		Source of the primary class to compile.
 	 *
 	 * @return Builder.
 	 */
@@ -43,6 +51,33 @@ public JavacArgumentsBuilder withClassSource(@Nonnull String classSource) {
 		return this;
 	}
 
+	/**
+	 * @param classSources
+	 * 		Sources of additional classes to compile, keyed by internal name.
+	 *
+	 * @return Builder.
+	 */
+	@Nonnull
+	public JavacArgumentsBuilder withClassSources(@Nonnull Map<String, String> classSources) {
+		// Skip if no additional sources are provided.
+		if (classSources.isEmpty())
+			return this;
+
+		// Add additional sources to the builder.
+		if (this.classSources == null)
+			this.classSources = new HashMap<>();
+		this.classSources.putAll(classSources);
+
+		// If the primary class is not set pick the first entry in the map as the primary class.
+		if (className == null) {
+			Map.Entry<String, String> entry = classSources.entrySet().iterator().next();
+			return withClassName(entry.getKey())
+					.withClassSource(entry.getValue());
+		}
+
+		return this;
+	}
+
 	/**
 	 * @param classPath
 	 * 		Classpath to use with compiler.
@@ -122,11 +157,20 @@ public JavacArgumentsBuilder withDebugSourceName(boolean debugSourceName) {
 	@Nonnull
 	public JavacArguments build() {
 		if (className == null)
-			throw new IllegalArgumentException("Class name must not be null");
-		if (classSource == null)
-			throw new IllegalArgumentException("Class source must not be null");
+			throw new IllegalArgumentException("Primary class name must not be null");
+
+		// Combine primary input with any additional sources.
+		Map<String, String> allClassSources = classSources == null ? new HashMap<>() : new HashMap<>(classSources);
+		if (classSource != null)
+			allClassSources.put(className, classSource);
+
+		// Validate that the primary class is present in the sources.
+		if (allClassSources.isEmpty())
+			throw new IllegalArgumentException("Class sources must not be empty");
+		if (!allClassSources.containsKey(className))
+			throw new IllegalArgumentException("Class sources must contain the primary class");
 
-		return new JavacArguments(className, classSource,
+		return new JavacArguments(className, allClassSources,
 				classPath, versionTarget, downsampleTarget,
 				debugVariables, debugLineNumbers, debugSourceName);
 	}

recaf-core/src/main/java/software/coley/recaf/services/compile/JavacCompiler.java

@@ -28,6 +28,8 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -96,11 +98,14 @@ public CompilerResult compile(@Nonnull JavacArguments arguments,
 		if (compiler == null)
 			return new CompilerResult(new IllegalStateException("Cannot load 'javac' compiler."));
 
-		String className = arguments.getClassName();
+		// Primary class to compile, plus any additional classes that are being compiled alongside it.
+		String className = arguments.getClassName(); // This is primarily used for logging, there's no other special treatment.
+		Map<String, String> classSources = arguments.getClassSources();
+		Set<String> sourceClassNames = classSources.keySet();
 
 		// Class input map
 		VirtualUnitMap unitMap = new VirtualUnitMap();
-		unitMap.addSource(className, arguments.getClassSource());
+		classSources.forEach(unitMap::addSource);
 
 		// Create a file manager to track files in-memory rather than on-disk
 		List<WorkspaceResource> virtualClassPath = workspace == null ?
@@ -111,8 +116,10 @@ public CompilerResult compile(@Nonnull JavacArguments arguments,
 		// Generate phantom classes if the workspace does not already have phantoms in it.
 		if (workspace != null && config.getGeneratePhantoms().getValue()
 				&& workspace.getSupportingResources().stream().noneMatch(resource -> resource instanceof GeneratedPhantomWorkspaceResource)) {
-			// Only scan the target class and any of its inner classes for content to fill in.
-			List<JvmClassInfo> classesToScan = workspace.findJvmClasses(c -> c.getName().equals(className) || c.isInnerClassOf(className)).stream()
+			// Only scan the target classes and any of their inner classes for content to fill in.
+			List<JvmClassInfo> classesToScan = workspace.findJvmClasses(c ->
+							sourceClassNames.contains(c.getName()) ||
+									sourceClassNames.stream().anyMatch(c::isInnerClassOf)).stream()
 					.map(p -> p.getValue().asJvmClass())
 					.collect(Collectors.toList());
 			if (!classesToScan.isEmpty()) {

recaf-core/src/test/java/software/coley/recaf/services/compile/JavacCompilerTest.java

@@ -184,4 +184,34 @@ public static void main(String[] args) {
 		assertEquals(0, result.getDiagnostics().size(), "There were unexpected diagnostic messages");
 		assertTrue(result.getCompilations().containsKey("HelloWorld"), "Class missing from compile map output");
 	}
+
+	@Test
+	void testJavacWithMultipleSources() {
+		// Specify multiple sources with a minimal dependency between them.
+		JavacArguments arguments = new JavacArgumentsBuilder()
+				.withClassSources(Map.of(
+						"HelloWorld",
+								"""
+								public class HelloWorld {
+									public static void main(String[] args) {
+										System.out.println(Helper.message());
+									}
+								}""",
+						// Supporting class
+						"Helper", """
+								public class Helper {
+									public static String message() {
+										return "Hello world";
+									}
+								}"""
+				))
+				.build();
+
+		// We should be able to compile both classes at once.
+		CompilerResult result = javac.compile(arguments, null, null);
+		assertTrue(result.wasSuccess(), "Result does not indicate success");
+		assertEquals(0, result.getDiagnostics().size(), "There were unexpected diagnostic messages");
+		assertTrue(result.getCompilations().containsKey("HelloWorld"), "Primary class missing from compile map output");
+		assertTrue(result.getCompilations().containsKey("Helper"), "Dependent class missing from compile map output");
+	}
 }