From d02055eb2323f2dfd26794424ef5e52f5e5e9894 Mon Sep 17 00:00:00 2001
From: modmuss50 <modmuss50@gmail.com>
Date: Wed, 3 Mar 2021 21:22:10 +0000
Subject: [PATCH] Fix custom sourceset not being correctly set for runconfig
 tasks (#357)

---
 .../loom/configuration/ide/RunConfig.java     |  2 ++
 .../fabricmc/loom/task/AbstractRunTask.java   | 28 +++----------------
 2 files changed, 6 insertions(+), 24 deletions(-)

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 3cabb33..2fc7990 100644
--- a/src/main/java/net/fabricmc/loom/configuration/ide/RunConfig.java
+++ b/src/main/java/net/fabricmc/loom/configuration/ide/RunConfig.java
@@ -57,6 +57,7 @@ public class RunConfig {
 	public String runDir;
 	public String vmArgs;
 	public String programArgs;
+	public SourceSet sourceSet;
 
 	public Element genRuns(Element doc) {
 		Element root = this.addXml(doc, "component", ImmutableMap.of("name", "ProjectRunConfigurationManager"));
@@ -197,6 +198,7 @@ public class RunConfig {
 		runConfig.ideaModuleName = getIdeaModuleName(project, sourceSet);
 		runConfig.runDirIdeaUrl = "file://$PROJECT_DIR$/" + runDir;
 		runConfig.runDir = runDir;
+		runConfig.sourceSet = sourceSet;
 
 		// Custom parameters
 		for (String progArg : settings.getProgramArgs()) {
diff --git a/src/main/java/net/fabricmc/loom/task/AbstractRunTask.java b/src/main/java/net/fabricmc/loom/task/AbstractRunTask.java
index b944f7a..8df8e80 100644
--- a/src/main/java/net/fabricmc/loom/task/AbstractRunTask.java
+++ b/src/main/java/net/fabricmc/loom/task/AbstractRunTask.java
@@ -32,31 +32,23 @@ import java.util.List;
 import java.util.function.Function;
 
 import org.gradle.api.Project;
-import org.gradle.api.plugins.JavaPlugin;
 import org.gradle.api.tasks.JavaExec;
 
-import net.fabricmc.loom.LoomGradleExtension;
 import net.fabricmc.loom.configuration.ide.RunConfig;
 
 public abstract class AbstractRunTask extends JavaExec {
-	private final Function<Project, RunConfig> configProvider;
-	private RunConfig config;
+	private final RunConfig config;
 
-	public AbstractRunTask(Function<Project, RunConfig> config) {
+	public AbstractRunTask(Function<Project, RunConfig> configProvider) {
 		super();
 		setGroup("fabric");
-		this.configProvider = config;
+		this.config = configProvider.apply(getProject());
 
-		setClasspath(getProject().getConfigurations().getByName(JavaPlugin.RUNTIME_CLASSPATH_CONFIGURATION_NAME));
-		classpath(this.getProject().getExtensions().getByType(LoomGradleExtension.class).getUnmappedModCollection());
+		setClasspath(config.sourceSet.getRuntimeClasspath());
 	}
 
 	@Override
 	public void exec() {
-		if (config == null) {
-			config = configProvider.apply(getProject());
-		}
-
 		List<String> argsSplit = new ArrayList<>();
 		String[] args = config.programArgs.split(" ");
 		int partPos = -1;
@@ -93,10 +85,6 @@ public abstract class AbstractRunTask extends JavaExec {
 
 	@Override
 	public void setWorkingDir(File dir) {
-		if (config == null) {
-			config = configProvider.apply(getProject());
-		}
-
 		if (!dir.exists()) {
 			dir.mkdirs();
 		}
@@ -106,19 +94,11 @@ public abstract class AbstractRunTask extends JavaExec {
 
 	@Override
 	public String getMain() {
-		if (config == null) {
-			config = configProvider.apply(getProject());
-		}
-
 		return config.mainClass;
 	}
 
 	@Override
 	public List<String> getJvmArgs() {
-		if (config == null) {
-			config = configProvider.apply(getProject());
-		}
-
 		List<String> superArgs = super.getJvmArgs();
 		List<String> args = new ArrayList<>(superArgs != null ? superArgs : Collections.emptyList());
 		args.addAll(Arrays.asList(config.vmArgs.split(" ")));