Reworked damage-modifying upgrades

src/main/java/WayofTime/bloodmagic/api/livingArmour/LivingArmourUpgrade.java

@@ -2,9 +2,11 @@ package WayofTime.bloodmagic.api.livingArmour;
 
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.Multimap;
+
 import net.minecraft.entity.EntityLivingBase;
 import net.minecraft.entity.ai.attributes.AttributeModifier;
 import net.minecraft.entity.player.EntityPlayer;
+import net.minecraft.item.ItemStack;
 import net.minecraft.nbt.NBTTagCompound;
 import net.minecraft.util.DamageSource;
 import net.minecraft.world.World;
@@ -33,6 +35,11 @@ public abstract class LivingArmourUpgrade
         this.level = Math.min(level, getMaxTier() - 1);
     }
 
+    public double getAdditionalDamageOnHit(double damage, EntityPlayer wearer, EntityLivingBase hitEntity, ItemStack weapon)
+    {
+        return 0;
+    }
+
     /**
      * Percentage of damage blocked. This stacks multiplicities with other
      * upgrades.

src/main/java/WayofTime/bloodmagic/livingArmour/LivingArmour.java

@@ -1,5 +1,19 @@
 package WayofTime.bloodmagic.livingArmour;
 
+import java.lang.reflect.Constructor;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map.Entry;
+
+import net.minecraft.entity.EntityLivingBase;
+import net.minecraft.entity.ai.attributes.AttributeModifier;
+import net.minecraft.entity.player.EntityPlayer;
+import net.minecraft.inventory.EntityEquipmentSlot;
+import net.minecraft.item.ItemStack;
+import net.minecraft.nbt.NBTTagCompound;
+import net.minecraft.nbt.NBTTagList;
+import net.minecraft.world.World;
 import WayofTime.bloodmagic.api.iface.IUpgradeTrainer;
 import WayofTime.bloodmagic.api.livingArmour.ILivingArmour;
 import WayofTime.bloodmagic.api.livingArmour.LivingArmourHandler;
@@ -12,20 +26,6 @@ import WayofTime.bloodmagic.util.helper.TextHelper;
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.Multimap;
 
-import net.minecraft.entity.ai.attributes.AttributeModifier;
-import net.minecraft.entity.player.EntityPlayer;
-import net.minecraft.inventory.EntityEquipmentSlot;
-import net.minecraft.item.ItemStack;
-import net.minecraft.nbt.NBTTagCompound;
-import net.minecraft.nbt.NBTTagList;
-import net.minecraft.world.World;
-
-import java.lang.reflect.Constructor;
-import java.util.ArrayList;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map.Entry;
-
 public class LivingArmour implements ILivingArmour
 {
     public static String chatBase = "chat.BloodMagic.livingArmour.";
@@ -46,6 +46,17 @@ public class LivingArmour implements ILivingArmour
         }
     }
 
+    public double getAdditionalDamageOnHit(double damage, EntityPlayer wearer, EntityLivingBase hitEntity, ItemStack weapon)
+    {
+        double total = 0;
+        for (Entry<String, LivingArmourUpgrade> entry : upgradeMap.entrySet())
+        {
+            total += entry.getValue().getAdditionalDamageOnHit(damage, wearer, hitEntity, weapon);
+        }
+
+        return total;
+    }
+
     @Override
     public Multimap<String, AttributeModifier> getAttributeModifiers()
     {

src/main/java/WayofTime/bloodmagic/livingArmour/upgrade/LivingArmourUpgradeGraveDigger.java

@@ -1,5 +1,9 @@
 package WayofTime.bloodmagic.livingArmour.upgrade;
 
+import net.minecraft.entity.EntityLivingBase;
+import net.minecraft.entity.player.EntityPlayer;
+import net.minecraft.item.ItemSpade;
+import net.minecraft.item.ItemStack;
 import net.minecraft.nbt.NBTTagCompound;
 import WayofTime.bloodmagic.api.Constants;
 import WayofTime.bloodmagic.api.livingArmour.LivingArmourUpgrade;
@@ -14,6 +18,17 @@ public class LivingArmourUpgradeGraveDigger extends LivingArmourUpgrade
         super(level);
     }
 
+    @Override
+    public double getAdditionalDamageOnHit(double damage, EntityPlayer wearer, EntityLivingBase hitEntity, ItemStack weapon)
+    {
+        if (weapon != null && weapon.getItem() instanceof ItemSpade)
+        {
+            return getDamageModifier();
+        }
+
+        return 0;
+    }
+
     public double getDamageModifier()
     {
         return damageBoost[this.level];

src/main/java/WayofTime/bloodmagic/livingArmour/upgrade/LivingArmourUpgradePoisonResist.java

@@ -1,16 +1,15 @@
 package WayofTime.bloodmagic.livingArmour.upgrade;
 
+import net.minecraft.entity.player.EntityPlayer;
+import net.minecraft.init.MobEffects;
+import net.minecraft.nbt.NBTTagCompound;
+import net.minecraft.potion.PotionEffect;
+import net.minecraft.world.World;
 import WayofTime.bloodmagic.api.Constants;
 import WayofTime.bloodmagic.api.livingArmour.ILivingArmour;
 import WayofTime.bloodmagic.api.livingArmour.LivingArmourUpgrade;
 import WayofTime.bloodmagic.util.ChatUtil;
 import WayofTime.bloodmagic.util.helper.TextHelper;
-import net.minecraft.entity.player.EntityPlayer;
-import net.minecraft.init.MobEffects;
-import net.minecraft.nbt.NBTTagCompound;
-import net.minecraft.potion.Potion;
-import net.minecraft.potion.PotionEffect;
-import net.minecraft.world.World;
 
 public class LivingArmourUpgradePoisonResist extends LivingArmourUpgrade
 {

src/main/java/WayofTime/bloodmagic/util/handler/EventHandler.java

@@ -102,7 +102,6 @@ import WayofTime.bloodmagic.livingArmour.tracker.StatTrackerSolarPowered;
 import WayofTime.bloodmagic.livingArmour.upgrade.LivingArmourUpgradeArrowShot;
 import WayofTime.bloodmagic.livingArmour.upgrade.LivingArmourUpgradeDigging;
 import WayofTime.bloodmagic.livingArmour.upgrade.LivingArmourUpgradeExperience;
-import WayofTime.bloodmagic.livingArmour.upgrade.LivingArmourUpgradeGraveDigger;
 import WayofTime.bloodmagic.livingArmour.upgrade.LivingArmourUpgradeGrimReaperSprint;
 import WayofTime.bloodmagic.livingArmour.upgrade.LivingArmourUpgradeJump;
 import WayofTime.bloodmagic.livingArmour.upgrade.LivingArmourUpgradeSelfSacrifice;
@@ -665,15 +664,7 @@ public class EventHandler
                 {
                     ItemStack mainWeapon = player.getItemStackFromSlot(EntityEquipmentSlot.MAINHAND);
 
-                    if (mainWeapon != null && mainWeapon.getItem() instanceof ItemSpade)
+                    event.setAmount((float) (event.getAmount() + lastPlayerSwingStrength * armour.getAdditionalDamageOnHit(event.getAmount(), player, attackedEntity, mainWeapon)));
-                    {
-                        LivingArmourUpgrade upgrade = ItemLivingArmour.getUpgrade(Constants.Mod.MODID + ".upgrade.graveDigger", chestStack);
-
-                        if (upgrade instanceof LivingArmourUpgradeGraveDigger)
-                        {
-                            event.setAmount((float) (event.getAmount() + lastPlayerSwingStrength * ((LivingArmourUpgradeGraveDigger) upgrade).getDamageModifier()));
-                        }
-                    }
 
                     float amount = Math.min(Utils.getModifiedDamage(attackedEntity, event.getSource(), event.getAmount()), attackedEntity.getHealth());