Allow configuring additional local sourceSets.

2 files changed, 45 insertions, 12 deletions

diff --git a/src/main/java/net/fabricmc/loom/LoomGradleExtension.java b/src/main/java/net/fabricmc/loom/LoomGradleExtension.java
index 3c93cf0d..c7344c4d 100644
--- a/src/main/java/net/fabricmc/loom/LoomGradleExtension.java
+++ b/src/main/java/net/fabricmc/loom/LoomGradleExtension.java
@@ -27,6 +27,7 @@ package net.fabricmc.loom;
 import java.io.File;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
@@ -42,11 +43,14 @@ import java.util.stream.Collectors;
 import com.google.gson.JsonObject;
 import org.cadixdev.lorenz.MappingSet;
 import org.cadixdev.mercury.Mercury;
+import org.gradle.api.Action;
 import org.gradle.api.Project;
 import org.gradle.api.artifacts.Configuration;
 import org.gradle.api.artifacts.Dependency;
 import org.gradle.api.file.ConfigurableFileCollection;
 import org.gradle.api.plugins.BasePluginConvention;
+import org.gradle.api.plugins.JavaPluginConvention;
+import org.gradle.api.tasks.SourceSet;
 import org.jetbrains.annotations.Nullable;
 
 import net.fabricmc.loom.api.decompilers.LoomDecompiler;
@@ -98,6 +102,12 @@ public class LoomGradleExtension {
 	private final LazyBool forge;
 	private Set<File> mixinMappings = Collections.synchronizedSet(new HashSet<>());
 	private final List<String> tasksBeforeRun = Collections.synchronizedList(new ArrayList<>());
+	public final List<Supplier<SourceSet>> forgeLocalMods = Collections.synchronizedList(new ArrayList<>(Arrays.asList(new Supplier<SourceSet>() {
+		@Override
+		public SourceSet get() {
+			return project.getConvention().getPlugin(JavaPluginConvention.class).getSourceSets().getByName("main");
+		}
+	})));
 
 	/**
 	 * Loom will generate a new genSources task (with a new name, based off of {@link LoomDecompiler#name()})
@@ -126,10 +136,28 @@ public class LoomGradleExtension {
 		return srcMercuryCache[id] != null ? srcMercuryCache[id] : (srcMercuryCache[id] = factory.get());
 	}
 
-	public void addTaskBeforeRun(String task) {
-		synchronized (this.tasksBeforeRun) {
-			this.tasksBeforeRun.add(task);
+	public void localMods(Action<SourceSetConsumer> action) {
+		if (!isForge()) {
+			throw new UnsupportedOperationException("Not running with Forge support.");
 		}
+
+		action.execute(new SourceSetConsumer());
+	}
+
+	public class SourceSetConsumer {
+		public void add(Object... sourceSets) {
+			for (Object sourceSet : sourceSets) {
+				if (sourceSet instanceof SourceSet) {
+					forgeLocalMods.add(() -> (SourceSet) sourceSet);
+				} else {
+					forgeLocalMods.add(() -> project.getConvention().getPlugin(JavaPluginConvention.class).getSourceSets().findByName(String.valueOf(forgeLocalMods)));
+				}
+			}
+		}
+	}
+
+	public void addTaskBeforeRun(String task) {
+		this.tasksBeforeRun.add(task);
 	}
 
 	public List<String> getTasksBeforeRun() {
diff --git a/src/main/java/net/fabricmc/loom/configuration/ide/RunConfig.java b/src/main/java/net/fabricmc/loom/configuration/ide/RunConfig.java
index 8553b18a..4b60be99 100644
--- a/src/main/java/net/fabricmc/loom/configuration/ide/RunConfig.java
+++ b/src/main/java/net/fabricmc/loom/configuration/ide/RunConfig.java
@@ -33,6 +33,8 @@ import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.UUID;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
@@ -45,7 +47,6 @@ import com.google.gson.JsonObject;
 import org.apache.commons.io.FileUtils;
 import org.apache.commons.io.IOUtils;
 import org.gradle.api.Project;
-import org.gradle.api.plugins.JavaPluginConvention;
 import org.gradle.api.tasks.SourceSet;
 import org.gradle.plugins.ide.eclipse.model.EclipseModel;
 import org.w3c.dom.Document;
@@ -156,16 +157,20 @@ public class RunConfig {
 		}
 
 		if (extension.isForge()) {
-			SourceSet main = project.getConvention().getPlugin(JavaPluginConvention.class).getSourceSets().getByName("main");
+			List<String> modClasses = new ArrayList<>();
 
-			String modClasses = Stream.concat(
-					Stream.of(main.getOutput().getResourcesDir().getAbsolutePath()),
-					StreamSupport.stream(main.getOutput().getClassesDirs().spliterator(), false)
-							.map(File::getAbsolutePath)
-			).map(s -> "loom%%" + s)
-					.collect(Collectors.joining(File.pathSeparator));
+			for (Supplier<SourceSet> sourceSetSupplier : extension.forgeLocalMods) {
+				SourceSet sourceSet = sourceSetSupplier.get();
+				String sourceSetName = sourceSet.getName() + "_" + UUID.randomUUID().toString().replace("-", "").substring(0, 7);
 
-			runConfig.envVariables.put("MOD_CLASSES", modClasses);
+				Stream.concat(
+						Stream.of(sourceSet.getOutput().getResourcesDir().getAbsolutePath()),
+						StreamSupport.stream(sourceSet.getOutput().getClassesDirs().spliterator(), false)
+								.map(File::getAbsolutePath)
+				).map(s -> sourceSetName + "%%" + s).collect(Collectors.toCollection(() -> modClasses));
+			}
+
+			runConfig.envVariables.put("MOD_CLASSES", String.join(File.pathSeparator, modClasses));
 		}
 
 		if (extension.getLoaderLaunchMethod().equals("launchwrapper")) {