Cancel stomps when moving upwards

src/main/java/com/minelittlepony/unicopia/ability/EarthPonyStompAbility.java

@@ -6,9 +6,11 @@ import org.jetbrains.annotations.Nullable;
 
 import com.minelittlepony.unicopia.AwaitTickQueue;
 import com.minelittlepony.unicopia.Race;
+import com.minelittlepony.unicopia.USounds;
 import com.minelittlepony.unicopia.UTags;
 import com.minelittlepony.unicopia.ability.data.Hit;
 import com.minelittlepony.unicopia.client.render.PlayerPoser.Animation;
+import com.minelittlepony.unicopia.entity.LandingEventHandler;
 import com.minelittlepony.unicopia.entity.Living;
 import com.minelittlepony.unicopia.entity.damage.UDamageTypes;
 import com.minelittlepony.unicopia.entity.player.Pony;
@@ -20,7 +22,6 @@ import com.minelittlepony.unicopia.server.world.BlockDestructionManager;
 import com.minelittlepony.unicopia.util.PosHelper;
 import com.minelittlepony.unicopia.util.VecHelper;
 
-import it.unimi.dsi.fastutil.floats.Float2FloatFunction;
 import net.minecraft.block.Block;
 import net.minecraft.block.BlockState;
 import net.minecraft.enchantment.EnchantmentHelper;
@@ -105,9 +106,17 @@ public class EarthPonyStompAbility implements Ability<Hit> {
 
     @Override
     public boolean apply(Pony iplayer, Hit data) {
-        PlayerEntity player = iplayer.asEntity();
+        final PlayerEntity player = iplayer.asEntity();
+        final double initialY = player.getY() + 5;
+
+        var r = new LandingEventHandler.Callback() {
+            @Override
+            public float dispatch(float fallDistance) {
+                // fail if landing above the starting position
+                if (player.getY() > initialY) {
+                    return fallDistance;
+                }
 
-        Float2FloatFunction r = fallDistance -> {
                 player.fallDistance = 0;
                 BlockPos center = PosHelper.findSolidGroundAt(player.getEntityWorld(), player.getBlockPos(), iplayer.getPhysics().getGravitySignum());
 
@@ -166,13 +175,20 @@ public class EarthPonyStompAbility implements Ability<Hit> {
                 iplayer.subtractEnergyCost(rad);
                 iplayer.asEntity().addExhaustion(3);
                 return 0F;
+            }
+
+            @Override
+            public void onCancelled() {
+                System.out.println("Cancel Stomp");
+                iplayer.playSound(USounds.GUI_ABILITY_FAIL, 1F);
+            }
         };
 
         if (iplayer.asEntity().isOnGround()) {
             iplayer.setAnimation(Animation.STOMP, Animation.Recipient.ANYONE, 10);
             iplayer.asEntity().jump();
             iplayer.updateVelocity();
-            AwaitTickQueue.scheduleTask(iplayer.asWorld(), w -> r.get(0F), 5);
+            AwaitTickQueue.scheduleTask(iplayer.asWorld(), w -> r.dispatch(0F), 5);
         } else {
             thrustDownwards(iplayer);
             iplayer.waitForFall(r);

src/main/java/com/minelittlepony/unicopia/entity/LandingEventHandler.java

@@ -0,0 +1,68 @@
+package com.minelittlepony.unicopia.entity;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+import org.jetbrains.annotations.Nullable;
+
+import com.minelittlepony.unicopia.util.Tickable;
+
+public class LandingEventHandler implements Tickable {
+
+    private final Living<?> living;
+
+    @Nullable
+    private final AtomicReference<Callback> callback = new AtomicReference<>();
+    private double prevY;
+    private float prevFallDistance;
+
+    public LandingEventHandler(Living<?> living) {
+        this.living = living;
+    }
+
+    public void setCallback(LandingEventHandler.Callback callback) {
+        if (living.asEntity().isOnGround()) {
+            callback.dispatch(0F);
+        } else {
+            updateCallback(callback);
+        }
+    }
+
+    public void beforeTick() {
+
+    }
+
+    @Override
+    public void tick() {
+        if (living.asEntity().getY() > prevY) {
+            discard();
+        }
+        prevY = living.asEntity().getY();
+
+        if (living.asEntity().isOnGround() && living.landedChanged()) {
+            fire(prevFallDistance);
+        }
+        prevFallDistance = living.asEntity().fallDistance;
+    }
+
+    float fire(float fallDistance) {
+        var event = callback.getAndSet(null);
+        return event == null ? fallDistance : event.dispatch(fallDistance);
+    }
+
+    void discard() {
+        updateCallback(null);
+    }
+
+    void updateCallback(@Nullable Callback callback) {
+        var event = this.callback.getAndSet(callback);
+        if (event != null) {
+            event.onCancelled();
+        }
+    }
+
+    public interface Callback {
+        float dispatch(float fallDistance);
+
+        void onCancelled();
+    }
+}

src/main/java/com/minelittlepony/unicopia/entity/Living.java

@@ -1,7 +1,6 @@
 package com.minelittlepony.unicopia.entity;
 
 import java.util.*;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
@@ -39,7 +38,6 @@ import com.minelittlepony.unicopia.projectile.ProjectileImpactListener;
 import com.minelittlepony.unicopia.server.world.DragonBreathStore;
 import com.minelittlepony.unicopia.util.*;
 
-import it.unimi.dsi.fastutil.floats.Float2FloatFunction;
 import it.unimi.dsi.fastutil.floats.Float2ObjectFunction;
 import net.fabricmc.fabric.api.util.TriState;
 import net.minecraft.block.BlockState;
@@ -84,10 +82,6 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
     private final Interactable landedHeuristic;
     private final Interactable jumpingHeuristic;
 
-    @Nullable
-    private final AtomicReference<Float2FloatFunction> landEvent = new AtomicReference<>();
-    private float prevFallDistance;
-
     private boolean invisible = false;
 
     @Nullable
@@ -101,6 +95,7 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
 
     private final List<Tickable> tickers = new ArrayList<>();
 
+    private final LandingEventHandler landEvent = addTicker(new LandingEventHandler(this));
     private final Enchantments enchants = addTicker(new Enchantments(this));
     private final ItemTracker armour = addTicker(new ItemTracker(this));
     //private final Transportation<T> transportation = new Transportation<>(this);
@@ -134,12 +129,8 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
         this.invisible = invisible;
     }
 
-    public void waitForFall(Float2FloatFunction action) {
-        if (entity.isOnGround()) {
-            action.get(0F);
-        } else {
-            landEvent.set(action);
-        }
+    public void waitForFall(LandingEventHandler.Callback callback) {
+        landEvent.setCallback(callback);
     }
 
     public boolean sneakingChanged() {
@@ -209,7 +200,7 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
 
     @Override
     public boolean beforeUpdate() {
-        prevFallDistance = entity.fallDistance;
+        landEvent.beforeTick();
         if (EffectUtils.getAmplifier(entity, UEffects.PARALYSIS) > 1 && entity.getVelocity().horizontalLengthSquared() > 0) {
             entity.setVelocity(entity.getVelocity().multiply(0, 1, 0));
             updateVelocity();
@@ -240,10 +231,6 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
             invinsibilityTicks--;
         }
 
-        if (entity.isOnGround() && landedChanged()) {
-            onLanded(prevFallDistance);
-        }
-
         if (entity.hasStatusEffect(UEffects.PARALYSIS) && entity.getVelocity().horizontalLengthSquared() > 0) {
             entity.setVelocity(entity.getVelocity().multiply(0, 1, 0));
             updateVelocity();
@@ -264,11 +251,6 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
         transportation.tick();
     }
 
-    private float onLanded(float fallDistance) {
-        var event = landEvent.getAndSet(null);
-        return event == null ? fallDistance : event.get(fallDistance);
-    }
-
     public void updateAttributeModifier(UUID id, EntityAttribute attribute, float desiredValue, Float2ObjectFunction<EntityAttributeModifier> modifierSupplier, boolean permanent) {
         @Nullable
         EntityAttributeInstance instance = asEntity().getAttributeInstance(attribute);
@@ -474,7 +456,7 @@ public abstract class Living<T extends LivingEntity> implements Equine<T>, Caste
     }
 
     public float onImpact(float distance, float damageMultiplier, DamageSource cause) {
-        float fallDistance = onLanded(getEffectiveFallDistance(distance));
+        float fallDistance = landEvent.fire(getEffectiveFallDistance(distance));
 
         getSpellSlot().get(SpellPredicate.IS_DISGUISE, false).ifPresent(spell -> {
             spell.getDisguise().onImpact(this, fallDistance, damageMultiplier, cause);