Add FabricLootTableBuilder.modifyPools (#2314)

* Add FabricLootTableBuilder.modifyPools

* Fix typo in javadoc

fabric-loot-api-v2/src/main/java/net/fabricmc/fabric/api/loot/v2/FabricLootTableBuilder.java

@@ -18,6 +18,7 @@ package net.fabricmc.fabric.api.loot.v2;
 
 import java.util.Collection;
 import java.util.List;
+import java.util.function.Consumer;
 
 import org.jetbrains.annotations.ApiStatus;
 
@@ -29,7 +30,7 @@ import net.fabricmc.fabric.mixin.loot.LootTableAccessor;
 
 /**
  * Convenience extensions to {@link LootTable.Builder}
- * for adding pre-built objects or collections.
+ * for adding pre-built objects or collections and modifying loot pools.
  *
  * <p>This interface is automatically injected to {@link LootTable.Builder}.
  */
@@ -75,6 +76,23 @@ public interface FabricLootTableBuilder {
 		throw new UnsupportedOperationException("Implemented via mixin");
 	}
 
+	/**
+	 * Modifies all loot pools already present in this builder.
+	 *
+	 * <p>This method can be used instead of simply adding a new pool
+	 * when you want the loot table to only drop items from one of the loot pool entries
+	 * instead of both.
+	 *
+	 * <p>Calling this method turns all pools into builders and rebuilds them back into loot pools afterwards,
+	 * so it is more efficient to do all transformations with one {@code modifyPools} call.
+	 *
+	 * @param modifier the modifying function
+	 * @return this builder
+	 */
+	default LootTable.Builder modifyPools(Consumer<? super LootPool.Builder> modifier) {
+		throw new UnsupportedOperationException("Implemented via mixin");
+	}
+
 	/**
 	 * Creates a builder copy of a loot table.
 	 *

fabric-loot-api-v2/src/main/java/net/fabricmc/fabric/api/loot/v2/LootTableEvents.java

@@ -54,6 +54,13 @@ public final class LootTableEvents {
 	 * They have the loot table source {@link LootTableSource#REPLACED}.
 	 *
 	 * <h2>Example: adding diamonds to the cobblestone loot table</h2>
+	 * We'll add a new diamond {@linkplain net.minecraft.loot.LootPool loot pool} to the cobblestone loot table
+	 * that will be dropped alongside the original cobblestone loot pool.
+	 *
+	 * <p>If you want only one of the items to drop, you can use
+	 * {@link FabricLootTableBuilder#modifyPools(java.util.function.Consumer)} to add the new item to
+	 * the original loot pool instead.
+	 *
 	 * <pre>
 	 * {@code
 	 * LootTableEvents.MODIFY.register((resourceManager, lootManager, id, tableBuilder, source) -> {

fabric-loot-api-v2/src/main/java/net/fabricmc/fabric/mixin/loot/LootTableBuilderMixin.java

@@ -18,6 +18,8 @@ package net.fabricmc.fabric.mixin.loot;
 
 import java.util.Collection;
 import java.util.List;
+import java.util.ListIterator;
+import java.util.function.Consumer;
 
 import org.spongepowered.asm.mixin.Final;
 import org.spongepowered.asm.mixin.Mixin;
@@ -28,6 +30,7 @@ import net.minecraft.loot.LootPool;
 import net.minecraft.loot.LootTable;
 import net.minecraft.loot.function.LootFunction;
 
+import net.fabricmc.fabric.api.loot.v2.FabricLootPoolBuilder;
 import net.fabricmc.fabric.api.loot.v2.FabricLootTableBuilder;
 
 /**
@@ -73,4 +76,17 @@ abstract class LootTableBuilderMixin implements FabricLootTableBuilder {
 		this.functions.addAll(functions);
 		return self();
 	}
+
+	@Override
+	public LootTable.Builder modifyPools(Consumer<? super LootPool.Builder> modifier) {
+		ListIterator<LootPool> iterator = pools.listIterator();
+
+		while (iterator.hasNext()) {
+			LootPool.Builder poolBuilder = FabricLootPoolBuilder.copyOf(iterator.next());
+			modifier.accept(poolBuilder);
+			iterator.set(poolBuilder.build());
+		}
+
+		return self();
+	}
 }

fabric-loot-api-v2/src/testmod/java/net/fabricmc/fabric/test/loot/LootTest.java

@@ -85,6 +85,12 @@ public class LootTest implements ModInitializer {
 			if (Blocks.RED_WOOL.getLootTableId().equals(id) && source != LootTableSource.MOD) {
 				throw new AssertionError("red wool loot table should have LootTableSource.MOD");
 			}
+
+			// Modify yellow wool to drop *either* yellow wool or emeralds by adding
+			// emeralds to the same loot pool.
+			if (Blocks.YELLOW_WOOL.getLootTableId().equals(id)) {
+				tableBuilder.modifyPools(poolBuilder -> poolBuilder.with(ItemEntry.builder(Items.EMERALD)));
+			}
 		});
 	}
 }