Merge branch 'pr/54'

src/main/java/net/frozenblock/lib/loot/LootTableModifier.java

@@ -0,0 +1,30 @@
+package net.frozenblock.lib.loot;
+
+import net.fabricmc.fabric.api.loot.v3.LootTableEvents;
+import net.minecraft.resources.ResourceKey;
+import net.minecraft.world.level.storage.loot.LootTable;
+
+public class LootTableModifier {
+
+	public static void editTable(ResourceKey<LootTable> targetLootTable, boolean requiresBuiltIn, Edit listener) {
+		LootTableEvents.Replace modification = (id, lootTable, source, registries) -> {
+			if ((requiresBuiltIn && !source.isBuiltin()) || !targetLootTable.equals(id)) return null;
+			MutableLootTable mutableLootTable = new MutableLootTable(lootTable);
+			listener.editLootTable(id, mutableLootTable);
+			return mutableLootTable.build();
+		};
+
+		LootTableEvents.REPLACE.register(modification);
+	}
+
+	@FunctionalInterface
+	public interface Edit {
+		/**
+		 * Edits loot tables.
+		 *
+		 * @param id The loot table key.
+		 * @param mutableLootTable The mutable copy of the loot table.
+		 */
+		void editLootTable(ResourceKey<LootTable> id, MutableLootTable mutableLootTable);
+	}
+}

src/main/java/net/frozenblock/lib/loot/MutableLootItem.java

@@ -0,0 +1,78 @@
+package net.frozenblock.lib.loot;
+
+import java.util.List;
+import net.minecraft.core.Holder;
+import net.minecraft.world.item.Item;
+import net.minecraft.world.level.ItemLike;
+import net.minecraft.world.level.storage.loot.entries.LootItem;
+import net.minecraft.world.level.storage.loot.functions.LootItemFunction;
+import net.minecraft.world.level.storage.loot.predicates.LootItemCondition;
+import org.jetbrains.annotations.NotNull;
+
+public class MutableLootItem {
+	public final List<LootItemCondition> conditions;
+	public int weight;
+	public int quality;
+	public final List<LootItemFunction> functions;
+	public Holder<Item> item;
+
+	public MutableLootItem(LootItem original) {
+		this.conditions = original.conditions;
+		this.weight = original.weight;
+		this.quality = original.quality;
+		this.functions = original.functions;
+		this.item = original.item;
+	}
+
+	public MutableLootItem(Holder<Item> item, int weight, int quality, List<LootItemCondition> conditions, List<LootItemFunction> functions) {
+		this.item = item;
+		this.conditions = conditions;
+		this.weight = weight;
+		this.quality = quality;
+		this.functions = functions;
+	}
+
+	public MutableLootItem(@NotNull ItemLike item, int weight, int quality, List<LootItemCondition> conditions, List<LootItemFunction> functions) {
+		this.item = item.asItem().builtInRegistryHolder();
+		this.conditions = conditions;
+		this.weight = weight;
+		this.quality = quality;
+		this.functions = functions;
+	}
+
+	public LootItem build() {
+		return new LootItem(item, weight, quality, conditions, functions);
+	}
+
+	public int getWeight() {
+		return weight;
+	}
+
+	public void setWeight(int weight) {
+		this.weight = weight;
+	}
+
+	public int getQuality() {
+		return quality;
+	}
+
+	public void setQuality(int quality) {
+		this.quality = quality;
+	}
+
+	public Holder<Item> getItemHolder() {
+		return item;
+	}
+
+	public Item getItem() {
+		return item.value();
+	}
+
+	public void setItem(Holder<Item> item) {
+		this.item = item;
+	}
+
+	public void setItem(@NotNull ItemLike item) {
+		this.item = item.asItem().builtInRegistryHolder();
+	}
+}

src/main/java/net/frozenblock/lib/loot/MutableLootPool.java

@@ -0,0 +1,226 @@
+package net.frozenblock.lib.loot;
+
+import java.util.ArrayList;
+import java.util.function.Predicate;
+import net.minecraft.tags.TagKey;
+import net.minecraft.world.item.Item;
+import net.minecraft.world.level.ItemLike;
+import net.minecraft.world.level.storage.loot.LootPool;
+import net.minecraft.world.level.storage.loot.entries.LootItem;
+import net.minecraft.world.level.storage.loot.entries.LootPoolEntryContainer;
+import net.minecraft.world.level.storage.loot.functions.LootItemFunction;
+import net.minecraft.world.level.storage.loot.predicates.LootItemCondition;
+import net.minecraft.world.level.storage.loot.providers.number.NumberProvider;
+import org.jetbrains.annotations.NotNull;
+
+public class MutableLootPool {
+	public ArrayList<LootPoolEntryContainer> entries = new ArrayList<>();
+	public ArrayList<LootItemCondition> conditions = new ArrayList<>();
+	public ArrayList<LootItemFunction> functions = new ArrayList<>();
+	public NumberProvider rolls;
+	public NumberProvider bonusRolls;
+
+	public MutableLootPool(@NotNull LootPool lootPool) {
+		entries.addAll(lootPool.entries);
+		conditions.addAll(lootPool.conditions);
+		functions.addAll(lootPool.functions);
+		rolls = lootPool.rolls;
+		bonusRolls = lootPool.bonusRolls;
+	}
+
+	public LootPool build() {
+		LootPool.Builder builder = LootPool.lootPool();
+		entries.forEach(builder.entries::add);
+		conditions.forEach(builder.conditions::add);
+		functions.forEach(builder.functions::add);
+		builder.setRolls(rolls);
+		builder.setBonusRolls(bonusRolls);
+		return builder.build();
+	}
+
+	/**
+	 * Adds one item to the loot pool
+	 *
+	 * @param item    item to add
+	 * @param weight  how likely the item is to get drawn from the pool
+	 * @param builder idk lol
+	 * @return this
+	 */
+	public MutableLootPool add(ItemLike item, int weight, LootItemFunction.Builder builder) {
+		entries.add(LootItem.lootTableItem(item).setWeight(weight).apply(builder).build());
+		return this;
+	}
+
+	/**
+	 * Adds one or more items to the loot pool with the same weight
+	 *
+	 * @param items   items to add
+	 * @param weight  how likely the items are to get drawn from the pool
+	 * @param builder idk lol
+	 * @return this
+	 */
+	public MutableLootPool addAll(int weight, LootItemFunction.Builder builder, ItemLike @NotNull ... items) {
+		for (ItemLike item : items) {
+			entries.add(LootItem.lootTableItem(item).setWeight(weight).apply(builder).build());
+		}
+		return this;
+	}
+
+	public MutableLootPool remove(ItemLike item) {
+		for (LootPoolEntryContainer entryContainer : entries) {
+			if (entryContainer instanceof LootItem lootItem) {
+				if (item.equals(lootItem.item.value())) {
+					entries.remove(entryContainer);
+					return this;
+				}
+			}
+		}
+		// Failed to remove item from the loot pool.
+		return this;
+	}
+
+	public MutableLootPool remove(ItemLike... items) {
+		ArrayList<LootPoolEntryContainer> toRemove = new ArrayList<>();
+		for (LootPoolEntryContainer entryContainer : entries) {
+			if (entryContainer instanceof LootItem lootItem) {
+				for (ItemLike item : items) {
+					if (item.equals(lootItem.item.value())) {
+						toRemove.add(entryContainer);
+					}
+				}
+			}
+		}
+		for (LootPoolEntryContainer entryContainer : toRemove) {
+			entries.remove(entryContainer);
+		}
+		return this;
+	}
+
+	public MutableLootPool remove(TagKey<Item> tag) {
+		ArrayList<LootPoolEntryContainer> toRemove = new ArrayList<>();
+		for (LootPoolEntryContainer entryContainer : entries) {
+			if (entryContainer instanceof LootItem lootItem) {
+				if (lootItem.item.value().builtInRegistryHolder().is(tag)) {
+					toRemove.add(entryContainer);
+				}
+			}
+		}
+		for (LootPoolEntryContainer entryContainer : toRemove) {
+			entries.remove(entryContainer);
+		}
+		return this;
+	}
+
+	public MutableLootPool remove(Predicate<Item> predicate) {
+		ArrayList<LootPoolEntryContainer> toRemove = new ArrayList<>();
+		for (LootPoolEntryContainer entryContainer : entries) {
+			if (entryContainer instanceof LootItem lootItem) {
+				if (predicate.test(lootItem.item.value())) {
+					toRemove.add(entryContainer);
+				}
+			}
+		}
+		for (LootPoolEntryContainer entryContainer : toRemove) {
+			entries.remove(entryContainer);
+		}
+		return this;
+	}
+
+	public MutableLootPool replace(ItemLike original, ItemLike replacement) {
+		for (int i = 0; i < entries.size(); i++) {
+			LootPoolEntryContainer entryContainer = entries.get(i);
+			if (entryContainer instanceof LootItem lootItem) {
+				if (original.equals(lootItem.item.value())) {
+					MutableLootItem mutableLootItem = new MutableLootItem(lootItem);
+					mutableLootItem.setItem(replacement);
+					entries.set(i, mutableLootItem.build());
+				}
+			}
+		}
+		return this;
+	}
+
+	/**
+	 * Returns if the loot pool contains the given item
+	 *
+	 * @param item item to check for
+	 * @return if item was found
+	 */
+	public boolean hasItem(Item item) {
+		for (LootPoolEntryContainer entryContainer : entries) {
+			if (entryContainer instanceof LootItem lootItem) {
+				if (item.equals(lootItem.item.value())) return true;
+			}
+		}
+		return false;
+	}
+
+	/**
+	 * Returns if the loot pool contains an item that matches the predicate
+	 *
+	 * @param predicate condition to check for
+	 * @return if item was found
+	 */
+	public boolean hasItem(Predicate<Item> predicate) {
+		for (LootPoolEntryContainer entryContainer : entries) {
+			if (entryContainer instanceof LootItem lootItem) {
+				if (predicate.test(lootItem.item.value())) return true;
+			}
+		}
+		return false;
+	}
+
+	/**
+	 * Returns if the loot pool contains an item with the given tag
+	 *
+	 * @param itemTagKey item tag to check for
+	 * @return if item was found
+	 */
+	public boolean hasItem(TagKey<Item> itemTagKey) {
+		for (LootPoolEntryContainer entryContainer : entries) {
+			if (entryContainer instanceof LootItem lootItem) {
+				if (lootItem.item.value().builtInRegistryHolder().is(itemTagKey)) return true;
+			}
+		}
+		return false;
+	}
+
+	/**
+	 * Returns if the loot pool contains any of the given items
+	 *
+	 * @param items items to check for
+	 * @return if any of the items were found
+	 */
+	public boolean hasAnyItems(Item... items) {
+		for (LootPoolEntryContainer entryContainer : entries) {
+			if (entryContainer instanceof LootItem lootItem) {
+				for (Item item : items) {
+					if (item.equals(lootItem.item.value())) return true;
+				}
+			}
+		}
+		return false;
+	}
+
+	/**
+	 * Returns if the loot table contains all the given items
+	 *
+	 * @param items items to check for
+	 * @return if all of the items were found
+	 */
+	public boolean hasAllItems(Item @NotNull ... items) {
+		for (Item item : items) {
+			boolean found = false;
+			for (LootPoolEntryContainer entryContainer : entries) {
+				if (entryContainer instanceof LootItem lootItem) {
+					if (item.equals(lootItem.item.value())) {
+						found = true;
+						break;
+					}
+				}
+			}
+			if (!found) return false;
+		}
+		return true;
+	}
+}