add choice type fixes + use custom schema

fabric-api-lookup-api-v1/src/testmod/java/net/fabricmc/fabric/test/lookup/FabricApiLookupTestDataFixer.java

@@ -16,8 +16,6 @@
 
 package net.fabricmc.fabric.test.lookup;
 
-
-
 import com.mojang.datafixers.schemas.Schema;
 
 import net.minecraft.datafixer.schema.Schema100;

fabric-data-fixer-api-v1/src/main/java/net/fabricmc/fabric/impl/datafixer/v1/FabricDataFixesInternals.java

@@ -37,8 +37,11 @@
 import net.minecraft.nbt.NbtCompound;
 import net.minecraft.nbt.NbtElement;
 
+import net.fabricmc.fabric.api.datafixer.v1.DataFixerEntrypoint;
 import net.fabricmc.fabric.api.datafixer.v1.SchemaRegistry;
+import net.fabricmc.loader.api.FabricLoader;
 import net.fabricmc.loader.api.ModContainer;
+import net.fabricmc.loader.api.entrypoint.EntrypointContainer;
 import net.fabricmc.loader.api.metadata.CustomValue;
 
 public abstract class FabricDataFixesInternals {
@@ -47,6 +50,7 @@ public abstract class FabricDataFixesInternals {
 	protected static final String DATA_VERSIONS_KEY = "_FabricDataVersions";
 	public static final String METADATA_VERSION_KEY = "fabric-data-fixer-api-v1:version";
 	public static final String METADATA_KEY_KEY = "fabric-data-fixer-api-v1:key";
+	private static final String ENTRYPOINT_KEY = "fabric-data-fixer";
 
 	public record DataFixerEntry(DataFixer dataFixer, int currentVersion, @Nullable String key) {
 	}
@@ -81,6 +85,7 @@ public static String getKeyFromMetadata(ModContainer mod) {
 	@Range(from = 0, to = Integer.MAX_VALUE)
 	public static int getModDataVersion(NbtCompound nbt, String modId, @Nullable String key) {
 		String nbtKey = modId;
+
 		if (key != null) {
 			nbtKey += ('_' + key);
 		}
@@ -89,6 +94,28 @@ public static int getModDataVersion(NbtCompound nbt, String modId, @Nullable Str
 		return dataVersions.getInt(nbtKey);
 	}
 
+	private static List<DataFixerEntrypoint> getEntrypoints() {
+		List<EntrypointContainer<DataFixerEntrypoint>> dataFixerEntrypoints = FabricLoader.getInstance()
+				.getEntrypointContainers(ENTRYPOINT_KEY, DataFixerEntrypoint.class);
+		return dataFixerEntrypoints.stream().map(EntrypointContainer::getEntrypoint).toList();
+	}
+
+	public static void registerBlockEntities(SchemaRegistry registry, Schema schema) {
+		List<DataFixerEntrypoint> entrypoints = getEntrypoints();
+
+		for (DataFixerEntrypoint entrypoint : entrypoints) {
+			entrypoint.onRegisterBlockEntities(registry, schema);
+		}
+	}
+
+	public static void registerEntities(SchemaRegistry registry, Schema schema) {
+		List<DataFixerEntrypoint> entrypoints = getEntrypoints();
+
+		for (DataFixerEntrypoint entrypoint : entrypoints) {
+			entrypoint.onRegisterEntities(registry, schema);
+		}
+	}
+
 	private static FabricDataFixesInternals instance;
 
 	public static FabricDataFixesInternals get() {
@@ -126,10 +153,6 @@ public static FabricDataFixesInternals get() {
 
 	public abstract NbtCompound addModDataVersions(NbtCompound nbt);
 
-	public abstract void registerBlockEntities(SchemaRegistry registry, Schema schema);
-
-	public abstract void registerEntities(SchemaRegistry registry, Schema schema);
-
 	public abstract void freeze();
 
 	@Contract(pure = true)

fabric-data-fixer-api-v1/src/main/java/net/fabricmc/fabric/impl/datafixer/v1/FabricDataFixesInternalsImpl.java

@@ -35,15 +35,8 @@
 import net.minecraft.nbt.NbtCompound;
 import net.minecraft.nbt.NbtElement;
 
-import net.fabricmc.fabric.api.datafixer.v1.DataFixerEntrypoint;
-import net.fabricmc.fabric.api.datafixer.v1.SchemaRegistry;
-import net.fabricmc.loader.api.FabricLoader;
-import net.fabricmc.loader.api.entrypoint.EntrypointContainer;
-
 public final class FabricDataFixesInternalsImpl extends FabricDataFixesInternals {
 	// From QSL.
-	private static final String ENTRYPOINT_KEY = "fabric-data-fixer";
-
 	private final Schema latestVanillaSchema;
 
 	private Map<String, List<DataFixerEntry>> modDataFixers;
@@ -118,30 +111,6 @@ public NbtCompound addModDataVersions(NbtCompound nbt) {
 		return nbt;
 	}
 
-	private static List<DataFixerEntrypoint> getEntrypoints() {
-		List<EntrypointContainer<DataFixerEntrypoint>> dataFixerEntrypoints = FabricLoader.getInstance()
-				.getEntrypointContainers(ENTRYPOINT_KEY, DataFixerEntrypoint.class);
-		return dataFixerEntrypoints.stream().map(EntrypointContainer::getEntrypoint).toList();
-	}
-
-	@Override
-	public void registerBlockEntities(SchemaRegistry registry, Schema schema) {
-		List<DataFixerEntrypoint> entrypoints = getEntrypoints();
-
-		for (DataFixerEntrypoint entrypoint : entrypoints) {
-			entrypoint.onRegisterBlockEntities(registry, schema);
-		}
-	}
-
-	@Override
-	public void registerEntities(SchemaRegistry registry, Schema schema) {
-		List<DataFixerEntrypoint> entrypoints = getEntrypoints();
-
-		for (DataFixerEntrypoint entrypoint : entrypoints) {
-			entrypoint.onRegisterEntities(registry, schema);
-		}
-	}
-
 	@Override
 	public void freeze() {
 		if (!this.frozen) {

fabric-data-fixer-api-v1/src/main/java/net/fabricmc/fabric/impl/datafixer/v1/FabricSchema1903.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.impl.datafixer.v1;
+
+import java.util.Map;
+import java.util.function.Supplier;
+
+import com.mojang.datafixers.schemas.Schema;
+import com.mojang.datafixers.types.templates.TypeTemplate;
+
+import net.minecraft.datafixer.schema.IdentifierNormalizingSchema;
+
+import net.fabricmc.fabric.api.datafixer.v1.SchemaRegistry;
+
+public class FabricSchema1903 extends IdentifierNormalizingSchema {
+	public SchemaRegistry registeredBlockEntities;
+	public SchemaRegistry registeredEntities;
+
+	public FabricSchema1903(int versionKey, Schema parent) {
+		super(versionKey, parent);
+	}
+
+	@Override
+	public Map<String, Supplier<TypeTemplate>> registerBlockEntities(Schema schema) {
+		Map<String, Supplier<TypeTemplate>> map = super.registerBlockEntities(schema);
+
+		this.registeredBlockEntities = new SchemaRegistryImpl();
+		FabricDataFixesInternals.registerBlockEntities(this.registeredBlockEntities, schema);
+
+		Map<String, Supplier<TypeTemplate>> modRegistry = this.registeredBlockEntities.get();
+
+		for (Map.Entry<String, Supplier<TypeTemplate>> entry : modRegistry.entrySet()) {
+			schema.register(map, entry.getKey(), entry.getValue());
+		}
+
+		return map;
+	}
+
+	@Override
+	public Map<String, Supplier<TypeTemplate>> registerEntities(Schema schema) {
+		Map<String, Supplier<TypeTemplate>> map = super.registerEntities(schema);
+
+		this.registeredEntities = new SchemaRegistryImpl();
+		FabricDataFixesInternals.registerEntities(this.registeredEntities, schema);
+
+		Map<String, Supplier<TypeTemplate>> modRegistry = this.registeredEntities.get();
+
+		for (Map.Entry<String, Supplier<TypeTemplate>> entry : modRegistry.entrySet()) {
+			schema.register(map, entry.getKey(), entry.getValue());
+		}
+
+		return map;
+	}
+}

fabric-data-fixer-api-v1/src/main/java/net/fabricmc/fabric/impl/datafixer/v1/NoOpFabricDataFixesInternals.java

@@ -32,7 +32,6 @@
 import net.minecraft.nbt.NbtElement;
 
 import net.fabricmc.fabric.api.datafixer.v1.EmptySchema;
-import net.fabricmc.fabric.api.datafixer.v1.SchemaRegistry;
 
 public final class NoOpFabricDataFixesInternals extends FabricDataFixesInternals {
 	// From QSL.
@@ -67,14 +66,6 @@ public NbtCompound addModDataVersions(NbtCompound nbt) {
 		return nbt;
 	}
 
-	@Override
-	public void registerBlockEntities(SchemaRegistry registry, Schema schema) {
-	}
-
-	@Override
-	public void registerEntities(SchemaRegistry registry, Schema schema) {
-	}
-
 	@Override
 	public void freeze() {
 		this.frozen = true;

fabric-data-fixer-api-v1/src/main/java/net/fabricmc/fabric/mixin/datafixer/v1/SchemasMixin.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.mixin.datafixer.v1;
+
+import com.mojang.datafixers.DataFixerBuilder;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.Slice;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+import net.minecraft.datafixer.Schemas;
+import net.minecraft.datafixer.TypeReferences;
+import net.minecraft.datafixer.fix.ChoiceTypesFix;
+
+import net.fabricmc.fabric.impl.datafixer.v1.FabricSchema1903;
+
+@Mixin(Schemas.class)
+public class SchemasMixin {
+	@Inject(
+			method = "build",
+			at = @At(
+					value = "INVOKE",
+					target = "Lcom/mojang/datafixers/DataFixerBuilder;addSchema(ILjava/util/function/BiFunction;)Lcom/mojang/datafixers/schemas/Schema;",
+					ordinal = 0
+			),
+			slice = @Slice(
+					from = @At(
+							value = "CONSTANT",
+							args = "intValue=1803"
+					)
+			),
+			remap = false
+	)
+	private static void addFabricFixers(DataFixerBuilder builder, CallbackInfo ci) {
+		FabricSchema1903 schema = (FabricSchema1903) builder.addSchema(1903, FabricSchema1903::new);
+
+		for (String id : schema.registeredBlockEntities.getKeys()) {
+			builder.addFixer(new ChoiceTypesFix(schema, "Add Fabric block entity " + id, TypeReferences.BLOCK_ENTITY));
+		}
+
+		for (String id : schema.registeredEntities.getKeys()) {
+			builder.addFixer(new ChoiceTypesFix(schema, "Add Fabric entity " + id, TypeReferences.ENTITY));
+		}
+
+		schema.registeredBlockEntities = null;
+		schema.registeredEntities = null;
+	}
+}

fabric-data-fixer-api-v1/src/main/resources/fabric-data-fixer-api-v1.mixins.json

@@ -5,8 +5,7 @@
   "mixins": [
     "DataFixTypesMixin",
     "NbtHelperMixin",
-    "Schema1904Mixin",
-    "Schema1906Mixin"
+    "SchemasMixin"
   ],
   "injectors": {
     "defaultRequire": 1