Fix registry resource conditions on loot tables. (#3432)

* Fix registry resource conditions on loot tables.

* assert we dont already have the DRM set on this thread.

* Use a sync map

(cherry picked from commit de5c6cc)

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

@@ -16,17 +16,27 @@
 
 package net.fabricmc.fabric.mixin.resource.conditions;
 
+import java.util.Collections;
+import java.util.IdentityHashMap;
 import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
 
 import com.google.gson.JsonElement;
 import com.google.gson.JsonObject;
 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.Inject;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
+import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
+import org.spongepowered.asm.mixin.injection.callback.LocalCapture;
 
 import net.minecraft.loot.LootDataType;
 import net.minecraft.loot.LootManager;
+import net.minecraft.registry.DynamicRegistryManager;
+import net.minecraft.resource.ResourceManager;
 import net.minecraft.util.Identifier;
 
 import net.fabricmc.fabric.api.resource.conditions.v1.ResourceConditions;
@@ -37,6 +47,22 @@
  */
 @Mixin(LootManager.class)
 public class LootManagerMixin {
+	// Keep track of the DynamicRegistryManager instance by assgining it to the map that is passed to the async runnable.
+	@Unique
+	private static final Map<Object, DynamicRegistryManager.Immutable> dynamicRegistryManagerMap = Collections.synchronizedMap(new IdentityHashMap<>());
+
+	@Inject(method = "load", at = @At(value = "INVOKE", target = "Ljava/util/concurrent/CompletableFuture;runAsync(Ljava/lang/Runnable;Ljava/util/concurrent/Executor;)Ljava/util/concurrent/CompletableFuture;"), locals = LocalCapture.CAPTURE_FAILHARD)
+	private static void load(LootDataType type, ResourceManager resourceManager, Executor executor, Map<LootDataType<?>, Map<Identifier, ?>> results, CallbackInfoReturnable<CompletableFuture<?>> cir, Map map) {
+		dynamicRegistryManagerMap.put(map, ResourceConditionsImpl.CURRENT_REGISTRIES.get());
+	}
+
+	// runAsync Runnable in load method
+	@Inject(method = "method_51189", at = @At("HEAD"))
+	private static void runAsync(ResourceManager resourceManager, LootDataType lootDataType, Map map, CallbackInfo ci) {
+		assert ResourceConditionsImpl.CURRENT_REGISTRIES.get() == null;
+		ResourceConditionsImpl.CURRENT_REGISTRIES.set(Objects.requireNonNull(dynamicRegistryManagerMap.remove(map)));
+	}
+
 	// forEach in load method
 	@Inject(method = "method_51195", at = @At("HEAD"), cancellable = true)
 	private static void applyResourceConditions(LootDataType lootDataType, Map map, Identifier id, JsonElement json, CallbackInfo ci) {
@@ -57,4 +83,10 @@ private static void applyResourceConditions(LootDataType lootDataType, Map map,
 			}
 		}
 	}
+
+	// runAsync Runnable in load method
+	@Inject(method = "method_51189", at = @At("RETURN"))
+	private static void runAsyncEnd(ResourceManager resourceManager, LootDataType lootDataType, Map map, CallbackInfo ci) {
+		ResourceConditionsImpl.CURRENT_REGISTRIES.remove();
+	}
 }

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

@@ -9,5 +9,8 @@
     "LootManagerMixin",
     "SinglePreparationResourceReloaderMixin",
     "TagManagerLoaderMixin"
-  ]
+  ],
+  "injectors": {
+    "defaultRequire": 1
+  }
 }

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

@@ -87,4 +87,19 @@ public void conditionalPredicates(TestContext context) {
 
 		context.complete();
 	}
+
+	@GameTest(templateName = FabricGameTest.EMPTY_STRUCTURE)
+	public void conditionalLootTables(TestContext context) {
+		LootManager manager = context.getWorld().getServer().getLootManager();
+
+		if (manager.getElementOptional(LootDataType.LOOT_TABLES, id("blocks/loaded")).isEmpty()) {
+			throw new AssertionError("loaded loot table should have been loaded.");
+		}
+
+		if (manager.getElementOptional(LootDataType.LOOT_TABLES, id("blocks/not_loaded")).isPresent()) {
+			throw new AssertionError("not_loaded loot table should not have been loaded.");
+		}
+
+		context.complete();
+	}
 }

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

@@ -0,0 +1,29 @@
+{
+  "type": "minecraft:block",
+  "pools": [
+    {
+      "bonus_rolls": 0.0,
+      "conditions": [
+        {
+          "condition": "minecraft:survives_explosion"
+        }
+      ],
+      "entries": [
+        {
+          "type": "minecraft:item",
+          "name": "minecraft:cobblestone"
+        }
+      ],
+      "rolls": 1.0
+    }
+  ],
+  "random_sequence": "minecraft:blocks/cobblestone",
+  "fabric:load_conditions": [
+    {
+      "condition": "fabric:registry_contains",
+      "values": [
+        "minecraft:cobblestone"
+      ]
+    }
+  ]
+}

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

@@ -0,0 +1,29 @@
+{
+  "type": "minecraft:block",
+  "pools": [
+    {
+      "bonus_rolls": 0.0,
+      "conditions": [
+        {
+          "condition": "minecraft:survives_explosion"
+        }
+      ],
+      "entries": [
+        {
+          "type": "minecraft:item",
+          "name": "minecraft:cobblestone"
+        }
+      ],
+      "rolls": 1.0
+    }
+  ],
+  "random_sequence": "minecraft:blocks/cobblestone",
+  "fabric:load_conditions": [
+    {
+      "condition": "fabric:registry_contains",
+      "values": [
+        "minecraft:does_not_exist"
+      ]
+    }
+  ]
+}