Remapper extension API

src/main/java/net/fabricmc/loom/LoomGradleExtension.java

@@ -45,6 +45,7 @@
 import net.fabricmc.loom.configuration.providers.minecraft.mapped.NamedMinecraftProvider;
 import net.fabricmc.loom.extension.LoomFiles;
 import net.fabricmc.loom.extension.MixinExtension;
+import net.fabricmc.loom.extension.RemapperExtensionHolder;
 import net.fabricmc.loom.util.download.DownloadBuilder;
 
 @ApiStatus.Internal
@@ -115,4 +116,6 @@ default List<Path> getMinecraftJars(MappingsNamespace mappingsNamespace) {
 	boolean multiProjectOptimisation();
 
 	ListProperty<LibraryProcessorManager.LibraryProcessorFactory> getLibraryProcessors();
+
+	ListProperty<RemapperExtensionHolder> getRemapperExtensions();
 }

src/main/java/net/fabricmc/loom/api/LoomGradleExtensionAPI.java

@@ -44,6 +44,8 @@
 import net.fabricmc.loom.api.mappings.intermediate.IntermediateMappingsProvider;
 import net.fabricmc.loom.api.mappings.layered.spec.LayeredMappingSpecBuilder;
 import net.fabricmc.loom.api.processor.MinecraftJarProcessor;
+import net.fabricmc.loom.api.remapping.RemapperExtension;
+import net.fabricmc.loom.api.remapping.RemapperParameters;
 import net.fabricmc.loom.configuration.ide.RunConfigSettings;
 import net.fabricmc.loom.configuration.processors.JarProcessor;
 import net.fabricmc.loom.configuration.providers.mappings.NoOpIntermediateMappingsProvider;
@@ -221,4 +223,6 @@ default void splitMinecraftJar() {
 	Property<Boolean> getRuntimeOnlyLog4j();
 
 	Property<Boolean> getSplitModDependencies();
+
+	<T extends RemapperParameters> void addRemapperExtension(Class<RemapperExtension<T>> remapperExtensionClass, Class<T> parametersClass, Action<T> parameterAction);
 }

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

@@ -0,0 +1,11 @@
+package net.fabricmc.loom.api.remapping;
+
+import org.objectweb.asm.commons.Remapper;
+
+public interface RemapperContext {
+	Remapper getRemapper();
+
+	String sourceNamespace();
+
+	String targetNamespace();
+}

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

@@ -0,0 +1,12 @@
+package net.fabricmc.loom.api.remapping;
+
+import javax.inject.Inject;
+
+import org.objectweb.asm.ClassVisitor;
+
+public interface RemapperExtension<T extends RemapperParameters> {
+	@Inject
+	T getParameters();
+
+	ClassVisitor insertVisitor(String className, RemapperContext remapperContext, ClassVisitor classVisitor);
+}

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

@@ -0,0 +1,4 @@
+package net.fabricmc.loom.api.remapping;
+
+public interface RemapperParameters {
+}

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

@@ -51,6 +51,7 @@
 import net.fabricmc.loom.api.mappings.layered.MappingsNamespace;
 import net.fabricmc.loom.configuration.mods.dependency.ModDependency;
 import net.fabricmc.loom.configuration.providers.mappings.MappingConfiguration;
+import net.fabricmc.loom.extension.RemapperExtensionHolder;
 import net.fabricmc.loom.task.RemapJarTask;
 import net.fabricmc.loom.util.Constants;
 import net.fabricmc.loom.util.Pair;
@@ -153,6 +154,10 @@ private void remapJars(List<ModDependency> remapList) throws IOException {
 		// Configure the mixin extension to remap mixins from mod jars detected not to contain refmaps.
 		builder.extension(new MixinExtension(hasMixinsWithoutRefmaps::contains));
 
+		for (RemapperExtensionHolder holder : extension.getRemapperExtensions().get()) {
+			holder.apply(builder);
+		}
+
 		final TinyRemapper remapper = builder.build();
 
 		for (Path minecraftJar : extension.getMinecraftJars(MappingsNamespace.INTERMEDIARY)) {

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

@@ -50,6 +50,8 @@
 import net.fabricmc.loom.api.mappings.intermediate.IntermediateMappingsProvider;
 import net.fabricmc.loom.api.mappings.layered.spec.LayeredMappingSpecBuilder;
 import net.fabricmc.loom.api.processor.MinecraftJarProcessor;
+import net.fabricmc.loom.api.remapping.RemapperExtension;
+import net.fabricmc.loom.api.remapping.RemapperParameters;
 import net.fabricmc.loom.configuration.RemapConfigurations;
 import net.fabricmc.loom.configuration.ide.RunConfigSettings;
 import net.fabricmc.loom.configuration.processors.JarProcessor;
@@ -91,6 +93,7 @@ public abstract class LoomGradleExtensionApiImpl implements LoomGradleExtensionA
 	private final NamedDomainObjectContainer<ModSettings> mods;
 	private final NamedDomainObjectList<RemapConfigurationSettings> remapConfigurations;
 	private final ListProperty<MinecraftJarProcessor<?>> minecraftJarProcessors;
+	protected final ListProperty<RemapperExtensionHolder> remapperExtensions;
 
 	// A common mistake with layered mappings is to call the wrong `officialMojangMappings` method, use this to keep track of when we are building a layered mapping spec.
 	protected final ThreadLocal<Boolean> layeredSpecBuilderScope = ThreadLocal.withInitial(() -> false);
@@ -150,6 +153,9 @@ protected LoomGradleExtensionApiImpl(Project project, LoomFiles directories) {
 		this.splitEnvironmentalSourceSet = project.getObjects().property(Boolean.class).convention(false);
 		this.splitEnvironmentalSourceSet.finalizeValueOnRead();
 
+		remapperExtensions = project.getObjects().listProperty(RemapperExtensionHolder.class);
+		remapperExtensions.finalizeValueOnRead();
+
 		// Enable dep iface injection by default
 		interfaceInjection(interfaceInjection -> {
 			interfaceInjection.getEnableDependencyInterfaceInjection().convention(true).finalizeValueOnRead();
@@ -382,6 +388,16 @@ public void createRemapConfigurations(SourceSet sourceSet) {
 		RemapConfigurations.setupForSourceSet(getProject(), sourceSet);
 	}
 
+	@Override
+	public <T extends RemapperParameters> void addRemapperExtension(Class<RemapperExtension<T>> remapperExtensionClass, Class<T> parametersClass, Action<T> parameterAction) {
+		T parameters = getProject().getObjects().newInstance(parametersClass);
+		parameterAction.execute(parameters);
+
+		RemapperExtensionHolder holder = getProject().getObjects().newInstance(RemapperExtensionHolder.class, parameters);
+		holder.getRemapperExtensionClassName().set(remapperExtensionClass.getName());
+		remapperExtensions.add(holder);
+	}
+
 	// This is here to ensure that LoomGradleExtensionApiImpl compiles without any unimplemented methods
 	private final class EnsureCompile extends LoomGradleExtensionApiImpl {
 		private EnsureCompile() {

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

@@ -247,6 +247,11 @@ public ListProperty<LibraryProcessorManager.LibraryProcessorFactory> getLibraryP
 		return libraryProcessorFactories;
 	}
 
+	@Override
+	public ListProperty<RemapperExtensionHolder> getRemapperExtensions() {
+		return remapperExtensions;
+	}
+
 	@Override
 	protected <T extends IntermediateMappingsProvider> void configureIntermediateMappingsProviderInternal(T provider) {
 		provider.getMinecraftVersion().set(getProject().provider(() -> getMinecraftProvider().minecraftVersion()));

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

@@ -0,0 +1,81 @@
+package net.fabricmc.loom.extension;
+
+import javax.inject.Inject;
+
+import org.gradle.api.provider.Property;
+import org.gradle.api.tasks.Input;
+import org.gradle.api.tasks.Nested;
+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.tinyremapper.TinyRemapper;
+import net.fabricmc.tinyremapper.api.TrClass;
+
+public abstract class RemapperExtensionHolder {
+	private final RemapperParameters remapperParameters;
+
+	@Inject
+	public RemapperExtensionHolder(RemapperParameters remapperParameters) {
+		this.remapperParameters = remapperParameters;
+	}
+
+	@Input
+	public abstract Property<String> getRemapperExtensionClassName();
+
+	@Nested
+	public RemapperParameters getRemapperParameters() {
+		return remapperParameters;
+	}
+
+	public void apply(TinyRemapper.Builder tinyRemapperBuilder) {
+		RemapperExtension<?> remapperExtension = newInstance();
+
+		tinyRemapperBuilder.extraPostApplyVisitor(new RemapperExtensionImpl(remapperExtension));
+
+		if (remapperExtension instanceof TinyRemapper.AnalyzeVisitorProvider analyzeVisitorProvider) {
+			tinyRemapperBuilder.extraAnalyzeVisitor(analyzeVisitorProvider);
+		}
+
+		if (remapperExtension instanceof TinyRemapper.ApplyVisitorProvider applyVisitorProvider) {
+			// TODO allow having a pre apply visitor?
+			tinyRemapperBuilder.extraPostApplyVisitor(applyVisitorProvider);
+		}
+	}
+
+	private RemapperExtension<?> newInstance() {
+		try {
+			return Class.forName(getRemapperExtensionClassName().get())
+					.asSubclass(RemapperExtension.class)
+					.getConstructor()
+					.newInstance();
+		} catch (Exception e) {
+			throw new RuntimeException("Failed to create remapper extension", e);
+		}
+	}
+
+	private record RemapperExtensionImpl(RemapperExtension<?> remapperExtension) implements TinyRemapper.ApplyVisitorProvider {
+		@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);
+		}
+	}
+}

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

@@ -44,6 +44,7 @@
 
 import net.fabricmc.loom.LoomGradleExtension;
 import net.fabricmc.loom.build.mixin.AnnotationProcessorInvoker;
+import net.fabricmc.loom.extension.RemapperExtensionHolder;
 import net.fabricmc.loom.task.AbstractRemapJarTask;
 import net.fabricmc.loom.util.gradle.GradleUtils;
 import net.fabricmc.loom.util.gradle.SourceSetHelper;
@@ -90,7 +91,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());
+			return new TinyRemapperService(mappings, !legacyMixin, kotlinClasspathService, extension.getKnownIndyBsms().get(), extension.getRemapperExtensions().get());
 		});
 
 		service.readClasspath(remapJarTask.getClasspath().getFiles().stream().map(File::toPath).filter(Files::exists).toList());
@@ -129,7 +130,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) {
+	public TinyRemapperService(List<IMappingProvider> mappings, boolean useMixinExtension, @Nullable KotlinClasspath kotlinClasspath, Set<String> knownIndyBsms, List<RemapperExtensionHolder> remapperExtensions) {
 		TinyRemapper.Builder builder = TinyRemapper.newRemapper().withKnownIndyBsm(knownIndyBsms);
 
 		for (IMappingProvider provider : mappings) {
@@ -145,6 +146,10 @@ public TinyRemapperService(List<IMappingProvider> mappings, boolean useMixinExte
 			builder.extension(kotlinRemapperClassloader.getTinyRemapperExtension());
 		}
 
+		for (RemapperExtensionHolder holder : remapperExtensions) {
+			holder.apply(builder);
+		}
+
 		tinyRemapper = builder.build();
 	}