Fix using item in off-hand slot not being swapped

src/main/java/net/shadowfacts/autoswap/mixin/MixinServerPlayerInteractionManager.java

@@ -1,11 +1,15 @@
 package net.shadowfacts.autoswap.mixin;
 
+import net.minecraft.entity.player.PlayerEntity;
 import net.minecraft.item.ItemStack;
 import net.minecraft.item.ItemUsageContext;
 import net.minecraft.server.network.ServerPlayerEntity;
 import net.minecraft.server.network.ServerPlayerInteractionManager;
 import net.minecraft.util.ActionResult;
+import net.minecraft.util.Hand;
+import net.minecraft.util.hit.BlockHitResult;
 import net.minecraft.util.math.BlockPos;
+import net.minecraft.world.World;
 import net.shadowfacts.autoswap.AutoSwap;
 import org.spongepowered.asm.mixin.Mixin;
 import org.spongepowered.asm.mixin.Shadow;
@@ -14,6 +18,8 @@ import org.spongepowered.asm.mixin.injection.Inject;
 import org.spongepowered.asm.mixin.injection.Redirect;
 import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;
 
+import java.awt.dnd.Autoscroll;
+
 /**
  * @author shadowfacts
  */
@@ -27,6 +33,7 @@ public abstract class MixinServerPlayerInteractionManager {
 	abstract boolean isCreative();
 
 	private ItemStack heldStackAtBeginningOfBreak = ItemStack.EMPTY;
+	private ItemStack heldStackAtBeginningOfUse = ItemStack.EMPTY;
 
 	@Inject(method = "tryBreakBlock", at = @At("HEAD"))
 	private void beginTryBreakBlock(BlockPos pos, CallbackInfoReturnable<Void> cb) {
@@ -39,17 +46,27 @@ public abstract class MixinServerPlayerInteractionManager {
 		heldStackAtBeginningOfBreak = ItemStack.EMPTY;
 	}
 
-	@Redirect(
+	@Inject(
 			method = "interactBlock",
 			at = @At(value = "INVOKE", target = "Lnet/minecraft/item/ItemStack;useOnBlock(Lnet/minecraft/item/ItemUsageContext;)Lnet/minecraft/util/ActionResult;")
 	)
-	private ActionResult useOnBlock(ItemStack self, ItemUsageContext context) {
-		ItemStack selfCopy = self.copy();
-		ActionResult result = self.useOnBlock(context);
-		if (!isCreative() && result == ActionResult.SUCCESS) {
-			AutoSwap.INSTANCE.onUseOnItem(player, player.inventory, selfCopy);
+	private void beforeUseOnBlock(PlayerEntity player, World world, ItemStack heldStack, Hand hand, BlockHitResult result, CallbackInfoReturnable<ActionResult> cb) {
+		heldStackAtBeginningOfUse = heldStack.copy();
+	}
+
+	@Inject(
+			method = "interactBlock",
+			at = @At(
+					value = "INVOKE",
+					target = "Lnet/minecraft/item/ItemStack;useOnBlock(Lnet/minecraft/item/ItemUsageContext;)Lnet/minecraft/util/ActionResult;",
+					shift = At.Shift.AFTER
+			)
+	)
+	private void afterUseOnBlock(PlayerEntity player, World world, ItemStack heldStack, Hand hand, BlockHitResult result, CallbackInfoReturnable<ActionResult> cb) {
+		if (!isCreative() && cb.getReturnValue() == ActionResult.SUCCESS) {
+			AutoSwap.INSTANCE.afterUseOnBlock(this.player, player.inventory, heldStackAtBeginningOfUse, hand);
 		}
-		return result;
+		heldStackAtBeginningOfUse = ItemStack.EMPTY;
 	}
 
 }

src/main/kotlin/net/shadowfacts/autoswap/AutoSwap.kt

@@ -8,6 +8,7 @@ import net.minecraft.entity.player.PlayerEntity
 import net.minecraft.entity.player.PlayerInventory
 import net.minecraft.item.ItemStack
 import net.minecraft.server.network.ServerPlayerEntity
+import net.minecraft.util.Hand
 import net.minecraft.util.math.BlockPos
 
 /**
@@ -30,16 +31,20 @@ object AutoSwap: ModInitializer {
 		}
 	}
 
-	fun onUseOnItem(player: ServerPlayerEntity, playerInv: PlayerInventory, heldStackBeforeUse: ItemStack) {
-		val heldStack = playerInv.mainHandStack
+	fun afterUseOnBlock(player: ServerPlayerEntity, playerInv: PlayerInventory, heldStackBeforeUse: ItemStack, hand: Hand) {
+		val heldStack = player.getStackInHand(hand)
 		if (!player.isCreative && heldStack.isEmpty) {
 			val index = playerInv.main.indexOfFirst {
 				ItemStack.areItemsEqualIgnoreDamage(heldStackBeforeUse, it)
 			}
 			if (index >= 0) {
-				playerInv.swapSlotWithHotbar(index)
-				val newStack = playerInv.mainHandStack
-				player.networkHandler.sendPacket(EntityEquipmentUpdateS2CPacket(player.entityId, EquipmentSlot.MAINHAND, newStack))
+				val newStack = playerInv.removeInvStack(index)
+				player.setStackInHand(hand, newStack)
+				val slot = when (hand) {
+					Hand.MAIN_HAND -> EquipmentSlot.MAINHAND
+					Hand.OFF_HAND -> EquipmentSlot.OFFHAND
+				}
+				player.networkHandler.sendPacket(EntityEquipmentUpdateS2CPacket(player.entityId, slot, newStack))
 			}
 		}
 	}