Add ability to remap annotations data (#1366)

* Add ability to remap annotations data

* Fix unpick remap test

src/main/java/net/fabricmc/loom/configuration/providers/mappings/LayeredMappingsFactory.java

@@ -134,13 +134,13 @@ public record LayeredMappingsFactory(LayeredMappingSpec spec) {
 	}
 
 	private void writeAnnotationData(LayeredMappingsProcessor processor, List<MappingLayer> layers, Path mappingsFile) throws IOException {
-		AnnotationsData annotationsData = processor.getAnnotationsData(layers);
+		List<AnnotationsData> annotationsData = processor.getAnnotationsData(layers);
 
-		if (annotationsData == null) {
+		if (annotationsData.isEmpty()) {
 			return;
 		}
 
-		byte[] data = AnnotationsData.GSON.toJson(annotationsData.toJson()).getBytes(StandardCharsets.UTF_8);
+		byte[] data = AnnotationsData.GSON.toJson(AnnotationsData.listToJson(annotationsData)).getBytes(StandardCharsets.UTF_8);
 
 		ZipUtils.add(mappingsFile, AnnotationsLayer.ANNOTATIONS_PATH, data);
 	}

src/main/java/net/fabricmc/loom/configuration/providers/mappings/LayeredMappingsProcessor.java

@@ -119,20 +119,15 @@ public class LayeredMappingsProcessor {
 		return mappingTree;
 	}
 
-	@Nullable
-	public AnnotationsData getAnnotationsData(List<MappingLayer> layers) throws IOException {
-		AnnotationsData result = null;
+	public List<AnnotationsData> getAnnotationsData(List<MappingLayer> layers) throws IOException {
+		List<AnnotationsData> result = new ArrayList<>();
 
 		for (MappingLayer layer : layers) {
 			if (layer instanceof AnnotationsLayer annotationsLayer) {
 				AnnotationsData annotationsData = annotationsLayer.getAnnotationsData();
 
 				if (annotationsData != null) {
-					if (result == null) {
-						result = annotationsData;
-					} else {
-						result = result.merge(annotationsData);
-					}
+					result.add(annotationsData);
 				}
 			}
 		}

src/main/java/net/fabricmc/loom/configuration/providers/mappings/MappingConfiguration.java

@@ -77,8 +77,7 @@ public class MappingConfiguration {
 	public final Path tinyMappingsJar;
 	private final Path unpickDefinitions;
 
-	@Nullable
-	private AnnotationsData annotationsData;
+	private List<AnnotationsData> annotationsData = List.of();
 	@Nullable
 	private UnpickMetadata unpickMetadata;
 	private Map<String, String> signatureFixes;
@@ -233,7 +232,7 @@ public class MappingConfiguration {
 		}
 
 		try (BufferedReader reader = Files.newBufferedReader(annotationsPath, StandardCharsets.UTF_8)) {
-			annotationsData = AnnotationsData.read(reader);
+			annotationsData = AnnotationsData.readList(reader);
 		}
 	}
 
@@ -312,8 +311,7 @@ public class MappingConfiguration {
 		return unpickMetadata != null;
 	}
 
-	@Nullable
-	public AnnotationsData getAnnotationsData() {
+	public List<AnnotationsData> getAnnotationsData() {
 		return annotationsData;
 	}
 

src/main/java/net/fabricmc/loom/configuration/providers/mappings/extras/annotations/AnnotationsData.java

@@ -24,19 +24,35 @@
 
 package net.fabricmc.loom.configuration.providers.mappings.extras.annotations;
 
+import java.io.IOException;
 import java.io.Reader;
+import java.lang.reflect.Type;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 
 import com.google.gson.FieldNamingPolicy;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
+import com.google.gson.JsonElement;
 import com.google.gson.JsonObject;
 import com.google.gson.JsonSyntaxException;
+import com.google.gson.reflect.TypeToken;
+import org.gradle.api.Project;
+import org.jetbrains.annotations.Nullable;
+import org.objectweb.asm.commons.AnnotationRemapper;
 import org.objectweb.asm.tree.AnnotationNode;
 import org.objectweb.asm.tree.TypeAnnotationNode;
 
-public record AnnotationsData(Map<String, ClassAnnotationData> classes) {
+import net.fabricmc.loom.api.mappings.layered.MappingsNamespace;
+import net.fabricmc.loom.configuration.providers.mappings.MappingConfiguration;
+import net.fabricmc.loom.util.TinyRemapperHelper;
+import net.fabricmc.loom.util.service.ServiceFactory;
+import net.fabricmc.tinyremapper.TinyRemapper;
+
+public record AnnotationsData(Map<String, ClassAnnotationData> classes, String namespace) {
 	public static final Gson GSON = new GsonBuilder()
 			.disableHtmlEscaping()
 			.setFieldNamingStrategy(FieldNamingPolicy.LOWER_CASE_WITH_UNDERSCORES)
@@ -45,34 +61,139 @@ public record AnnotationsData(Map<String, ClassAnnotationData> classes) {
 			.registerTypeAdapter(AnnotationNode.class, new AnnotationNodeSerializer())
 			.registerTypeAdapterFactory(new SkipEmptyTypeAdapterFactory())
 			.create();
+	private static final Type LIST_TYPE = new TypeToken<List<AnnotationNode>>() { }.getType();
+	private static final int CURRENT_VERSION = 1;
+
+	public AnnotationsData {
+		if (namespace == null) {
+			namespace = MappingsNamespace.NAMED.toString();
+		}
+	}
 
 	public static AnnotationsData read(Reader reader) {
 		JsonObject json = GSON.fromJson(reader, JsonObject.class);
+		checkVersion(json);
+		return GSON.fromJson(json, AnnotationsData.class);
+	}
 
+	public static List<AnnotationsData> readList(Reader reader) {
+		JsonObject json = GSON.fromJson(reader, JsonObject.class);
+		checkVersion(json);
+		JsonElement values = json.get("values");
+
+		if (values == null || values.isJsonNull()) {
+			return List.of(GSON.fromJson(json, AnnotationsData.class));
+		}
+
+		return GSON.fromJson(values, LIST_TYPE);
+	}
+
+	private static void checkVersion(JsonObject json) {
 		if (!json.has("version")) {
 			throw new JsonSyntaxException("Missing annotations version");
 		}
 
 		int version = json.getAsJsonPrimitive("version").getAsInt();
 
-		if (version != 1) {
+		if (version != CURRENT_VERSION) {
 			throw new JsonSyntaxException("Invalid annotations version " + version + ". Try updating loom");
 		}
-
-		return GSON.fromJson(json, AnnotationsData.class);
 	}
 
 	public JsonObject toJson() {
 		JsonObject json = GSON.toJsonTree(this).getAsJsonObject();
 		JsonObject result = new JsonObject();
-		result.addProperty("version", 1);
+		result.addProperty("version", CURRENT_VERSION);
 		result.asMap().putAll(json.asMap());
 		return result;
 	}
 
+	public static JsonObject listToJson(List<AnnotationsData> annotationsData) {
+		if (annotationsData.size() == 1) {
+			return annotationsData.getFirst().toJson();
+		}
+
+		JsonObject result = new JsonObject();
+		result.addProperty("version", CURRENT_VERSION);
+		result.add("values", GSON.toJsonTree(annotationsData));
+		return result;
+	}
+
 	public AnnotationsData merge(AnnotationsData other) {
+		if (!namespace.equals(other.namespace)) {
+			throw new IllegalArgumentException("Cannot merge annotations from namespace " + other.namespace + " into annotations from namespace " + this.namespace);
+		}
+
 		Map<String, ClassAnnotationData> newClassData = new LinkedHashMap<>(classes);
 		other.classes.forEach((key, value) -> newClassData.merge(key, value, ClassAnnotationData::merge));
-		return new AnnotationsData(newClassData);
+		return new AnnotationsData(newClassData, namespace);
+	}
+
+	public AnnotationsData remap(TinyRemapper remapper, String newNamespace) {
+		return new AnnotationsData(
+				remapMap(
+						classes,
+						entry -> remapper.getEnvironment().getRemapper().map(entry.getKey()),
+						entry -> entry.getValue().remap(entry.getKey(), remapper)
+				),
+				newNamespace
+		);
+	}
+
+	static AnnotationNode remap(AnnotationNode node, TinyRemapper remapper) {
+		AnnotationNode remapped = new AnnotationNode(remapper.getEnvironment().getRemapper().mapDesc(node.desc));
+		node.accept(new AnnotationRemapper(node.desc, remapped, remapper.getEnvironment().getRemapper()));
+		return remapped;
+	}
+
+	static TypeAnnotationNode remap(TypeAnnotationNode node, TinyRemapper remapper) {
+		TypeAnnotationNode remapped = new TypeAnnotationNode(node.typeRef, node.typePath, remapper.getEnvironment().getRemapper().mapDesc(node.desc));
+		node.accept(new AnnotationRemapper(node.desc, remapped, remapper.getEnvironment().getRemapper()));
+		return remapped;
+	}
+
+	static <K, V> Map<K, V> remapMap(Map<K, V> map, Function<Map.Entry<K, V>, K> keyRemapper, Function<Map.Entry<K, V>, V> valueRemapper) {
+		Map<K, V> result = LinkedHashMap.newLinkedHashMap(map.size());
+
+		for (Map.Entry<K, V> entry : map.entrySet()) {
+			if (result.put(keyRemapper.apply(entry), valueRemapper.apply(entry)) != null) {
+				throw new IllegalStateException("Remapping annotations resulted in duplicate key: " + keyRemapper.apply(entry));
+			}
+		}
+
+		return result;
+	}
+
+	@Nullable
+	public static AnnotationsData getRemappedAnnotations(MappingsNamespace targetNamespace, MappingConfiguration mappingConfiguration, Project project, ServiceFactory serviceFactory, String newNamespace) throws IOException {
+		List<AnnotationsData> datas = mappingConfiguration.getAnnotationsData();
+
+		if (datas.isEmpty()) {
+			return null;
+		}
+
+		Map<String, TinyRemapper> existingRemappers = new HashMap<>();
+		AnnotationsData result = datas.getFirst().remap(targetNamespace, project, serviceFactory, newNamespace, existingRemappers);
+
+		for (int i = 1; i < datas.size(); i++) {
+			result = result.merge(datas.get(i).remap(targetNamespace, project, serviceFactory, newNamespace, existingRemappers));
+		}
+
+		return result;
+	}
+
+	private AnnotationsData remap(MappingsNamespace targetNamespace, Project project, ServiceFactory serviceFactory, String newNamespace, Map<String, TinyRemapper> existingRemappers) throws IOException {
+		if (namespace.equals(targetNamespace.toString())) {
+			return this;
+		}
+
+		TinyRemapper remapper = existingRemappers.get(namespace);
+
+		if (remapper == null) {
+			remapper = TinyRemapperHelper.getTinyRemapper(project, serviceFactory, namespace, newNamespace);
+			existingRemappers.put(namespace, remapper);
+		}
+
+		return remap(remapper, newNamespace);
 	}
 }

src/main/java/net/fabricmc/loom/configuration/providers/mappings/extras/annotations/ClassAnnotationData.java

@@ -30,6 +30,7 @@ import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import com.google.gson.annotations.SerializedName;
 import org.jetbrains.annotations.Nullable;
@@ -37,6 +38,9 @@ import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.tree.AnnotationNode;
 import org.objectweb.asm.tree.TypeAnnotationNode;
 
+import net.fabricmc.tinyremapper.TinyRemapper;
+import net.fabricmc.tinyremapper.api.TrRemapper;
+
 public record ClassAnnotationData(
 		@SerializedName("remove")
 		Set<String> annotationsToRemove,
@@ -90,6 +94,49 @@ public record ClassAnnotationData(
 		return new ClassAnnotationData(newAnnotationsToRemove, newAnnotationsToAdd, newTypeAnnotationsToRemove, newTypeAnnotationsToAdd, newFields, newMethods);
 	}
 
+	ClassAnnotationData remap(String className, TinyRemapper remapper) {
+		return new ClassAnnotationData(
+				annotationsToRemove.stream().map(remapper.getEnvironment().getRemapper()::map).collect(Collectors.toCollection(LinkedHashSet::new)),
+				annotationsToAdd.stream().map(ann -> AnnotationsData.remap(ann, remapper)).collect(Collectors.toCollection(ArrayList::new)),
+				typeAnnotationsToRemove.stream().map(key -> key.remap(remapper)).collect(Collectors.toCollection(LinkedHashSet::new)),
+				typeAnnotationsToAdd.stream().map(ann -> AnnotationsData.remap(ann, remapper)).collect(Collectors.toCollection(ArrayList::new)),
+				AnnotationsData.remapMap(
+						fields,
+						entry -> remapField(className, entry.getKey(), remapper),
+						entry -> entry.getValue().remap(remapper)
+				),
+				AnnotationsData.remapMap(
+						methods,
+						entry -> remapMethod(className, entry.getKey(), remapper),
+						entry -> entry.getValue().remap(remapper)
+				)
+		);
+	}
+
+	private static String remapField(String className, String field, TinyRemapper remapper) {
+		String[] nameDesc = field.split(":", 2);
+
+		if (nameDesc.length != 2) {
+			return field;
+		}
+
+		TrRemapper trRemapper = remapper.getEnvironment().getRemapper();
+		return trRemapper.mapFieldName(className, nameDesc[0], nameDesc[1]) + ":" + trRemapper.mapDesc(nameDesc[1]);
+	}
+
+	private static String remapMethod(String className, String method, TinyRemapper remapper) {
+		int parenIndex = method.indexOf('(');
+
+		if (parenIndex == -1) {
+			return method;
+		}
+
+		String name = method.substring(0, parenIndex);
+		String desc = method.substring(parenIndex);
+		TrRemapper trRemapper = remapper.getEnvironment().getRemapper();
+		return trRemapper.mapMethodName(className, name, desc) + trRemapper.mapMethodDesc(desc);
+	}
+
 	public int modifyAccessFlags(int access) {
 		if (annotationsToRemove.contains("java/lang/Deprecated")) {
 			access &= ~Opcodes.ACC_DEPRECATED;

src/main/java/net/fabricmc/loom/configuration/providers/mappings/extras/annotations/GenericAnnotationData.java

@@ -28,12 +28,15 @@ import java.util.ArrayList;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import com.google.gson.annotations.SerializedName;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.tree.AnnotationNode;
 import org.objectweb.asm.tree.TypeAnnotationNode;
 
+import net.fabricmc.tinyremapper.TinyRemapper;
+
 public record GenericAnnotationData(
 		@SerializedName("remove")
 		Set<String> annotationsToRemove,
@@ -74,6 +77,15 @@ public record GenericAnnotationData(
 		return new GenericAnnotationData(newAnnotationToRemove, newAnnotationsToAdd, newTypeAnnotationsToRemove, newTypeAnnotationsToAdd);
 	}
 
+	GenericAnnotationData remap(TinyRemapper remapper) {
+		return new GenericAnnotationData(
+				annotationsToRemove.stream().map(remapper.getEnvironment().getRemapper()::map).collect(Collectors.toCollection(LinkedHashSet::new)),
+				annotationsToAdd.stream().map(ann -> AnnotationsData.remap(ann, remapper)).collect(Collectors.toCollection(ArrayList::new)),
+				typeAnnotationsToRemove.stream().map(key -> key.remap(remapper)).collect(Collectors.toCollection(LinkedHashSet::new)),
+				typeAnnotationsToAdd.stream().map(ann -> AnnotationsData.remap(ann, remapper)).collect(Collectors.toCollection(ArrayList::new))
+		);
+	}
+
 	public int modifyAccessFlags(int access) {
 		if (annotationsToRemove.contains("java/lang/Deprecated")) {
 			access &= ~Opcodes.ACC_DEPRECATED;

src/main/java/net/fabricmc/loom/configuration/providers/mappings/extras/annotations/MethodAnnotationData.java

@@ -30,12 +30,15 @@ import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import com.google.gson.annotations.SerializedName;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.tree.AnnotationNode;
 import org.objectweb.asm.tree.TypeAnnotationNode;
 
+import net.fabricmc.tinyremapper.TinyRemapper;
+
 public record MethodAnnotationData(
 		@SerializedName("remove")
 		Set<String> annotationsToRemove,
@@ -83,6 +86,20 @@ public record MethodAnnotationData(
 		return new MethodAnnotationData(newAnnotationsToRemove, newAnnotationsToAdd, newTypeAnnotationsToRemove, newTypeAnnotationsToAdd, newParameters);
 	}
 
+	MethodAnnotationData remap(TinyRemapper remapper) {
+		return new MethodAnnotationData(
+				annotationsToRemove.stream().map(remapper.getEnvironment().getRemapper()::map).collect(Collectors.toCollection(LinkedHashSet::new)),
+				annotationsToAdd.stream().map(ann -> AnnotationsData.remap(ann, remapper)).collect(Collectors.toCollection(ArrayList::new)),
+				typeAnnotationsToRemove.stream().map(key -> key.remap(remapper)).collect(Collectors.toCollection(LinkedHashSet::new)),
+				typeAnnotationsToAdd.stream().map(ann -> AnnotationsData.remap(ann, remapper)).collect(Collectors.toCollection(ArrayList::new)),
+				AnnotationsData.remapMap(
+						parameters,
+						Map.Entry::getKey,
+						entry -> entry.getValue().remap(remapper)
+				)
+		);
+	}
+
 	public int modifyAccessFlags(int access) {
 		if (annotationsToRemove.contains("java/lang/Deprecated")) {
 			access &= ~Opcodes.ACC_DEPRECATED;

src/main/java/net/fabricmc/loom/configuration/providers/mappings/extras/annotations/TypeAnnotationKey.java

@@ -24,5 +24,10 @@
 
 package net.fabricmc.loom.configuration.providers.mappings.extras.annotations;
 
+import net.fabricmc.tinyremapper.TinyRemapper;
+
 public record TypeAnnotationKey(int typeRef, String typePath, String name) {
+	TypeAnnotationKey remap(TinyRemapper remapper) {
+		return new TypeAnnotationKey(typeRef, typePath, remapper.getEnvironment().getRemapper().map(name));
+	}
 }

src/main/java/net/fabricmc/loom/configuration/providers/minecraft/mapped/AbstractMappedMinecraftProvider.java

@@ -244,15 +244,14 @@ public abstract class AbstractMappedMinecraftProvider<M extends MinecraftProvide
 
 		Files.deleteIfExists(remappedJars.outputJarPath());
 
+		final AnnotationsData remappedAnnotations = AnnotationsData.getRemappedAnnotations(getTargetNamespace(), mappingConfiguration, getProject(), configContext.serviceFactory(), toM);
 		final Map<String, String> remappedSignatures = SignatureFixerApplyVisitor.getRemappedSignatures(getTargetNamespace() == MappingsNamespace.INTERMEDIARY, mappingConfiguration, getProject(), configContext.serviceFactory(), toM);
 		final MinecraftVersionMeta.JavaVersion javaVersion = minecraftProvider.getVersionInfo().javaVersion();
 		final boolean fixRecords = javaVersion != null && javaVersion.majorVersion() >= 16;
 
 		TinyRemapper remapper = TinyRemapperHelper.getTinyRemapper(getProject(), configContext.serviceFactory(), fromM, toM, fixRecords, (builder) -> {
-			AnnotationsData annotationsData = mappingConfiguration.getAnnotationsData();
-
-			if (annotationsData != null) {
-				builder.extraPostApplyVisitor(new AnnotationsApplyVisitor(annotationsData));
+			if (remappedAnnotations != null) {
+				builder.extraPostApplyVisitor(new AnnotationsApplyVisitor(remappedAnnotations));
 			}
 
 			builder.extraPostApplyVisitor(new SignatureFixerApplyVisitor(remappedSignatures));