Improve default enchantment API (#1698)

* Make sure to not modify the tool property's enchantments * Add capability for "default enchantment level increases" * Further modification protection and EnchantmentLevel * Apply ALSON's suggestions * Expose to CraftTweaker
GregTechCEu · Apr 30, 2023 · 4baec8b · 4baec8b
1 parent 03f3257
commit 4baec8b
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 20 deletions.
diff --git a/src/main/java/gregtech/api/items/toolitem/EnchantmentLevel.java b/src/main/java/gregtech/api/items/toolitem/EnchantmentLevel.java
@@ -0,0 +1,15 @@
+package gregtech.api.items.toolitem;
+
+public class EnchantmentLevel {
+    private double level;
+    private double levelGrowth;
+
+    public EnchantmentLevel(double level, double levelGrowth) {
+        this.level = level;
+        this.levelGrowth = levelGrowth;
+    }
+
+    public int getLevel(int harvestTier) {
+        return (int) Math.min(Byte.MAX_VALUE, Math.floor(this.level + this.levelGrowth * harvestTier));
+    }
+}
diff --git a/src/main/java/gregtech/api/items/toolitem/IGTTool.java b/src/main/java/gregtech/api/items/toolitem/IGTTool.java
@@ -36,6 +36,7 @@
 import gregtech.client.utils.TooltipHelper;
 import gregtech.common.ConfigHolder;
 import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
 import net.minecraft.block.state.IBlockState;
 import net.minecraft.client.renderer.block.model.ModelResourceLocation;
 import net.minecraft.client.resources.I18n;
@@ -173,8 +174,9 @@ default ItemStack get(Material material) {
         }
 
         // Set tool and material enchantments
-        Object2IntMap<Enchantment> enchantments = toolProperty.getEnchantments();
-        enchantments.putAll(toolStats.getDefaultEnchantments(stack));
+        Object2IntMap<Enchantment> enchantments = new Object2IntOpenHashMap<>();
+        toolProperty.getEnchantments().forEach((enchantment, level) -> enchantments.put(enchantment, level.getLevel(toolProperty.getToolHarvestLevel())));
+        toolStats.getDefaultEnchantments(stack).forEach((enchantment, level) -> enchantments.put(enchantment, level.getLevel(toolProperty.getToolHarvestLevel())));
         enchantments.forEach((enchantment, level) -> {
             if (stack.getItem().canApplyAtEnchantingTable(stack, enchantment)) {
                 stack.addEnchantment(enchantment, level);

diff --git a/src/main/java/gregtech/api/items/toolitem/IGTToolDefinition.java b/src/main/java/gregtech/api/items/toolitem/IGTToolDefinition.java
@@ -2,7 +2,7 @@
 
 import gregtech.api.items.toolitem.aoe.AoESymmetrical;
 import gregtech.api.items.toolitem.behavior.IToolBehavior;
-import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2ObjectMap;
 import net.minecraft.block.state.IBlockState;
 import net.minecraft.enchantment.Enchantment;
 import net.minecraft.item.ItemStack;
@@ -93,7 +93,7 @@ default boolean isEnchantable(ItemStack stack) {
 
     boolean canApplyEnchantment(ItemStack stack, Enchantment enchantment);
 
-    Object2IntMap<Enchantment> getDefaultEnchantments(ItemStack stack);
+    Object2ObjectMap<Enchantment, EnchantmentLevel> getDefaultEnchantments(ItemStack stack);
 
     /**
      * Misc

diff --git a/src/main/java/gregtech/api/items/toolitem/ToolDefinitionBuilder.java b/src/main/java/gregtech/api/items/toolitem/ToolDefinitionBuilder.java
@@ -3,9 +3,7 @@
 import com.google.common.collect.ImmutableList;
 import gregtech.api.items.toolitem.aoe.AoESymmetrical;
 import gregtech.api.items.toolitem.behavior.IToolBehavior;
-import it.unimi.dsi.fastutil.objects.Object2IntArrayMap;
-import it.unimi.dsi.fastutil.objects.Object2IntMap;
-import it.unimi.dsi.fastutil.objects.ObjectOpenHashSet;
+import it.unimi.dsi.fastutil.objects.*;
 import net.minecraft.block.Block;
 import net.minecraft.block.material.Material;
 import net.minecraft.block.state.IBlockState;
@@ -44,7 +42,7 @@ public class ToolDefinitionBuilder {
     private final Set<Block> effectiveBlocks = new ObjectOpenHashSet<>();
     private final Set<Material> effectiveMaterials = new ObjectOpenHashSet<>();
     private Predicate<IBlockState> effectiveStates;
-    private Object2IntMap<Enchantment> defaultEnchantments = new Object2IntArrayMap<>();
+    private Object2ObjectMap<Enchantment, EnchantmentLevel> defaultEnchantments = new Object2ObjectArrayMap<>();
 
     public ToolDefinitionBuilder behaviors(IToolBehavior... behaviours) {
         Collections.addAll(this.behaviours, behaviours);
@@ -184,10 +182,15 @@ public ToolDefinitionBuilder effectiveStates(Predicate<IBlockState> effectiveSta
     }
 
     public ToolDefinitionBuilder defaultEnchantment(Enchantment enchantment, int level) {
-        this.defaultEnchantments.put(enchantment, level);
+        return this.defaultEnchantment(enchantment, level, 0);
+    }
+
+    public ToolDefinitionBuilder defaultEnchantment(Enchantment enchantment, double level, double growth) {
+        this.defaultEnchantments.put(enchantment, new EnchantmentLevel(level, growth));
         return this;
     }
 
+
     public IGTToolDefinition build() {
         return new IGTToolDefinition() {
 
@@ -210,8 +213,7 @@ public IGTToolDefinition build() {
             private final Supplier<ItemStack> brokenStack = ToolDefinitionBuilder.this.brokenStack;
             private final AoESymmetrical aoeSymmetrical = ToolDefinitionBuilder.this.aoeSymmetrical;
             private final Predicate<IBlockState> effectiveStatePredicate;
-            private final Object2IntMap<Enchantment> defaultEnchantments = ToolDefinitionBuilder.this.defaultEnchantments;
-
+            private final Object2ObjectMap<Enchantment, EnchantmentLevel> defaultEnchantments = ToolDefinitionBuilder.this.defaultEnchantments;
 
             {
                 Set<Block> effectiveBlocks = ToolDefinitionBuilder.this.effectiveBlocks;
@@ -308,8 +310,8 @@ public boolean canApplyEnchantment(ItemStack stack, Enchantment enchantment) {
             }
 
             @Override
-            public Object2IntMap<Enchantment> getDefaultEnchantments(ItemStack stack) {
-                return this.defaultEnchantments;
+            public Object2ObjectMap<Enchantment, EnchantmentLevel> getDefaultEnchantments(ItemStack stack) {
+                return Object2ObjectMaps.unmodifiable(this.defaultEnchantments);
             }
 
             @Override

diff --git a/src/main/java/gregtech/api/items/toolitem/ToolHelper.java b/src/main/java/gregtech/api/items/toolitem/ToolHelper.java
@@ -186,7 +186,7 @@ public static ItemStack getAndSetToolData(IGTTool tool, Material material, int m
         if (toolProperty != null) {
             toolProperty.getEnchantments().forEach((enchantment, level) -> {
                 if (stack.getItem().canApplyAtEnchantingTable(stack, enchantment)) {
-                    stack.addEnchantment(enchantment, level);
+                    stack.addEnchantment(enchantment, level.getLevel(toolProperty.getToolHarvestLevel()));
                 }
             });
         }

diff --git a/src/main/java/gregtech/api/unification/crafttweaker/MaterialExpansion.java b/src/main/java/gregtech/api/unification/crafttweaker/MaterialExpansion.java
@@ -196,11 +196,17 @@ public static int toolEnchant(Material m) {
     @ZenMethod
     @net.minecraftforge.fml.common.Optional.Method(modid = GTValues.MODID_CT)
     public static void addToolEnchantment(Material m, IEnchantment enchantment) {
+        addScaledToolEnchantment(m, enchantment, 0);
+    }
+
+    @ZenMethod
+    @net.minecraftforge.fml.common.Optional.Method(modid = GTValues.MODID_CT)
+    public static void addScaledToolEnchantment(Material m, IEnchantment enchantment, double levelGrowth) {
         if (checkFrozen("add tool enchantment")) return;
         ToolProperty prop = m.getProperty(PropertyKey.TOOL);
         if (prop != null) {
             Enchantment enchantmentType = (Enchantment) enchantment.getDefinition().getInternal();
-            prop.addEnchantmentForTools(enchantmentType, enchantment.getLevel());
+            prop.addEnchantmentForTools(enchantmentType, enchantment.getLevel(), levelGrowth);
         } else logError(m, "change tool enchantments", "Tool");
     }
 

diff --git a/src/main/java/gregtech/api/unification/material/properties/ToolProperty.java b/src/main/java/gregtech/api/unification/material/properties/ToolProperty.java
@@ -1,7 +1,8 @@
 package gregtech.api.unification.material.properties;
 
-import it.unimi.dsi.fastutil.objects.Object2IntArrayMap;
-import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import gregtech.api.items.toolitem.EnchantmentLevel;
+import it.unimi.dsi.fastutil.objects.Object2ObjectArrayMap;
+import it.unimi.dsi.fastutil.objects.Object2ObjectMap;
 import net.minecraft.enchantment.Enchantment;
 
 public class ToolProperty implements IMaterialProperty<ToolProperty> {
@@ -73,7 +74,7 @@ public class ToolProperty implements IMaterialProperty<ToolProperty> {
     /**
      * Enchantment to be applied to tools made from this Material.
      */
-    private final Object2IntMap<Enchantment> enchantments = new Object2IntArrayMap<>();
+    private final Object2ObjectMap<Enchantment, EnchantmentLevel> enchantments = new Object2ObjectArrayMap<>();
 
     public ToolProperty(float harvestSpeed, float attackDamage, int durability, int harvestLevel) {
         this.harvestSpeed = harvestSpeed;
@@ -150,7 +151,7 @@ public void setUnbreakable(boolean isUnbreakable) {
         this.isUnbreakable = isUnbreakable;
     }
 
-    public Object2IntMap<Enchantment> getEnchantments() {
+    public Object2ObjectMap<Enchantment, EnchantmentLevel> getEnchantments() {
         return enchantments;
     }
 
@@ -176,7 +177,11 @@ public void verifyProperty(MaterialProperties properties) {
     }
 
     public void addEnchantmentForTools(Enchantment enchantment, int level) {
-        enchantments.put(enchantment, level);
+        this.addEnchantmentForTools(enchantment, level, 0);
+    }
+
+    public void addEnchantmentForTools(Enchantment enchantment, double level, double levelGrowth) {
+        enchantments.put(enchantment, new EnchantmentLevel(level, levelGrowth));
     }
 
     public static class Builder {