Added code to check if the map contains the entity (#122)

TeamMetallurgy · May 26, 2019 · e0b7f5e · e0b7f5e
1 parent 1caf2fc
commit e0b7f5e
Showing 1 changed file with 1 addition and 6 deletions.
diff --git a/src/main/java/com/teammetallurgy/atum/items/tools/ItemHammer.java b/src/main/java/com/teammetallurgy/atum/items/tools/ItemHammer.java
@@ -58,15 +58,10 @@ protected void onStun(EntityLivingBase target) {
 
     @SubscribeEvent
     public void onLivingUpdate(LivingEvent.LivingUpdateEvent event) {
-        if (stun.isEmpty()) return;
+        if (stun.isEmpty() || !stun.containsKey(event.getEntityLiving())) return;
         ModifiableAttributeInstance attribute = (ModifiableAttributeInstance) event.getEntityLiving().getEntityAttribute(SharedMonsterAttributes.MOVEMENT_SPEED);
         if (attribute.hasModifier(STUN)) {
             EntityLivingBase entity = event.getEntityLiving();
-            if(!stun.containsKey(entity)){
-                attribute.removeModifier(STUN);
-                return;
-            }
-
             int stunTime = stun.get(entity);
             if (stunTime <= 1) {
                 stun.remove(entity);