More work

src/main/java/net/fabricmc/loom/api/remapping/RemapperContext.java

@@ -26,10 +26,22 @@
 
 import org.objectweb.asm.commons.Remapper;
 
+/**
+ * Context for a {@link RemapperExtension}.
+ */
 public interface RemapperContext {
-	Remapper getRemapper();
+	/**
+	 * @return The {@link Remapper} instance
+	 */
+	Remapper remapper();
 
+	/**
+	 * @return the source namespace
+	 */
 	String sourceNamespace();
 
+	/**
+	 * @return the target namespace
+	 */
 	String targetNamespace();
 }

src/main/java/net/fabricmc/loom/api/remapping/RemapperExtension.java

@@ -26,11 +26,35 @@
 
 import javax.inject.Inject;
 
+import org.jetbrains.annotations.ApiStatus;
 import org.objectweb.asm.ClassVisitor;
 
+/**
+ * A remapper extension can be used to add extra processing to the remapping process.
+ *
+ * <p>Implementations of RemapperExtension's are subject to the following constraints:
+ * <ul>
+ *     <li>Do not implement {@link #getParameters()} in your class, the method will be implemented by Gradle.</li>
+ *     <li>Constructors must be annotated with {@link Inject}.</li>
+ * </ul>
+ *
+ * @param <T> Parameter type for the extension. Should be {@link RemapperParameters.None} if the action does not have parameters.
+ */
 public interface RemapperExtension<T extends RemapperParameters> {
+	/**
+	 * @return The parameters for this extension.
+	 */
+	@ApiStatus.NonExtendable
 	@Inject
 	T getParameters();
 
+	/**
+	 * Return a {@link ClassVisitor} that will be used when remapping the given class.
+	 *
+	 * @param className The name of the class being remapped
+	 * @param remapperContext The remapper context
+	 * @param classVisitor The parent class visitor
+	 * @return A {@link ClassVisitor} that will be used when remapping the given class, or the given {@code classVisitor} if no extra processing is required for this class.
+	 */
 	ClassVisitor insertVisitor(String className, RemapperContext remapperContext, ClassVisitor classVisitor);
 }

src/main/java/net/fabricmc/loom/api/remapping/RemapperParameters.java

@@ -24,5 +24,14 @@
 
 package net.fabricmc.loom.api.remapping;
 
+/**
+ * Marker interface for parameter objects to {@link RemapperExtension}s.
+ *
+ * <p>Design based off of Gradle's {@link org.gradle.workers.WorkParameters}.
+ */
 public interface RemapperParameters {
+	final class None implements RemapperParameters {
+		private None() {
+		}
+	}
 }

src/main/java/net/fabricmc/loom/api/remapping/TinyRemapperExtension.java

@@ -0,0 +1,68 @@
+/*
+ * This file is part of fabric-loom, licensed under the MIT License (MIT).
+ *
+ * Copyright (c) 2023 FabricMC
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package net.fabricmc.loom.api.remapping;
+
+import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.Nullable;
+
+import net.fabricmc.tinyremapper.TinyRemapper;
+
+/**
+ * A remapper extension, that has direct access to the TinyRemapper APIs.
+ *
+ * <p>This API is not stable and may change without notice.
+ */
+@ApiStatus.Experimental
+public interface TinyRemapperExtension {
+	/**
+	 * See: {@link TinyRemapper.Builder#extraAnalyzeVisitor(TinyRemapper.AnalyzeVisitorProvider)}.
+	 *
+	 * @return A {@link TinyRemapper.AnalyzeVisitorProvider} or {@code null}.
+	 */
+	@Nullable
+	default TinyRemapper.AnalyzeVisitorProvider getAnalyzeVisitorProvider() {
+		return null;
+	}
+
+	/**
+	 * See: {@link TinyRemapper.Builder#extraPreApplyVisitor(TinyRemapper.ApplyVisitorProvider)}.
+	 *
+	 * @return A {@link TinyRemapper.ApplyVisitorProvider} or {@code null}.
+	 */
+	@Nullable
+	default TinyRemapper.ApplyVisitorProvider getPreApplyVisitor() {
+		return null;
+	}
+
+	/**
+	 * See: {@link TinyRemapper.Builder#extraPostApplyVisitor(TinyRemapper.ApplyVisitorProvider)}.
+	 *
+	 * @return A {@link TinyRemapper.ApplyVisitorProvider} or {@code null}.
+	 */
+	@Nullable
+	default TinyRemapper.ApplyVisitorProvider getPostApplyVisitor() {
+		return null;
+	}
+}

src/main/java/net/fabricmc/loom/configuration/mods/ModProcessor.java

@@ -159,7 +159,7 @@ private void remapJars(List<ModDependency> remapList) throws IOException {
 		}
 
 		for (RemapperExtensionHolder holder : extension.getRemapperExtensions().get()) {
-			holder.apply(builder);
+			holder.apply(builder, fromM, toM, project.getObjects());
 		}
 
 		final TinyRemapper remapper = builder.build();

src/main/java/net/fabricmc/loom/extension/RemapperExtensionHolder.java

@@ -26,15 +26,18 @@
 
 import javax.inject.Inject;
 
+import org.gradle.api.model.ObjectFactory;
 import org.gradle.api.provider.Property;
 import org.gradle.api.tasks.Input;
 import org.gradle.api.tasks.Nested;
+import org.jetbrains.annotations.Nullable;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.commons.Remapper;
 
 import net.fabricmc.loom.api.remapping.RemapperContext;
 import net.fabricmc.loom.api.remapping.RemapperExtension;
 import net.fabricmc.loom.api.remapping.RemapperParameters;
+import net.fabricmc.loom.api.remapping.TinyRemapperExtension;
 import net.fabricmc.tinyremapper.TinyRemapper;
 import net.fabricmc.tinyremapper.api.TrClass;
 
@@ -54,52 +57,60 @@ public RemapperParameters getRemapperParameters() {
 		return remapperParameters;
 	}
 
-	public void apply(TinyRemapper.Builder tinyRemapperBuilder) {
-		RemapperExtension<?> remapperExtension = newInstance();
+	public void apply(TinyRemapper.Builder tinyRemapperBuilder, String sourceNamespace, String targetNamespace, ObjectFactory objectFactory) {
+		final RemapperExtension<?> remapperExtension = newInstance(objectFactory);
 
-		tinyRemapperBuilder.extraPostApplyVisitor(new RemapperExtensionImpl(remapperExtension));
+		tinyRemapperBuilder.extraPostApplyVisitor(new RemapperExtensionImpl(remapperExtension, sourceNamespace, targetNamespace));
 
-		if (remapperExtension instanceof TinyRemapper.AnalyzeVisitorProvider analyzeVisitorProvider) {
-			tinyRemapperBuilder.extraAnalyzeVisitor(analyzeVisitorProvider);
-		}
+		if (remapperExtension instanceof TinyRemapperExtension tinyRemapperExtension) {
+			if (tinyRemapperExtension.getAnalyzeVisitorProvider() != null) {
+				tinyRemapperBuilder.extraAnalyzeVisitor(tinyRemapperExtension.getAnalyzeVisitorProvider());
+			}
+
+			if (tinyRemapperExtension.getPreApplyVisitor() != null) {
+				tinyRemapperBuilder.extraPreApplyVisitor(tinyRemapperExtension.getPreApplyVisitor());
+			}
 
-		if (remapperExtension instanceof TinyRemapper.ApplyVisitorProvider applyVisitorProvider) {
-			// TODO allow having a pre apply visitor?
-			tinyRemapperBuilder.extraPostApplyVisitor(applyVisitorProvider);
+			if (tinyRemapperExtension.getPostApplyVisitor() != null) {
+				tinyRemapperBuilder.extraPostApplyVisitor(tinyRemapperExtension.getPostApplyVisitor());
+			}
 		}
 	}
 
-	private RemapperExtension<?> newInstance() {
+	private RemapperExtension<?> newInstance(ObjectFactory objectFactory) {
 		try {
-			return Class.forName(getRemapperExtensionClassName().get())
-					.asSubclass(RemapperExtension.class)
-					.getConstructor()
-					.newInstance();
+			Class<? extends RemapperExtension> remapperExtensionClass = Class.forName(getRemapperExtensionClassName().get())
+					.asSubclass(RemapperExtension.class);
+			return objectFactory.newInstance(remapperExtensionClass, getRemapperParameters());
 		} catch (Exception e) {
 			throw new RuntimeException("Failed to create remapper extension", e);
 		}
 	}
 
-	private record RemapperExtensionImpl(RemapperExtension<?> remapperExtension) implements TinyRemapper.ApplyVisitorProvider {
+	private static final class RemapperExtensionImpl implements TinyRemapper.ApplyVisitorProvider {
+		private final RemapperExtension<?> remapperExtension;
+		private final String sourceNamespace;
+		private final String targetNamespace;
+
+		@Nullable
+		private RemapperContext context;
+
+		private RemapperExtensionImpl(RemapperExtension<?> remapperExtension, String sourceNamespace, String targetNamespace) {
+			this.remapperExtension = remapperExtension;
+			this.sourceNamespace = sourceNamespace;
+			this.targetNamespace = targetNamespace;
+		}
+
 		@Override
 		public ClassVisitor insertApplyVisitor(TrClass cls, ClassVisitor next) {
-			return remapperExtension.insertVisitor(cls.getName(), new RemapperContext() {
-				// TODO dont create a new instance of this for every class
-				@Override
-				public Remapper getRemapper() {
-					return cls.getEnvironment().getRemapper();
-				}
-
-				@Override
-				public String sourceNamespace() {
-					return null;
-				}
-
-				@Override
-				public String targetNamespace() {
-					return null;
-				}
-			}, next);
+			if (context == null) {
+				context = new RemapperContextImpl(cls.getEnvironment().getRemapper(), sourceNamespace, targetNamespace);
+			}
+
+			return remapperExtension.insertVisitor(cls.getName(), context, next);
 		}
 	}
+
+	private record RemapperContextImpl(Remapper remapper, String sourceNamespace, String targetNamespace) implements RemapperContext {
+	}
 }

src/main/java/net/fabricmc/loom/task/service/TinyRemapperService.java

@@ -39,6 +39,7 @@
 
 import org.gradle.api.Project;
 import org.gradle.api.invocation.Gradle;
+import org.gradle.api.model.ObjectFactory;
 import org.gradle.api.tasks.SourceSet;
 import org.jetbrains.annotations.Nullable;
 
@@ -91,7 +92,7 @@ public static synchronized TinyRemapperService getOrCreate(SharedServiceManager
 				mappings.add(gradleMixinMappingProvider(serviceManager, project.getGradle(), extension.getMappingConfiguration().mappingsIdentifier, from, to));
 			}
 
-			return new TinyRemapperService(mappings, !legacyMixin, kotlinClasspathService, extension.getKnownIndyBsms().get(), extension.getRemapperExtensions().get());
+			return new TinyRemapperService(mappings, !legacyMixin, kotlinClasspathService, extension.getKnownIndyBsms().get(), extension.getRemapperExtensions().get(), from, to, project.getObjects());
 		});
 
 		service.readClasspath(remapJarTask.getClasspath().getFiles().stream().map(File::toPath).filter(Files::exists).toList());
@@ -130,7 +131,7 @@ private static IMappingProvider gradleMixinMappingProvider(SharedServiceManager
 	// Set to true once remapping has started, once set no inputs can be read.
 	private boolean isRemapping = false;
 
-	public TinyRemapperService(List<IMappingProvider> mappings, boolean useMixinExtension, @Nullable KotlinClasspath kotlinClasspath, Set<String> knownIndyBsms, List<RemapperExtensionHolder> remapperExtensions) {
+	private TinyRemapperService(List<IMappingProvider> mappings, boolean useMixinExtension, @Nullable KotlinClasspath kotlinClasspath, Set<String> knownIndyBsms, List<RemapperExtensionHolder> remapperExtensions, String sourceNamespace, String targetNamespace, ObjectFactory objectFactory) {
 		TinyRemapper.Builder builder = TinyRemapper.newRemapper().withKnownIndyBsm(knownIndyBsms);
 
 		for (IMappingProvider provider : mappings) {
@@ -147,7 +148,7 @@ public TinyRemapperService(List<IMappingProvider> mappings, boolean useMixinExte
 		}
 
 		for (RemapperExtensionHolder holder : remapperExtensions) {
-			holder.apply(builder);
+			holder.apply(builder, sourceNamespace, targetNamespace, objectFactory);
 		}
 
 		tinyRemapper = builder.build();

src/test/groovy/net/fabricmc/loom/test/integration/SimpleProjectTest.groovy

@@ -42,6 +42,7 @@ class SimpleProjectTest extends Specification implements GradleProjectTestTrait
 	def "build and run (gradle #version)"() {
 		setup:
 		def gradle = gradleProject(project: "simple", version: version)
+		gradle.buildSrc("remapext") // apply the remap extension plugin
 
 		def server = ServerRunner.create(gradle.projectDir, "1.16.5")
 				.withMod(gradle.getOutputFile("fabric-example-mod-1.0.0.jar"))
@@ -60,6 +61,7 @@ class SimpleProjectTest extends Specification implements GradleProjectTestTrait
 
 		serverResult.successful()
 		serverResult.output.contains("Hello simple Fabric mod") // A check to ensure our mod init was actually called
+		serverResult.output.contains("Hello Loom!") // Check that the remapper extension worked
 		where:
 		version << STANDARD_TEST_VERSIONS
 	}

src/test/groovy/net/fabricmc/loom/test/integration/buildSrc/remapext/StringReplacementClassVisitor.groovy

@@ -0,0 +1,61 @@
+/*
+ * This file is part of fabric-loom, licensed under the MIT License (MIT).
+ *
+ * Copyright (c) 2023 FabricMC
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+package net.fabricmc.loom.test.integration.buildSrc.remapext
+
+import org.objectweb.asm.ClassVisitor
+import org.objectweb.asm.MethodVisitor
+
+class StringReplacementClassVisitor extends ClassVisitor {
+	final Map<String, String> replacements
+
+	StringReplacementClassVisitor(int api, ClassVisitor classVisitor, Map<String, String> replacements) {
+		super(api, classVisitor)
+		this.replacements = replacements
+	}
+
+	@Override
+	MethodVisitor visitMethod(int access, String name, String descriptor, String signature, String[] exceptions) {
+		def methodVisitor = super.visitMethod(access, name, descriptor, signature, exceptions)
+		return new StringReplacementMethodVisitor(api, methodVisitor)
+	}
+
+	class StringReplacementMethodVisitor extends MethodVisitor {
+		StringReplacementMethodVisitor(int api, MethodVisitor methodVisitor) {
+			super(api, methodVisitor)
+		}
+
+		@Override
+		void visitLdcInsn(Object value) {
+			if (value instanceof String) {
+				String replacement = replacements.get(value)
+				if (replacement != null) {
+					value = replacement
+				}
+			}
+
+			super.visitLdcInsn(value)
+		}
+	}
+}