Use custom version of InventoryStorage because fabric's one is shit

src/main/java/slimeknights/mantle/block/entity/InventoryBlockEntity.java

@@ -1,7 +1,7 @@
 package slimeknights.mantle.block.entity;
 
+import io.github.fabricators_of_create.porting_lib.transfer.item.SlottedStackStorage;
 import lombok.Getter;
-import net.fabricmc.fabric.api.transfer.v1.item.InventoryStorage;
 import net.fabricmc.fabric.api.transfer.v1.item.ItemVariant;
 import net.fabricmc.fabric.api.transfer.v1.storage.SlottedStorage;
 import net.fabricmc.fabric.api.transfer.v1.storage.Storage;
@@ -22,6 +22,8 @@
 import net.minecraft.world.level.block.entity.BlockEntityType;
 import net.minecraft.world.level.block.state.BlockState;
 import org.jetbrains.annotations.Nullable;
+import slimeknights.mantle.fabric.transfer.IInventoryStorage;
+import slimeknights.mantle.fabric.transfer.InventoryStorage;
 import slimeknights.mantle.util.ItemStackList;
 
 import javax.annotation.Nonnull;
@@ -37,7 +39,7 @@ public abstract class InventoryBlockEntity extends NameableBlockEntity implement
   private final boolean saveSizeToNBT;
   protected int stackSizeLimit;
   @Getter
-  protected SlottedStorage<ItemVariant> itemHandler;
+  protected SlottedStackStorage itemHandler;
 
   /**
    * @param name Localization String for the inventory title. Can be overridden through setCustomName

src/main/java/slimeknights/mantle/fabric/transfer/IInventoryStorage.java

@@ -0,0 +1,28 @@
+package slimeknights.mantle.fabric.transfer;
+
+import io.github.fabricators_of_create.porting_lib.transfer.item.SlottedStackStorage;
+import net.fabricmc.fabric.api.transfer.v1.item.ItemVariant;
+import net.fabricmc.fabric.api.transfer.v1.storage.base.SingleSlotStorage;
+import org.jetbrains.annotations.UnmodifiableView;
+
+import java.util.List;
+
+public interface IInventoryStorage extends SlottedStackStorage {
+  /**
+   * Retrieve an unmodifiable list of the wrappers for the slots in this inventory.
+   * Each wrapper corresponds to a single slot in the inventory.
+   */
+  @Override
+  @UnmodifiableView
+  List<SingleSlotStorage<ItemVariant>> getSlots();
+
+  @Override
+  default int getSlotCount() {
+    return getSlots().size();
+  }
+
+  @Override
+  default SingleSlotStorage<ItemVariant> getSlot(int slot) {
+    return getSlots().get(slot);
+  }
+}

src/main/java/slimeknights/mantle/fabric/transfer/InventorySlotWrapper.java

@@ -0,0 +1,155 @@
+package slimeknights.mantle.fabric.transfer;
+
+import net.fabricmc.fabric.api.transfer.v1.item.ItemVariant;
+import net.fabricmc.fabric.api.transfer.v1.item.base.SingleStackStorage;
+import net.fabricmc.fabric.api.transfer.v1.transaction.TransactionContext;
+import net.fabricmc.fabric.impl.transfer.DebugMessages;
+import net.fabricmc.fabric.impl.transfer.item.ItemVariantImpl;
+import net.fabricmc.fabric.impl.transfer.item.SpecialLogicInventory;
+import net.minecraft.core.BlockPos;
+import net.minecraft.world.item.ItemStack;
+import net.minecraft.world.item.Items;
+import net.minecraft.world.level.block.ChestBlock;
+import net.minecraft.world.level.block.entity.AbstractFurnaceBlockEntity;
+import net.minecraft.world.level.block.entity.BrewingStandBlockEntity;
+import net.minecraft.world.level.block.entity.ChestBlockEntity;
+import net.minecraft.world.level.block.entity.ShulkerBoxBlockEntity;
+import net.minecraft.world.level.block.state.properties.ChestType;
+
+/**
+ * A wrapper around a single slot of an inventory.
+ * We must ensure that only one instance of this class exists for every inventory slot,
+ * or the transaction logic will not work correctly.
+ * This is handled by the Map in InventoryStorageImpl.
+ */
+class InventorySlotWrapper extends SingleStackStorage {
+  /**
+   * The strong reference to the InventoryStorageImpl ensures that the weak value doesn't get GC'ed when individual slots are still being accessed.
+   */
+  private final InventoryStorage storage;
+  final int slot;
+  private final SpecialLogicInventory specialInv;
+  private ItemStack lastReleasedSnapshot = null;
+
+  InventorySlotWrapper(InventoryStorage storage, int slot) {
+    this.storage = storage;
+    this.slot = slot;
+    this.specialInv = storage.inventory instanceof SpecialLogicInventory specialInv ? specialInv : null;
+  }
+
+  @Override
+  protected ItemStack getStack() {
+    return storage.inventory.getItem(slot);
+  }
+
+  @Override
+  protected void setStack(ItemStack stack) {
+    if (specialInv == null) {
+      storage.inventory.setItem(slot, stack);
+    } else {
+      specialInv.fabric_setSuppress(true);
+
+      try {
+        storage.inventory.setItem(slot, stack);
+      } finally {
+        specialInv.fabric_setSuppress(false);
+      }
+    }
+  }
+
+  @Override
+  public long insert(ItemVariant insertedVariant, long maxAmount, TransactionContext transaction) {
+    if (!canInsert(slot, ((ItemVariantImpl) insertedVariant).getCachedStack())) {
+      return 0;
+    }
+
+    long ret = super.insert(insertedVariant, maxAmount, transaction);
+    if (specialInv != null && ret > 0) specialInv.fabric_onTransfer(slot, transaction);
+    return ret;
+  }
+
+  private boolean canInsert(int slot, ItemStack stack) {
+    if (storage.inventory instanceof ShulkerBoxBlockEntity shulker) {
+      // Shulkers override canInsert but not isValid.
+      return shulker.canPlaceItemThroughFace(slot, stack, null);
+    } else {
+      return storage.inventory.canPlaceItem(slot, stack);
+    }
+  }
+
+  @Override
+  public long extract(ItemVariant variant, long maxAmount, TransactionContext transaction) {
+    long ret = super.extract(variant, maxAmount, transaction);
+    if (specialInv != null && ret > 0) specialInv.fabric_onTransfer(slot, transaction);
+    return ret;
+  }
+
+  /**
+   * Special cases because vanilla checks the current stack in the following functions (which it shouldn't):
+   * <ul>
+   *     <li>{@link AbstractFurnaceBlockEntity#canPlaceItem(int, ItemStack)}.</li>
+   *     <li>{@link BrewingStandBlockEntity#canPlaceItem(int, ItemStack)}.</li>
+   * </ul>
+   */
+  @Override
+  public int getCapacity(ItemVariant variant) {
+    // Special case to limit buckets to 1 in furnace fuel inputs.
+    if (storage.inventory instanceof AbstractFurnaceBlockEntity && slot == 1 && variant.isOf(Items.BUCKET)) {
+      return 1;
+    }
+
+    // Special case to limit brewing stand "bottle inputs" to 1.
+    if (storage.inventory instanceof BrewingStandBlockEntity && slot < 3) {
+      return 1;
+    }
+
+    return Math.min(storage.inventory.getMaxStackSize(), variant.getItem().getMaxStackSize());
+  }
+
+  // We override updateSnapshots to also schedule a markDirty call for the backing inventory.
+  @Override
+  public void updateSnapshots(TransactionContext transaction) {
+    storage.markDirtyParticipant.updateSnapshots(transaction);
+    super.updateSnapshots(transaction);
+
+    // For chests: also schedule a markDirty call for the other half
+    if (storage.inventory instanceof ChestBlockEntity chest && chest.getBlockState().getValue(ChestBlock.TYPE) != ChestType.SINGLE) {
+      BlockPos otherChestPos = chest.getBlockPos().relative(ChestBlock.getConnectedDirection(chest.getBlockState()));
+
+      if (chest.getLevel().getBlockEntity(otherChestPos) instanceof ChestBlockEntity otherChest) {
+        ((InventoryStorage) InventoryStorage.of(otherChest, null)).markDirtyParticipant.updateSnapshots(transaction);
+      }
+    }
+  }
+
+  @Override
+  protected void releaseSnapshot(ItemStack snapshot) {
+    lastReleasedSnapshot = snapshot;
+  }
+
+  @Override
+  protected void onFinalCommit() {
+    // Try to apply the change to the original stack
+    ItemStack original = lastReleasedSnapshot;
+    ItemStack currentStack = getStack();
+
+    if (storage.inventory instanceof SpecialLogicInventory specialLogicInv) {
+      specialLogicInv.fabric_onFinalCommit(slot, original, currentStack);
+    }
+
+    if (!original.isEmpty() && original.getItem() == currentStack.getItem()) {
+      // None is empty and the items match: just update the amount and NBT, and reuse the original stack.
+      original.setCount(currentStack.getCount());
+      original.setTag(currentStack.hasTag() ? currentStack.getTag().copy() : null);
+      setStack(original);
+    } else {
+      // Otherwise assume everything was taken from original so empty it.
+      original.setCount(0);
+    }
+  }
+
+  @Override
+  public String toString() {
+    return "InventorySlotWrapper[%s#%d]".formatted(DebugMessages.forInventory(storage.inventory), slot);
+  }
+}

src/main/java/slimeknights/mantle/fabric/transfer/InventoryStorage.java

@@ -0,0 +1,135 @@
+package slimeknights.mantle.fabric.transfer;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+
+import com.google.common.collect.MapMaker;
+import io.github.fabricators_of_create.porting_lib.transfer.item.SlottedStackStorage;
+import net.minecraft.world.item.ItemStack;
+import org.jetbrains.annotations.Nullable;
+import net.fabricmc.fabric.api.transfer.v1.item.ItemVariant;
+import net.fabricmc.fabric.api.transfer.v1.storage.base.CombinedStorage;
+import net.fabricmc.fabric.api.transfer.v1.storage.base.SingleSlotStorage;
+import net.fabricmc.fabric.api.transfer.v1.transaction.base.SnapshotParticipant;
+import net.fabricmc.fabric.impl.transfer.DebugMessages;
+import net.minecraft.core.Direction;
+import net.minecraft.world.Container;
+import net.minecraft.world.WorldlyContainer;
+import net.minecraft.world.entity.player.Inventory;
+
+/**
+ * Implementation of {@link InventoryStorage}.
+ * Note on thread-safety: we assume that Inventory's are inherently single-threaded, and no attempt is made at synchronization.
+ * However, the access to implementations can happen on multiple threads concurrently, which is why we use a thread-safe wrapper map.
+ */
+public class InventoryStorage extends CombinedStorage<ItemVariant, SingleSlotStorage<ItemVariant>> implements IInventoryStorage {
+  /**
+   * Global wrapper concurrent map.
+   *
+   * <p>A note on GC: weak keys alone are not suitable as the InventoryStorage slots strongly reference the Inventory keys.
+   * Weak values are suitable, but we have to ensure that the InventoryStorageImpl remains strongly reachable as long as
+   * one of the slot wrappers refers to it, hence the {@code strongRef} field in {@link InventorySlotWrapper}.
+   */
+  // TODO: look into promoting the weak reference to a soft reference if building the wrappers becomes a performance bottleneck.
+  // TODO: should have identity semantics?
+  private static final Map<Container, InventoryStorage> WRAPPERS = new MapMaker().weakValues().makeMap();
+
+  public static IInventoryStorage of(Container inventory, @Nullable Direction direction) {
+    InventoryStorage storage = WRAPPERS.computeIfAbsent(inventory, inv -> {
+      if (inv instanceof Inventory playerInventory) {
+        return new PlayerInventoryStorage(playerInventory);
+      } else {
+        return new InventoryStorage(inv);
+      }
+    });
+    storage.resizeSlotList();
+    return storage.getSidedWrapper(direction);
+  }
+
+  final Container inventory;
+  /**
+   * This {@code backingList} is the real list of wrappers.
+   * The {@code parts} in the superclass is the public-facing unmodifiable sublist with exactly the right amount of slots.
+   */
+  final List<InventorySlotWrapper> backingList;
+  /**
+   * This participant ensures that markDirty is only called once for the entire inventory.
+   */
+  final MarkDirtyParticipant markDirtyParticipant = new MarkDirtyParticipant();
+
+  InventoryStorage(Container inventory) {
+    super(Collections.emptyList());
+    this.inventory = inventory;
+    this.backingList = new ArrayList<>();
+  }
+
+  @Override
+  public List<SingleSlotStorage<ItemVariant>> getSlots() {
+    return parts;
+  }
+
+  /**
+   * Resize slot list to match the current size of the inventory.
+   */
+  private void resizeSlotList() {
+    int inventorySize = inventory.getContainerSize();
+
+    // If the public-facing list must change...
+    if (inventorySize != parts.size()) {
+      // Ensure we have enough wrappers in the backing list.
+      while (backingList.size() < inventorySize) {
+        backingList.add(new InventorySlotWrapper(this, backingList.size()));
+      }
+
+      // Update the public-facing list.
+      parts = Collections.unmodifiableList(backingList.subList(0, inventorySize));
+    }
+  }
+
+  private IInventoryStorage getSidedWrapper(@Nullable Direction direction) {
+    if (inventory instanceof WorldlyContainer && direction != null) {
+      return new SidedInventoryStorage(this, direction);
+    } else {
+      return this;
+    }
+  }
+
+  @Override
+  public String toString() {
+    return "InventoryStorage[" + DebugMessages.forInventory(inventory) + "]";
+  }
+
+  @Override
+  public ItemStack getStackInSlot(int slot) {
+    return inventory.getItem(slot);
+  }
+
+  @Override
+  public void setStackInSlot(int slot, ItemStack stack) {
+    inventory.setItem(slot, stack);
+  }
+
+  @Override
+  public int getSlotLimit(int slot) {
+    return inventory.getItem(slot).getMaxStackSize();
+  }
+
+  // Boolean is used to prevent allocation. Null values are not allowed by SnapshotParticipant.
+  class MarkDirtyParticipant extends SnapshotParticipant<Boolean> {
+    @Override
+    protected Boolean createSnapshot() {
+      return Boolean.TRUE;
+    }
+
+    @Override
+    protected void readSnapshot(Boolean snapshot) {
+    }
+
+    @Override
+    protected void onFinalCommit() {
+      inventory.setChanged();
+    }
+  }
+}