Implement more functions

GoldenStack · Aug 29, 2024 · 5b31fa7 · 5b31fa7
1 parent 322ac4c
commit 5b31fa7
Show file tree

Hide file tree

Showing 3 changed files with 111 additions and 12 deletions.
diff --git a/src/main/java/net/goldenstack/loot/LootFunction.java b/src/main/java/net/goldenstack/loot/LootFunction.java
@@ -1,13 +1,13 @@
 package net.goldenstack.loot;
 
-
 import net.goldenstack.loot.util.*;
 import net.goldenstack.loot.util.nbt.NBTPath;
 import net.goldenstack.loot.util.nbt.NBTReference;
 import net.goldenstack.loot.util.nbt.NBTUtils;
 import net.goldenstack.loot.util.predicate.ItemPredicate;
 import net.kyori.adventure.nbt.*;
 import net.kyori.adventure.text.Component;
+import net.minestom.server.MinecraftServer;
 import net.minestom.server.ServerFlag;
 import net.minestom.server.component.DataComponent;
 import net.minestom.server.entity.Entity;
@@ -23,6 +23,7 @@
 import net.minestom.server.potion.PotionEffect;
 import net.minestom.server.potion.PotionType;
 import net.minestom.server.registry.DynamicRegistry;
+import net.minestom.server.registry.Registries;
 import net.minestom.server.tag.Tag;
 import net.minestom.server.utils.NamespaceID;
 import net.minestom.server.utils.nbt.BinaryTagSerializer;
@@ -68,7 +69,9 @@ public interface LootFunction {
                     Template.entry("set_enchantments", SetEnchantments.class, SetEnchantments.SERIALIZER),
                     Template.entry("enchant_with_levels", EnchantWithLevels.class, EnchantWithLevels.SERIALIZER),
                     Template.entry("set_book_cover", SetBookCover.class, SetBookCover.SERIALIZER),
-                    Template.entry("fill_player_head", FillPlayerHead.class, FillPlayerHead.SERIALIZER)
+                    Template.entry("fill_player_head", FillPlayerHead.class, FillPlayerHead.SERIALIZER),
+                    Template.entry("enchant_randomly", EnchantRandomly.class, EnchantRandomly.SERIALIZER),
+                    Template.entry("furnace_smelt", FurnaceSmelt.class, FurnaceSmelt.SERIALIZER)
             )
     );
 
@@ -709,14 +712,6 @@ record SetEnchantments(@NotNull List<LootPredicate> predicates, @NotNull Map<Dyn
         public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
             if (!LootPredicate.all(predicates, context)) return input;
 
-            if (input.material().equals(Material.BOOK)) {
-                input = input.builder()
-                        .material(Material.ENCHANTED_BOOK)
-                        .set(ItemComponent.STORED_ENCHANTMENTS, input.get(ItemComponent.ENCHANTMENTS, EnchantmentList.EMPTY))
-                        .remove(ItemComponent.ENCHANTMENTS)
-                        .build();
-            }
-
             return EnchantmentUtils.modifyItem(input, map -> {
                 this.enchantments.forEach((enchantment, number) -> {
                     int count = number.getInt(context);
@@ -734,7 +729,7 @@ record EnchantWithLevels(@NotNull List<LootPredicate> predicates, @NotNull LootN
         public static final @NotNull BinaryTagSerializer<EnchantWithLevels> SERIALIZER = Template.template(
                 "conditions", Serial.lazy(() -> LootPredicate.SERIALIZER).list().optional(List.of()), EnchantWithLevels::predicates,
                 "levels", LootNumber.SERIALIZER, EnchantWithLevels::levels,
-                "options", Template.todo("enchantwithlevels list"), EnchantWithLevels::options,
+                "options", EnchantmentUtils.TAG_LIST, EnchantWithLevels::options,
                 EnchantWithLevels::new
         );
 
@@ -799,4 +794,59 @@ record FillPlayerHead(@NotNull List<LootPredicate> predicates, @NotNull Relevant
             return input.with(ItemComponent.PROFILE, new HeadProfile(skin));
         }
     }
+
+    record EnchantRandomly(@NotNull List<LootPredicate> predicates, @Nullable List<DynamicRegistry.Key<Enchantment>> options, boolean onlyCompatible) implements LootFunction {
+
+        public static final @NotNull BinaryTagSerializer<EnchantRandomly> SERIALIZER = Template.template(
+                "conditions", Serial.lazy(() -> LootPredicate.SERIALIZER).list().optional(List.of()), EnchantRandomly::predicates,
+                "options", BinaryTagSerializer.registryKey(Registries::enchantment).list().optional(), EnchantRandomly::options,
+                "only_compatible", BinaryTagSerializer.BOOLEAN.optional(true), EnchantRandomly::onlyCompatible,
+                EnchantRandomly::new
+        );
+
+        @Override
+        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
+            var reg = MinecraftServer.getEnchantmentRegistry();
+
+            List<DynamicRegistry.Key<Enchantment>> values = new ArrayList<>();
+
+            if (options == null) {
+                reg.values().forEach(value -> values.add(reg.getKey(value)));
+            } else {
+                values.addAll(options);
+            }
+
+            if (onlyCompatible && !input.material().equals(Material.BOOK)) {
+                values.removeIf(ench -> !reg.get(ench).supportedItems().contains(input.material()));
+            }
+
+            if (values.isEmpty()) return input;
+
+            Random rng = context.require(LootContext.RANDOM);
+
+            DynamicRegistry.Key<Enchantment> chosen = values.get(rng.nextInt(values.size()));
+
+            int level = rng.nextInt(reg.get(chosen).maxLevel() + 1);
+
+            return EnchantmentUtils.modifyItem(input, map -> map.put(chosen, level));
+        }
+    }
+
+    record FurnaceSmelt(@NotNull List<LootPredicate> predicates) implements LootFunction {
+
+        public static final @NotNull BinaryTagSerializer<FurnaceSmelt> SERIALIZER = Template.template(
+                "conditions", Serial.lazy(() -> LootPredicate.SERIALIZER).list().optional(List.of()), FurnaceSmelt::predicates,
+                FurnaceSmelt::new
+        );
+
+        @Override
+        public @NotNull ItemStack apply(@NotNull ItemStack input, @NotNull LootContext context) {
+            if (!LootPredicate.all(predicates, context)) return input;
+
+            ItemStack smelted = context.require(LootContext.VANILLA_INTERFACE).smelt(input);
+
+            return smelted != null ? smelted.withAmount(input.amount()) : input;
+        }
+    }
+
 }
diff --git a/src/main/java/net/goldenstack/loot/util/EnchantmentUtils.java b/src/main/java/net/goldenstack/loot/util/EnchantmentUtils.java
@@ -1,26 +1,64 @@
 package net.goldenstack.loot.util;
 
+import net.kyori.adventure.nbt.BinaryTag;
+import net.kyori.adventure.nbt.ListBinaryTag;
+import net.kyori.adventure.nbt.StringBinaryTag;
+import net.minestom.server.MinecraftServer;
 import net.minestom.server.component.DataComponent;
 import net.minestom.server.entity.Entity;
 import net.minestom.server.entity.EquipmentSlot;
 import net.minestom.server.entity.LivingEntity;
+import net.minestom.server.gamedata.tags.Tag;
 import net.minestom.server.item.ItemComponent;
 import net.minestom.server.item.ItemStack;
 import net.minestom.server.item.Material;
 import net.minestom.server.item.component.EnchantmentList;
 import net.minestom.server.item.enchant.Enchantment;
 import net.minestom.server.registry.DynamicRegistry;
+import net.minestom.server.utils.nbt.BinaryTagSerializer;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.function.Consumer;
 
+@SuppressWarnings("UnstableApiUsage")
 public class EnchantmentUtils {
 
     private EnchantmentUtils() {}
 
+    public static final BinaryTagSerializer<List<DynamicRegistry.Key<Enchantment>>> KEYS_RAW = Serial.<Enchantment>key().list();
+
+    public static final @NotNull BinaryTagSerializer<List<DynamicRegistry.Key<Enchantment>>> TAG_LIST = new BinaryTagSerializer<>() {
+        @Override
+        public @NotNull BinaryTag write(@NotNull Context context, @NotNull List<DynamicRegistry.Key<Enchantment>> value) {
+            return KEYS_RAW.write(context, value);
+        }
+
+        @Override
+        public @NotNull List<DynamicRegistry.Key<Enchantment>> read(@NotNull Context context, @NotNull BinaryTag tag) {
+            return switch (tag) {
+                case StringBinaryTag string -> {
+                    if (string.value().startsWith("#")) {
+                        List<DynamicRegistry.Key<Enchantment>> values = new ArrayList<>();
+                        MinecraftServer.getTagManager()
+                                .getTag(Tag.BasicType.ENCHANTMENTS, string.value().substring(1))
+                                .getValues()
+                                .forEach(value -> values.add(DynamicRegistry.Key.of(value)));
+                        yield values;
+                    } else {
+                        yield List.of(DynamicRegistry.Key.of(string.value()));
+                    }
+                }
+                case ListBinaryTag list -> KEYS_RAW.read(context, list);
+                default -> throw new IllegalArgumentException();
+            };
+        }
+    };
+
     public static int level(@Nullable ItemStack item, @NotNull DynamicRegistry.Key<Enchantment> key) {
         if (item == null) return 0;
 
@@ -52,7 +90,16 @@ public static int level(@Nullable Entity entity, @NotNull DynamicRegistry.Key<En
         var map = new HashMap<>(component.enchantments());
         enchantments.accept(map);
 
-        return item.with(type, new EnchantmentList(map, component.showInTooltip()));
+        // Make the book enchanted!
+        if (!map.isEmpty() && item.material().equals(Material.BOOK)) {
+            return item.builder()
+                    .material(Material.ENCHANTED_BOOK)
+                    .set(ItemComponent.STORED_ENCHANTMENTS, new EnchantmentList(map, component.showInTooltip()))
+                    .build();
+        } else {
+            return item.with(type, new EnchantmentList(map, component.showInTooltip()));
+        }
+
     }
 
 }
diff --git a/src/main/java/net/goldenstack/loot/util/VanillaInterface.java b/src/main/java/net/goldenstack/loot/util/VanillaInterface.java
@@ -25,4 +25,6 @@ public interface VanillaInterface {
 
     @NotNull List<ItemStack> getDynamicDrops(@NotNull NamespaceID choiceID, @NotNull CompoundBinaryTag blockNBT);
 
+    @Nullable ItemStack smelt(@NotNull ItemStack input);
+
 }
Original file line number	Diff line number	Diff line change
Expand Up		@@ -25,4 +25,6 @@ public interface VanillaInterface {

		@NotNull List<ItemStack> getDynamicDrops(@NotNull NamespaceID choiceID, @NotNull CompoundBinaryTag blockNBT);

		@Nullable ItemStack smelt(@NotNull ItemStack input);

		}