Rework stopRiding Overwrite

Author: aromaa

src/mixins/java/org/spongepowered/common/mixin/core/world/entity/EntityMixin.java

@@ -98,7 +98,6 @@
 import org.spongepowered.common.SpongeCommon;
 import org.spongepowered.common.accessor.server.level.ChunkMapAccessor;
 import org.spongepowered.common.accessor.server.level.ChunkMap_TrackedEntityAccessor;
-import org.spongepowered.common.accessor.world.entity.EntityAccessor;
 import org.spongepowered.common.accessor.world.entity.PortalProcessorAccessor;
 import org.spongepowered.common.bridge.commands.CommandSourceProviderBridge;
 import org.spongepowered.common.bridge.data.DataCompoundHolder;
@@ -182,6 +181,7 @@ public abstract class EntityMixin implements EntityBridge, PlatformEntityBridge,
     @Shadow public abstract boolean shadow$onGround();
     @Shadow @Nullable protected abstract String shadow$getEncodeId();
     @Shadow @javax.annotation.Nullable public PortalProcessor portalProcess;
+    @Shadow public abstract void shadow$stopRiding();
     // @formatter:on
 
     private boolean impl$isConstructing = true;
@@ -306,23 +306,15 @@ public abstract class EntityMixin implements EntityBridge, PlatformEntityBridge,
 
     @Override
     public boolean bridge$dismountRidingEntity(final DismountType type) {
-        if (!this.shadow$level().isClientSide && ShouldFire.RIDE_ENTITY_EVENT_DISMOUNT) {
-            try (final CauseStackManager.StackFrame frame = PhaseTracker.getInstance().pushCauseFrame()) {
-                frame.pushCause(this);
-                frame.addContext(EventContextKeys.DISMOUNT_TYPE, type);
-                if (SpongeCommon.post(SpongeEventFactory.
-                        createRideEntityEventDismount(frame.currentCause(), (org.spongepowered.api.entity.Entity) this.shadow$getVehicle()))) {
-                    return false;
-                }
-            }
-        }
+        final Entity vehicle = this.shadow$getVehicle();
+
+        try (final CauseStackManager.StackFrame frame = PhaseTracker.getInstance().pushCauseFrame()) {
+            frame.addContext(EventContextKeys.DISMOUNT_TYPE, type);
 
-        final Entity tempEntity = this.shadow$getVehicle();
-        if (tempEntity != null) {
-            this.vehicle = null;
-            ((EntityAccessor) tempEntity).invoker$removePassenger((Entity) (Object) this);
+            this.shadow$stopRiding();
         }
-        return true;
+
+        return vehicle != this.shadow$getVehicle();
     }
 
     @Override
@@ -616,18 +608,26 @@ public Entity changeDimension(final DimensionTransition transition) {
         }
     }
 
-    /**
-     * @author rexbut - December 16th, 2016
-     * @reason - adjusted to support {@link DismountTypes}
-     */
-    @Overwrite
-    public void stopRiding() {
+    @Inject(method = "removeVehicle",
+        at = @At(value = "FIELD", target = "Lnet/minecraft/world/entity/Entity;vehicle:Lnet/minecraft/world/entity/Entity;", opcode = Opcodes.PUTFIELD), cancellable = true)
+    private void impl$onRemoveVehicle(final CallbackInfo ci) {
         final Entity vehicle = this.shadow$getVehicle();
         if (vehicle != null) {
-            if (vehicle.isRemoved()) {
-                this.bridge$dismountRidingEntity(DismountTypes.DEATH.get());
-            } else {
-                this.bridge$dismountRidingEntity(DismountTypes.PLAYER.get());
+            if (this.shadow$level().isClientSide || !ShouldFire.RIDE_ENTITY_EVENT_DISMOUNT) {
+                return;
+            }
+
+            try (final CauseStackManager.StackFrame frame = PhaseTracker.getInstance().pushCauseFrame()) {
+                frame.pushCause(this);
+                if (!frame.currentContext().containsKey(EventContextKeys.DISMOUNT_TYPE)) {
+                    frame.addContext(EventContextKeys.DISMOUNT_TYPE, vehicle.isRemoved()
+                        ? DismountTypes.DEATH.get()
+                        : DismountTypes.PLAYER.get());
+                }
+                if (SpongeCommon.post(SpongeEventFactory.
+                    createRideEntityEventDismount(frame.currentCause(), (org.spongepowered.api.entity.Entity) this.shadow$getVehicle()))) {
+                    ci.cancel();
+                }
             }
         }
     }