Registry resource conditions (#3752)

* Add resource conditions to dynamic registries

* Support conditions in FabricDynamicRegistryProvider

 - Add matching overloads for all add(...) methods which take a varargs
   list of ResourceConditions.

 - Add an additional overload for registering RegistryEntry.References
   directly. This makes it a little easier to register directly from a
   RegistryBuilder.

* Throw error if we cannot add resource conditions

* Rename EntryWithConditions to ConditionalEntry

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

@@ -31,6 +31,7 @@ import com.mojang.serialization.DynamicOps;
 import com.mojang.serialization.Encoder;
 import com.mojang.serialization.JsonOps;
 import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -52,6 +53,8 @@ import net.minecraft.world.gen.feature.PlacedFeature;
 
 import net.fabricmc.fabric.api.datagen.v1.FabricDataOutput;
 import net.fabricmc.fabric.api.event.registry.DynamicRegistries;
+import net.fabricmc.fabric.api.resource.conditions.v1.ResourceCondition;
+import net.fabricmc.fabric.impl.datagen.FabricDataGenHelper;
 import net.fabricmc.fabric.impl.registry.sync.DynamicRegistriesImpl;
 
 /**
@@ -129,10 +132,45 @@ public abstract class FabricDynamicRegistryProvider implements DataProvider {
 		/**
 		 * Adds a new object to be data generated.
 		 *
+		 * @param key    The key of the resource to register.
+		 * @param object The object to register.
 		 * @return a reference to it for use in other objects.
 		 */
-		public <T> RegistryEntry<T> add(RegistryKey<T> registry, T object) {
-			return getQueuedEntries(registry).add(registry.getValue(), object);
+		public <T> RegistryEntry<T> add(RegistryKey<T> key, T object) {
+			return getQueuedEntries(key).add(key, object, null);
+		}
+
+		/**
+		 * Adds a new object to be data generated with several resource conditions.
+		 *
+		 * @param key        The key of the resource to register.
+		 * @param object     The object to register.
+		 * @param conditions Conditions that must be satisfied to load this object.
+		 * @return a reference to it for use in other objects.
+		 */
+		public <T> RegistryEntry<T> add(RegistryKey<T> key, T object, ResourceCondition... conditions) {
+			return getQueuedEntries(key).add(key, object, conditions);
+		}
+
+		/**
+		 * Adds a new object to be data generated.
+		 *
+		 * @param object The object to generate. This registry entry must have both a
+		 *               {@linkplain RegistryEntry#hasKeyAndValue() key and value}.
+		 */
+		public <T> void add(RegistryEntry.Reference<T> object) {
+			add(object.registryKey(), object.value());
+		}
+
+		/**
+		 * Adds a new object to be data generated with several resource conditions.
+		 *
+		 * @param object     The object to generate. This registry entry must have both a
+		 *                   {@linkplain RegistryEntry#hasKeyAndValue() key and value}.
+		 * @param conditions Conditions that must be satisfied to load this object.
+		 */
+		public <T> void add(RegistryEntry.Reference<T> object, ResourceCondition... conditions) {
+			add(object.registryKey(), object.value(), conditions);
 		}
 
 		/**
@@ -144,6 +182,16 @@ public abstract class FabricDynamicRegistryProvider implements DataProvider {
 			return add(valueKey, registry.getOrThrow(valueKey).value());
 		}
 
+		/**
+		 * Adds a new {@link RegistryKey} from a given {@link RegistryWrapper.Impl} to be data generated.
+		 *
+		 * @param conditions Conditions that must be satisfied to load this object.
+		 * @return a reference to it for use in other objects.
+		 */
+		public <T> RegistryEntry<T> add(RegistryWrapper.Impl<T> registry, RegistryKey<T> valueKey, ResourceCondition... conditions) {
+			return add(valueKey, registry.getOrThrow(valueKey).value(), conditions);
+		}
+
 		/**
 		 * All the registry entries whose namespace matches the current effective mod ID will be data generated.
 		 */
@@ -166,11 +214,14 @@ public abstract class FabricDynamicRegistryProvider implements DataProvider {
 		}
 	}
 
+	private record ConditionalEntry<T>(T value, @Nullable ResourceCondition... conditions) {
+	}
+
 	private static class RegistryEntries<T> {
 		final RegistryEntryOwner<T> lookup;
 		final RegistryKey<? extends Registry<T>> registry;
 		final Codec<T> elementCodec;
-		Map<RegistryKey<T>, T> entries = new IdentityHashMap<>();
+		Map<RegistryKey<T>, ConditionalEntry<T>> entries = new IdentityHashMap<>();
 
 		RegistryEntries(RegistryEntryOwner<T> lookup,
 						RegistryKey<? extends Registry<T>> registry,
@@ -185,17 +236,13 @@ public abstract class FabricDynamicRegistryProvider implements DataProvider {
 			return new RegistryEntries<>(lookup, loaderEntry.key(), loaderEntry.elementCodec());
 		}
 
-		public RegistryEntry<T> add(RegistryKey<T> key, T value) {
-			if (entries.put(key, value) != null) {
+		RegistryEntry<T> add(RegistryKey<T> key, T value, @Nullable ResourceCondition[] conditions) {
+			if (entries.put(key, new ConditionalEntry<>(value, conditions)) != null) {
 				throw new IllegalArgumentException("Trying to add registry key " + key + " more than once.");
 			}
 
 			return RegistryEntry.Reference.standAlone(lookup, key);
 		}
-
-		public RegistryEntry<T> add(Identifier id, T value) {
-			return add(RegistryKey.of(registry, id), value);
-		}
 	}
 
 	@Override
@@ -227,21 +274,31 @@ public abstract class FabricDynamicRegistryProvider implements DataProvider {
 		final DataOutput.PathResolver pathResolver = output.getResolver(DataOutput.OutputType.DATA_PACK, directoryName);
 		final List<CompletableFuture<?>> futures = new ArrayList<>();
 
-		for (Map.Entry<RegistryKey<T>, T> entry : entries.entries.entrySet()) {
+		for (Map.Entry<RegistryKey<T>, ConditionalEntry<T>> entry : entries.entries.entrySet()) {
 			Path path = pathResolver.resolveJson(entry.getKey().getValue());
-			futures.add(writeToPath(path, writer, ops, entries.elementCodec, entry.getValue()));
+			futures.add(writeToPath(path, writer, ops, entries.elementCodec, entry.getValue().value(), entry.getValue().conditions()));
 		}
 
 		return CompletableFuture.allOf(futures.toArray(CompletableFuture[]::new));
 	}
 
-	private static <E> CompletableFuture<?> writeToPath(Path path, DataWriter cache, DynamicOps<JsonElement> json, Encoder<E> encoder, E value) {
+	private static <E> CompletableFuture<?> writeToPath(Path path, DataWriter cache, DynamicOps<JsonElement> json, Encoder<E> encoder, E value, @Nullable ResourceCondition[] conditions) {
 		Optional<JsonElement> optional = encoder.encodeStart(json, value).resultOrPartial((error) -> {
 			LOGGER.error("Couldn't serialize element {}: {}", path, error);
 		});
 
 		if (optional.isPresent()) {
-			return DataProvider.writeToPath(cache, optional.get(), path);
+			JsonElement jsonElement = optional.get();
+
+			if (conditions != null && conditions.length > 0) {
+				if (!jsonElement.isJsonObject()) {
+					throw new IllegalStateException("Cannot add conditions to " + path + ": JSON is a non-object value");
+				} else {
+					FabricDataGenHelper.addConditions(jsonElement.getAsJsonObject(), conditions);
+				}
+			}
+
+			return DataProvider.writeToPath(cache, jsonElement, path);
 		}
 
 		return CompletableFuture.completedFuture(null);

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

@@ -434,7 +434,10 @@ public class DataGeneratorTestEntrypoint implements DataGeneratorEntrypoint {
 
 		@Override
 		protected void configure(RegistryWrapper.WrapperLookup registries, Entries entries) {
-			entries.add(registries.getWrapperOrThrow(TEST_DATAGEN_DYNAMIC_REGISTRY_KEY), TEST_DYNAMIC_REGISTRY_ITEM_KEY);
+			entries.add(
+					registries.getWrapperOrThrow(TEST_DATAGEN_DYNAMIC_REGISTRY_KEY), TEST_DYNAMIC_REGISTRY_ITEM_KEY,
+					ResourceConditions.allModsLoaded(MOD_ID)
+			);
 		}
 
 		@Override

fabric-resource-conditions-api-v1/src/main/java/net/fabricmc/fabric/mixin/resource/conditions/RegistryLoaderMixin.java

@@ -0,0 +1,82 @@
+/*
+ * 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.resource.conditions;
+
+import java.io.IOException;
+import java.io.Reader;
+import java.util.List;
+
+import com.google.gson.JsonElement;
+import com.llamalad7.mixinextras.injector.wrapoperation.Operation;
+import com.llamalad7.mixinextras.injector.wrapoperation.WrapOperation;
+import com.llamalad7.mixinextras.sugar.Local;
+import com.mojang.serialization.Decoder;
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.Unique;
+import org.spongepowered.asm.mixin.injection.At;
+import org.spongepowered.asm.mixin.injection.Coerce;
+import org.spongepowered.asm.mixin.injection.Inject;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+
+import net.minecraft.registry.DynamicRegistryManager;
+import net.minecraft.registry.MutableRegistry;
+import net.minecraft.registry.RegistryKey;
+import net.minecraft.registry.RegistryLoader;
+import net.minecraft.registry.RegistryOps;
+import net.minecraft.registry.entry.RegistryEntryInfo;
+import net.minecraft.resource.Resource;
+
+import net.fabricmc.fabric.impl.resource.conditions.ResourceConditionsImpl;
+
+@Mixin(RegistryLoader.class)
+public class RegistryLoaderMixin {
+	@Unique
+	private static final ThreadLocal<DynamicRegistryManager> REGISTRIES = new ThreadLocal<>();
+
+	/**
+	 * Capture the current registries, so they can be passed to the resource conditions.
+	 */
+	@WrapOperation(method = "loadFromResource(Lnet/minecraft/resource/ResourceManager;Lnet/minecraft/registry/DynamicRegistryManager;Ljava/util/List;)Lnet/minecraft/registry/DynamicRegistryManager$Immutable;", at = @At(value = "INVOKE", target = "Lnet/minecraft/registry/RegistryLoader;load(Lnet/minecraft/registry/RegistryLoader$RegistryLoadable;Lnet/minecraft/registry/DynamicRegistryManager;Ljava/util/List;)Lnet/minecraft/registry/DynamicRegistryManager$Immutable;"))
+	private static DynamicRegistryManager.Immutable captureRegistries(@Coerce Object registryLoadable, DynamicRegistryManager baseRegistryManager, List<RegistryLoader.Entry<?>> entries, Operation<DynamicRegistryManager.Immutable> original) {
+		try {
+			REGISTRIES.set(baseRegistryManager);
+			return original.call(registryLoadable, baseRegistryManager, entries);
+		} finally {
+			REGISTRIES.remove();
+		}
+	}
+
+	@Inject(
+			method = "Lnet/minecraft/registry/RegistryLoader;parseAndAdd(Lnet/minecraft/registry/MutableRegistry;Lcom/mojang/serialization/Decoder;Lnet/minecraft/registry/RegistryOps;Lnet/minecraft/registry/RegistryKey;Lnet/minecraft/resource/Resource;Lnet/minecraft/registry/entry/RegistryEntryInfo;)V",
+			at = @At(value = "INVOKE_ASSIGN", target = "Lcom/google/gson/JsonParser;parseReader(Ljava/io/Reader;)Lcom/google/gson/JsonElement;", remap = false),
+			cancellable = true
+	)
+	private static <E> void checkResourceCondition(
+			MutableRegistry<E> registry, Decoder<E> decoder, RegistryOps<JsonElement> ops, RegistryKey<E> key, Resource resource, RegistryEntryInfo entryInfo,
+			CallbackInfo ci, @Local Reader reader, @Local JsonElement json
+	) throws IOException {
+		// This method is called both on the server (when loading resources) and on the client (when syncing from the
+		// server). We only want to apply resource conditions when loading via loadFromResource.
+		DynamicRegistryManager registries = REGISTRIES.get();
+		if (registries == null) return;
+
+		if (json.isJsonObject() && !ResourceConditionsImpl.applyResourceConditions(json.getAsJsonObject(), key.getRegistry().toString(), key.getValue(), registries)) {
+			reader.close();
+			ci.cancel();
+		}
+	}
+}

fabric-resource-conditions-api-v1/src/main/resources/fabric-resource-conditions-api-v1.mixins.json

@@ -7,6 +7,7 @@
     "DataProviderMixin",
     "JsonDataLoaderMixin",
     "RecipeManagerMixin",
+    "RegistryLoaderMixin",
     "ReloadableRegistriesMixin",
     "ServerAdvancementLoaderMixin",
     "SinglePreparationResourceReloaderMixin",

fabric-resource-conditions-api-v1/src/testmod/java/net/fabricmc/fabric/test/resource/conditions/ConditionalResourcesTest.java

@@ -16,8 +16,10 @@
 
 package net.fabricmc.fabric.test.resource.conditions;
 
+import net.minecraft.block.entity.BannerPattern;
 import net.minecraft.loot.LootTable;
 import net.minecraft.recipe.RecipeManager;
+import net.minecraft.registry.Registry;
 import net.minecraft.registry.RegistryKey;
 import net.minecraft.registry.RegistryKeys;
 import net.minecraft.registry.ReloadableRegistries;
@@ -104,4 +106,19 @@ public class ConditionalResourcesTest {
 
 		context.complete();
 	}
+
+	@GameTest(templateName = FabricGameTest.EMPTY_STRUCTURE)
+	public void conditionalDynamicRegistry(TestContext context) {
+		Registry<BannerPattern> registry = context.getWorld().getRegistryManager().get(RegistryKeys.BANNER_PATTERN);
+
+		if (registry.get(id("loaded")) == null) {
+			throw new AssertionError("loaded banner pattern should have been loaded.");
+		}
+
+		if (registry.get(id("not_loaded")) != null) {
+			throw new AssertionError("not_loaded banner pattern should not have been loaded.");
+		}
+
+		context.complete();
+	}
 }

fabric-resource-conditions-api-v1/src/testmod/resources/data/fabric-resource-conditions-api-v1-testmod/banner_pattern/loaded.json

@@ -0,0 +1,12 @@
+{
+  "asset_id": "minecraft:loaded",
+  "translation_key": "block.minecraft.banner.loaded",
+  "fabric:load_conditions": [
+    {
+      "condition": "fabric:all_mods_loaded",
+      "values": [
+        "fabric-resource-conditions-api-v1"
+      ]
+    }
+  ]
+}

fabric-resource-conditions-api-v1/src/testmod/resources/data/fabric-resource-conditions-api-v1-testmod/banner_pattern/not_loaded.json

@@ -0,0 +1,21 @@
+{
+  "asset_id": "minecraft:not_loaded",
+  "translation_key": "block.minecraft.banner.not_loaded",
+  "fabric:load_conditions": [
+    {
+      "condition": "fabric:not",
+      "value": {
+        "condition": "fabric:all_mods_loaded",
+        "values": [
+          "fabric-resource-conditions-api-v1"
+        ]
+      }
+    },
+    {
+      "condition": "fabric:all_mods_loaded",
+      "values": [
+        "fabric-resource-conditions-api-v1"
+      ]
+    }
+  ]
+}