Call modifier removal hook on modifier traits

src/main/java/slimeknights/tconstruct/library/modifiers/hook/build/ModifierRemovalHook.java

@@ -2,6 +2,8 @@
 
 import net.minecraft.network.chat.Component;
 import slimeknights.tconstruct.library.modifiers.Modifier;
+import slimeknights.tconstruct.library.modifiers.ModifierEntry;
+import slimeknights.tconstruct.library.modifiers.ModifierHooks;
 import slimeknights.tconstruct.library.tools.nbt.IToolStackView;
 import slimeknights.tconstruct.library.utils.RestrictedCompoundTag;
 
@@ -41,4 +43,25 @@ public Component onRemoved(IToolStackView tool, Modifier modifier) {
       return null;
     }
   }
+
+  /**
+   * Calls the hook for all modifiers that were on the original but not on the updated tool.
+   * Handles any modifier update including standard removal, part swapping, and even modifier traits.
+   * @param original  Tool before changes were made
+   * @param updated   Tool after changes were made
+   * @return  Error message if a removed modifier errored, or null if no error
+   */
+  @Nullable
+  static Component onRemoved(IToolStackView original, IToolStackView updated) {
+    for (ModifierEntry entry : original.getModifierList()) {
+      ModifierRemovalHook hook = entry.getModifier().getHooks().getOrNull(ModifierHooks.REMOVE);
+      if (hook != null && updated.getModifierLevel(entry.getId()) == 0) {
+        Component error = hook.onRemoved(updated, entry.getModifier());
+        if (error != null) {
+          return error;
+        }
+      }
+    }
+    return null;
+  }
 }

src/main/java/slimeknights/tconstruct/shared/command/subcommand/ModifiersCommand.java

@@ -17,6 +17,7 @@
 import slimeknights.tconstruct.TConstruct;
 import slimeknights.tconstruct.library.modifiers.Modifier;
 import slimeknights.tconstruct.library.modifiers.ModifierHooks;
+import slimeknights.tconstruct.library.modifiers.hook.build.ModifierRemovalHook;
 import slimeknights.tconstruct.library.tools.nbt.ToolStack;
 import slimeknights.tconstruct.shared.command.HeldModifiableItemIterator;
 import slimeknights.tconstruct.shared.command.argument.ModifierArgument;
@@ -89,18 +90,18 @@ private static int remove(CommandContext<CommandSourceStack> context, int level)
     MutableInt maxRemove = new MutableInt(1);
     List<LivingEntity> successes = HeldModifiableItemIterator.apply(context, (living, stack) -> {
       // add modifier
-      ToolStack tool = ToolStack.from(stack);
+      ToolStack original = ToolStack.from(stack);
 
       // first, see if the modifier exists
-      int currentLevel = tool.getUpgrades().getLevel(modifier.getId());
+      int currentLevel = original.getUpgrades().getLevel(modifier.getId());
       if (currentLevel == 0) {
         throw CANNOT_REMOVE.create(modifier.getDisplayName(level), living.getName());
       }
       int removeLevel = level == -1 ? currentLevel : level;
       if (removeLevel > maxRemove.intValue()) {
         maxRemove.setValue(removeLevel);
       }
-      tool = tool.copy();
+      ToolStack tool = original.copy();
 
       // first remove hook, primarily for removing raw NBT which is highly discouraged using
       int newLevel = currentLevel - removeLevel;
@@ -117,13 +118,12 @@ private static int remove(CommandContext<CommandSourceStack> context, int level)
         throw MODIFIER_ERROR.create(validated);
       }
 
-      // if this was the last level, validate the tool is still valid without it
-      if (newLevel <= 0) {
-        validated = modifier.getHook(ModifierHooks.REMOVE).onRemoved(tool, modifier);
-        if (validated != null) {
-          throw MODIFIER_ERROR.create(validated);
-        }
+      // ask modifiers if it's okay to remove them
+      validated = ModifierRemovalHook.onRemoved(original, tool);
+      if (validated != null) {
+        throw MODIFIER_ERROR.create(validated);
       }
+
       // if successful, update held item
       living.setItemInHand(InteractionHand.MAIN_HAND, tool.createStack(stack.getCount()));
       return true;

src/main/java/slimeknights/tconstruct/tables/recipe/TinkerStationPartSwapping.java

@@ -10,14 +10,12 @@
 import net.minecraft.world.level.Level;
 import slimeknights.tconstruct.TConstruct;
 import slimeknights.tconstruct.common.TinkerTags;
-import slimeknights.tconstruct.library.materials.IMaterialRegistry;
-import slimeknights.tconstruct.library.materials.MaterialRegistry;
 import slimeknights.tconstruct.library.materials.definition.IMaterial;
 import slimeknights.tconstruct.library.materials.definition.MaterialVariant;
 import slimeknights.tconstruct.library.materials.definition.MaterialVariantId;
-import slimeknights.tconstruct.library.modifiers.Modifier;
 import slimeknights.tconstruct.library.modifiers.ModifierEntry;
 import slimeknights.tconstruct.library.modifiers.ModifierHooks;
+import slimeknights.tconstruct.library.modifiers.hook.build.ModifierRemovalHook;
 import slimeknights.tconstruct.library.recipe.RecipeResult;
 import slimeknights.tconstruct.library.recipe.casting.material.MaterialCastingLookup;
 import slimeknights.tconstruct.library.recipe.material.MaterialRecipe;
@@ -31,12 +29,7 @@
 import slimeknights.tconstruct.library.tools.part.IToolPart;
 import slimeknights.tconstruct.tables.TinkerTables;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
-import java.util.Map.Entry;
 import java.util.stream.IntStream;
 
 /**
@@ -93,8 +86,8 @@ public ItemStack getResultItem() {
   public RecipeResult<ItemStack> getValidatedResult(ITinkerStationContainer inv) {
     // copy the tool NBT to ensure the original tool is intact
     ItemStack tinkerable = inv.getTinkerableStack();
-    ToolStack tool = ToolStack.from(tinkerable);
-    List<IToolPart> parts = ToolPartsHook.parts(tool.getDefinition());
+    ToolStack original = ToolStack.from(tinkerable);
+    List<IToolPart> parts = ToolPartsHook.parts(original.getDefinition());
 
     // prevent part swapping on large tools in small tables
     if (parts.size() > inv.getInputCount()) {
@@ -130,49 +123,18 @@ public RecipeResult<ItemStack> getValidatedResult(ITinkerStationContainer inv) {
         }
 
         // ensure there is a change in the part or we are repairing the tool, note we compare variants so you could swap oak head for birch head
-        MaterialVariant toolVariant = tool.getMaterial(index);
+        MaterialVariant toolVariant = original.getMaterial(index);
         boolean didChange = !toolVariant.sameVariant(partVariant);
         float repairDurability = MaterialRepairModule.getDurability(null, partVariant.getId(), part.getStatType());
-        if (!didChange && (tool.getDamage() == 0 || repairDurability == 0)) {
+        if (!didChange && (original.getDamage() == 0 || repairDurability == 0)) {
           return RecipeResult.pass();
         }
 
         // actual update
-        tool = tool.copy();
+        ToolStack tool = original.copy();
 
         // determine which modifiers are going to be removed
-        List<Modifier> actuallyRemoved = Collections.emptyList();
         if (didChange) {
-          Map<Modifier,Integer> removedTraits = new HashMap<>();
-          // start with a map of all modifiers on the old part
-          // TODO: this logic looks correct, but I feel like it might be more complicated than needed
-          // basically, if the new part has the modifier, its not going to be removed no matter how the levels differ, a set should suffice
-          IMaterialRegistry materialRegistry = MaterialRegistry.getInstance();
-          for (ModifierEntry entry : materialRegistry.getTraits(toolVariant.getId(), part.getStatType())) {
-            removedTraits.put(entry.getModifier(), entry.getLevel());
-          }
-          // subtract any modifiers on the new part
-          for (ModifierEntry entry : materialRegistry.getTraits(partVariant.getId(), part.getStatType())) {
-            Modifier modifier = entry.getModifier();
-            if (removedTraits.containsKey(modifier)) {
-              int value = removedTraits.get(modifier) - entry.getLevel();
-              if (value <= 0) {
-                removedTraits.remove(modifier);
-              } else {
-                removedTraits.put(modifier, value);
-              }
-            }
-          }
-          // for the remainder, fill a list as we have another hooks to call with them
-          actuallyRemoved = new ArrayList<>();
-          for (Entry<Modifier,Integer> entry : removedTraits.entrySet()) {
-            Modifier modifier = entry.getKey();
-            if (tool.getModifierLevel(modifier) <= entry.getValue()) {
-              modifier.getHook(ModifierHooks.RAW_DATA).removeRawData(tool, modifier, tool.getRestrictedNBT());
-              actuallyRemoved.add(modifier);
-            }
-          }
-
           // do the actual part replacement
           tool.replaceMaterial(index, partVariant);
         }
@@ -205,9 +167,8 @@ public RecipeResult<ItemStack> getValidatedResult(ITinkerStationContainer inv) {
         if (error != null) {
           return RecipeResult.failure(error);
         }
-        // finally, validate removed modifiers
-        for (Modifier modifier : actuallyRemoved) {
-          error = modifier.getHook(ModifierHooks.REMOVE).onRemoved(tool, modifier);
+        if (didChange) {
+          error = ModifierRemovalHook.onRemoved(original, tool);
           if (error != null) {
             return RecipeResult.failure(error);
           }

src/main/java/slimeknights/tconstruct/tools/recipe/ModifierRemovalRecipe.java

@@ -16,8 +16,9 @@
 import slimeknights.tconstruct.library.json.predicate.modifier.ModifierPredicate;
 import slimeknights.tconstruct.library.modifiers.Modifier;
 import slimeknights.tconstruct.library.modifiers.ModifierEntry;
-import slimeknights.tconstruct.library.modifiers.ModifierId;
 import slimeknights.tconstruct.library.modifiers.ModifierHooks;
+import slimeknights.tconstruct.library.modifiers.ModifierId;
+import slimeknights.tconstruct.library.modifiers.hook.build.ModifierRemovalHook;
 import slimeknights.tconstruct.library.recipe.ITinkerableContainer;
 import slimeknights.tconstruct.library.recipe.RecipeResult;
 import slimeknights.tconstruct.library.recipe.modifiers.ModifierRecipeLookup;
@@ -103,10 +104,10 @@ public Component getDescription(@Nullable ITinkerableContainer inv) {
 
   @Override
   public RecipeResult<ToolStack> getResult(ITinkerableContainer inv, ModifierEntry entry) {
-    ToolStack tool = inv.getTinkerable();
+    ToolStack original = inv.getTinkerable();
 
     // salvage
-    tool = tool.copy();
+    ToolStack tool = original.copy();
     ModifierId modifierId = entry.getId();
     ModifierSalvage salvage = ModifierRecipeLookup.getSalvage(inv.getTinkerableStack(), tool, modifierId, entry.getLevel());
 
@@ -130,12 +131,9 @@ public RecipeResult<ToolStack> getResult(ITinkerableContainer inv, ModifierEntry
     if (error != null) {
       return RecipeResult.failure(error);
     }
-    // if this was the last level, validate the tool is still valid without it
-    if (newLevel <= 0) {
-      error = modifier.getHook(ModifierHooks.REMOVE).onRemoved(tool, modifier);
-      if (error != null) {
-        return RecipeResult.failure(error);
-      }
+    error = ModifierRemovalHook.onRemoved(original, tool);
+    if (error != null) {
+      return RecipeResult.failure(error);
     }
     // successfully removed
     return RecipeResult.success(tool);