Add datagen support for dynamic registry tags (#1890)

* Add datagen support for dynamic registry tags

* Update according review

* Make FabricTagBuilder final

* Throw AssertionError if expected errors didn't happen

fabric-data-generation-api-v1/build.gradle

@@ -7,6 +7,10 @@ moduleDependencies(project, [
 		'fabric-networking-api-v1'
 ])
 
+dependencies {
+	testmodImplementation project(path: ':fabric-tag-extensions-v0', configuration: 'namedElements')
+}
+
 sourceSets {
 	testmod {
 		resources {

fabric-data-generation-api-v1/src/main/java/net/fabricmc/fabric/api/datagen/v1/provider/FabricTagProvider.java

@@ -20,6 +20,7 @@ import java.nio.file.Path;
 import java.util.Objects;
 import java.util.function.Function;
 
+import com.google.common.base.Preconditions;
 import org.jetbrains.annotations.Nullable;
 
 import net.minecraft.block.Block;
@@ -30,10 +31,14 @@ import net.minecraft.fluid.Fluid;
 import net.minecraft.item.Item;
 import net.minecraft.tag.Tag;
 import net.minecraft.util.Identifier;
+import net.minecraft.util.registry.BuiltinRegistries;
 import net.minecraft.util.registry.Registry;
+import net.minecraft.util.registry.RegistryKey;
 import net.minecraft.world.event.GameEvent;
 
 import net.fabricmc.fabric.api.datagen.v1.FabricDataGenerator;
+import net.fabricmc.fabric.impl.datagen.FabricDataGenHelper;
+import net.fabricmc.fabric.mixin.datagen.DynamicRegistryManagerAccessor;
 
 /**
  * Implement this class (or one of the inner classes) to generate a tag list.
@@ -47,6 +52,7 @@ import net.fabricmc.fabric.api.datagen.v1.FabricDataGenerator;
  * @see FluidTagProvider
  * @see EntityTypeTagProvider
  * @see GameEventTagProvider
+ * @see DynamicRegistryTagProvider
  */
 public abstract class FabricTagProvider<T> extends AbstractTagProvider<T> {
 	private final String path;
@@ -62,10 +68,15 @@ public abstract class FabricTagProvider<T> extends AbstractTagProvider<T> {
 	 * @param path The directory name to write the tag file names. Example: "blocks" or "items"
 	 * @param name The name used for {@link DataProvider#getName()}
 	 */
+	@SuppressWarnings({"unchecked", "rawtypes"})
 	protected FabricTagProvider(FabricDataGenerator dataGenerator, Registry<T> registry, String path, String name) {
 		super(dataGenerator, registry);
 		this.path = path;
 		this.name = name;
+
+		if (!(this instanceof DynamicRegistryTagProvider) && BuiltinRegistries.REGISTRIES.contains((RegistryKey) registry.getKey())) {
+			throw new IllegalArgumentException("Using FabricTagProvider to generate dynamic registry tags is not supported, Use DynamicRegistryTagProvider instead.");
+		}
 	}
 
 	/**
@@ -179,10 +190,29 @@ public abstract class FabricTagProvider<T> extends AbstractTagProvider<T> {
 		}
 	}
 
+	/**
+	 * Extend this class to create dynamic registry tags.
+	 */
+	public abstract static class DynamicRegistryTagProvider<T> extends FabricTagProvider<T> {
+		/**
+		 * Construct a new {@link DynamicRegistryTagProvider}.
+		 *
+		 * @param dataGenerator The data generator instance
+		 * @param registryKey The registry key of the dynamic registry
+		 * @param path The directory name to write the tag file names
+		 * @param name The name used for {@link DataProvider#getName()}
+		 * @throws IllegalArgumentException if the registry is static registry
+		 */
+		protected DynamicRegistryTagProvider(FabricDataGenerator dataGenerator, RegistryKey<? extends Registry<T>> registryKey, String path, String name) {
+			super(dataGenerator, FabricDataGenHelper.getFakeDynamicRegistry(), path, name);
+			Preconditions.checkArgument(DynamicRegistryManagerAccessor.getInfos().containsKey(registryKey), "Only dynamic registries are supported in this tag provider.");
+		}
+	}
+
 	/**
 	 * An extension to {@link net.minecraft.data.server.AbstractTagProvider.ObjectBuilder} that provides additional functionality.
 	 */
-	public static class FabricTagBuilder<T> extends ObjectBuilder<T> {
+	public final class FabricTagBuilder<T> extends ObjectBuilder<T> {
 		private final AbstractTagProvider.ObjectBuilder<T> parent;
 
 		private FabricTagBuilder(ObjectBuilder<T> parent) {
@@ -206,13 +236,35 @@ public abstract class FabricTagProvider<T> extends AbstractTagProvider<T> {
 		 * Add a single element to the tag.
 		 *
 		 * @return the {@link FabricTagBuilder} instance
+		 * @throws UnsupportedOperationException if the provider is an instance of {@link DynamicRegistryTagProvider}
+		 * @see #add(Identifier)
 		 */
 		@Override
 		public FabricTagBuilder<T> add(T element) {
+			assertStaticRegistry();
 			parent.add(element);
 			return this;
 		}
 
+		/**
+		 * Add a single element to the tag.
+		 *
+		 * @return the {@link FabricTagBuilder} instance
+		 */
+		public FabricTagBuilder<T> add(Identifier id) {
+			builder.add(id, source);
+			return this;
+		}
+
+		/**
+		 * Add a single element to the tag.
+		 *
+		 * @return the {@link FabricTagBuilder} instance
+		 */
+		public FabricTagBuilder<T> add(RegistryKey<? extends T> registryKey) {
+			return add(registryKey.getValue());
+		}
+
 		/**
 		 * Add an optional {@link Identifier} to the tag.
 		 *
@@ -256,5 +308,56 @@ public abstract class FabricTagProvider<T> extends AbstractTagProvider<T> {
 			parent.addOptionalTag(id);
 			return this;
 		}
+
+		/**
+		 * Add multiple elements to this tag.
+		 *
+		 * @return the {@link FabricTagBuilder} instance
+		 * @throws UnsupportedOperationException if the provider is an instance of {@link DynamicRegistryTagProvider}
+		 */
+		@SafeVarargs
+		@Override
+		public final FabricTagBuilder<T> add(T... elements) {
+			assertStaticRegistry();
+
+			for (T element : elements) {
+				add(element);
+			}
+
+			return this;
+		}
+
+		/**
+		 * Add multiple elements to this tag.
+		 *
+		 * @return the {@link FabricTagBuilder} instance
+		 */
+		public FabricTagBuilder<T> add(Identifier... ids) {
+			for (Identifier id : ids) {
+				add(id);
+			}
+
+			return this;
+		}
+
+		/**
+		 * Add multiple elements to this tag.
+		 *
+		 * @return the {@link FabricTagBuilder} instance
+		 */
+		@SafeVarargs
+		public final FabricTagBuilder<T> add(RegistryKey<? extends T>... registryKeys) {
+			for (RegistryKey<? extends T> registryKey : registryKeys) {
+				add(registryKey);
+			}
+
+			return this;
+		}
+
+		private void assertStaticRegistry() {
+			if (FabricTagProvider.this instanceof DynamicRegistryTagProvider) {
+				throw new UnsupportedOperationException("Adding object instances is not supported for DynamicRegistryTagProvider.");
+			}
+		}
 	}
 }

fabric-data-generation-api-v1/src/main/java/net/fabricmc/fabric/impl/datagen/FabricDataGenHelper.java

@@ -22,13 +22,21 @@ import java.nio.file.Paths;
 import java.util.List;
 import java.util.Objects;
 
+import com.mojang.serialization.Lifecycle;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Nullable;
 
+import net.minecraft.data.server.AbstractTagProvider;
+import net.minecraft.util.Identifier;
+import net.minecraft.util.registry.Registry;
+import net.minecraft.util.registry.RegistryKey;
+import net.minecraft.util.registry.SimpleRegistry;
+
 import net.fabricmc.fabric.api.datagen.v1.DataGeneratorEntrypoint;
 import net.fabricmc.fabric.api.datagen.v1.FabricDataGenerator;
+import net.fabricmc.fabric.api.datagen.v1.provider.FabricTagProvider.DynamicRegistryTagProvider;
 import net.fabricmc.loader.api.FabricLoader;
 import net.fabricmc.loader.api.entrypoint.EntrypointContainer;
 
@@ -62,6 +70,21 @@ public final class FabricDataGenHelper {
 	 */
 	private static final String ENTRYPOINT_KEY = "fabric-datagen";
 
+	/**
+	 * A fake registry instance to be used for {@link DynamicRegistryTagProvider}.
+	 *
+	 * <p>In {@link AbstractTagProvider#run}, it checks for whether the registry has all the elements added to the builder.
+	 * This would be fine for static registry, but there won't be any instance dynamic registry available.
+	 * Therefore, this simply return true for all {@link Registry#containsId} call.
+	 */
+	@SuppressWarnings("rawtypes")
+	private static final Registry FAKE_DYNAMIC_REGISTRY = new SimpleRegistry<>(RegistryKey.ofRegistry(new Identifier("fabric:fake_dynamic_registry")), Lifecycle.experimental()) {
+		@Override
+		public boolean containsId(Identifier id) {
+			return true;
+		}
+	};
+
 	private FabricDataGenHelper() {
 	}
 
@@ -89,4 +112,9 @@ public final class FabricDataGenHelper {
 			dataGenerator.run();
 		}
 	}
+
+	@SuppressWarnings("unchecked")
+	public static <T> Registry<T> getFakeDynamicRegistry() {
+		return FAKE_DYNAMIC_REGISTRY;
+	}
 }

fabric-data-generation-api-v1/src/main/java/net/fabricmc/fabric/mixin/datagen/DynamicRegistryManagerAccessor.java

@@ -0,0 +1,34 @@
+/*
+ * 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.datagen;
+
+import java.util.Map;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Accessor;
+
+import net.minecraft.util.registry.DynamicRegistryManager;
+import net.minecraft.util.registry.Registry;
+import net.minecraft.util.registry.RegistryKey;
+
+@Mixin(DynamicRegistryManager.class)
+public interface DynamicRegistryManagerAccessor {
+	@Accessor("INFOS")
+	static Map<RegistryKey<? extends Registry<?>>, ?> getInfos() {
+		throw new AssertionError();
+	}
+}

fabric-data-generation-api-v1/src/main/resources/fabric-data-generation-api-v1.accesswidener

@@ -3,6 +3,7 @@ accessWidener	v2	named
 accessible  field   net/minecraft/data/server/RecipesProvider   root    Lnet/minecraft/data/DataGenerator;
 
 extendable  method  net/minecraft/data/server/AbstractTagProvider$ObjectBuilder <init>  (Lnet/minecraft/tag/Tag$Builder;Lnet/minecraft/util/registry/Registry;Ljava/lang/String;)V
+extendable  method  net/minecraft/data/server/AbstractTagProvider$ObjectBuilder add     ([Ljava/lang/Object;)Lnet/minecraft/data/server/AbstractTagProvider$ObjectBuilder;
 accessible  field   net/minecraft/data/server/AbstractTagProvider$ObjectBuilder builder Lnet/minecraft/tag/Tag$Builder;
 accessible  field   net/minecraft/data/server/AbstractTagProvider$ObjectBuilder registry    Lnet/minecraft/util/registry/Registry;
 accessible  field   net/minecraft/data/server/AbstractTagProvider$ObjectBuilder source  Ljava/lang/String;

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

@@ -4,6 +4,7 @@
   "compatibilityLevel": "JAVA_16",
   "mixins": [
     "BlockStateDefinitionProviderMixin",
+    "DynamicRegistryManagerAccessor",
     "TagBuilderMixin"
   ],
   "client": [

fabric-data-generation-api-v1/src/testmod/java/net/fabricmc/fabric/test/datagen/DataGeneratorTestEntrypoint.java

@@ -39,6 +39,11 @@ import net.minecraft.tag.BlockTags;
 import net.minecraft.tag.ItemTags;
 import net.minecraft.text.TranslatableText;
 import net.minecraft.util.Identifier;
+import net.minecraft.util.registry.BuiltinRegistries;
+import net.minecraft.util.registry.Registry;
+import net.minecraft.world.biome.Biome;
+import net.minecraft.world.biome.BiomeKeys;
+import net.minecraft.world.biome.BuiltinBiomes;
 
 import net.fabricmc.fabric.api.datagen.v1.DataGeneratorEntrypoint;
 import net.fabricmc.fabric.api.datagen.v1.FabricDataGenerator;
@@ -48,6 +53,7 @@ import net.fabricmc.fabric.api.datagen.v1.provider.FabricBlockStateDefinitionPro
 import net.fabricmc.fabric.api.datagen.v1.provider.FabricRecipesProvider;
 import net.fabricmc.fabric.api.datagen.v1.provider.FabricTagProvider;
 import net.fabricmc.fabric.api.datagen.v1.provider.SimpleFabricLootTableProvider;
+import net.fabricmc.fabric.api.tag.TagFactory;
 
 public class DataGeneratorTestEntrypoint implements DataGeneratorEntrypoint {
 	@Override
@@ -60,6 +66,29 @@ public class DataGeneratorTestEntrypoint implements DataGeneratorEntrypoint {
 
 		TestBlockTagsProvider blockTagsProvider = dataGenerator.addProvider(TestBlockTagsProvider::new);
 		dataGenerator.addProvider(new TestItemTagsProvider(dataGenerator, blockTagsProvider));
+		dataGenerator.addProvider(TestBiomeTagsProvider::new);
+
+		try {
+			new FabricTagProvider<>(dataGenerator, BuiltinRegistries.BIOME, "biomes", "Biome Tags") {
+				@Override
+				protected void generateTags() {
+				}
+			};
+			throw new AssertionError("Using FabricTagProvider with built-in registry didn't throw an exception!");
+		} catch (IllegalArgumentException e) {
+			// no-op
+		}
+
+		try {
+			new FabricTagProvider.DynamicRegistryTagProvider<>(dataGenerator, Registry.ITEM_KEY, "items", "Item Tags") {
+				@Override
+				protected void generateTags() {
+				}
+			};
+			throw new AssertionError("Using DynamicRegistryTagProvider with static registry didn't throw an exception!");
+		} catch (IllegalArgumentException e) {
+			// no-op
+		}
 	}
 
 	private static class TestRecipeProvider extends FabricRecipesProvider {
@@ -113,6 +142,26 @@ public class DataGeneratorTestEntrypoint implements DataGeneratorEntrypoint {
 		}
 	}
 
+	private static class TestBiomeTagsProvider extends FabricTagProvider.DynamicRegistryTagProvider<Biome> {
+		private TestBiomeTagsProvider(FabricDataGenerator dataGenerator) {
+			super(dataGenerator, Registry.BIOME_KEY, "biomes", "Biome Tags");
+		}
+
+		@Override
+		protected void generateTags() {
+			FabricTagBuilder<Biome> builder = getOrCreateTagBuilder(TagFactory.BIOME.create(new Identifier(MOD_ID, "biome_tag_test")))
+					.add(BiomeKeys.BADLANDS, BiomeKeys.BAMBOO_JUNGLE)
+					.add(BiomeKeys.BASALT_DELTAS);
+
+			try {
+				builder.add(BuiltinBiomes.PLAINS);
+				throw new AssertionError("Adding built-in entry to dynamic registry tag builder didn't throw an exception!");
+			} catch (UnsupportedOperationException e) {
+				// no-op
+			}
+		}
+	}
+
 	private static class TestAdvancementsProvider extends FabricAdvancementsProvider {
 		private TestAdvancementsProvider(FabricDataGenerator dataGenerator) {
 			super(dataGenerator);

fabric-data-generation-api-v1/template.accesswidener

@@ -3,6 +3,7 @@ accessWidener	v2	named
 accessible  field   net/minecraft/data/server/RecipesProvider   root    Lnet/minecraft/data/DataGenerator;
 
 extendable  method  net/minecraft/data/server/AbstractTagProvider$ObjectBuilder <init>  (Lnet/minecraft/tag/Tag$Builder;Lnet/minecraft/util/registry/Registry;Ljava/lang/String;)V
+extendable  method  net/minecraft/data/server/AbstractTagProvider$ObjectBuilder add     ([Ljava/lang/Object;)Lnet/minecraft/data/server/AbstractTagProvider$ObjectBuilder;
 accessible  field   net/minecraft/data/server/AbstractTagProvider$ObjectBuilder builder Lnet/minecraft/tag/Tag$Builder;
 accessible  field   net/minecraft/data/server/AbstractTagProvider$ObjectBuilder registry    Lnet/minecraft/util/registry/Registry;
 accessible  field   net/minecraft/data/server/AbstractTagProvider$ObjectBuilder source  Ljava/lang/String;