Attempt at splitting up pocket computer logic

Oh, I hate the pocket computer code so much. Minecraft was really not
designed to attach this sort of behaviour to computers. This commit is
an attempt of cleaning this up[^1].

Firstly, we move the the pocket computer state (upgrades, light) out of
PocketServerComputer and into a new PocketBrain class. This now acts as
the sole source-of-truth, with all state being synced back to the
original item stack on the entity tick.

This also adds a new PocketHolder interface, which generalises over the
various types that can hold a pocket computer (players and item
entities right now, possibly lecterns in the future).

[^1]: I'd say simplifying, but this would be a lie.

projects/common/src/main/java/dan200/computercraft/shared/network/client/PocketComputerDataMessage.java

@@ -27,7 +27,7 @@ public class PocketComputerDataMessage implements NetworkMessage<ClientNetworkCo
     public PocketComputerDataMessage(PocketServerComputer computer, boolean sendTerminal) {
         clientId = computer.getInstanceUUID();
         state = computer.getState();
-        lightState = computer.getLight();
+        lightState = computer.getBrain().getLight();
         terminal = sendTerminal ? computer.getTerminalState() : null;
     }
 

projects/common/src/main/java/dan200/computercraft/shared/pocket/core/PocketBrain.java

@@ -0,0 +1,167 @@
+// SPDX-FileCopyrightText: 2024 The CC: Tweaked Developers
+//
+// SPDX-License-Identifier: MPL-2.0
+
+package dan200.computercraft.shared.pocket.core;
+
+import dan200.computercraft.api.peripheral.IPeripheral;
+import dan200.computercraft.api.pocket.IPocketAccess;
+import dan200.computercraft.api.pocket.IPocketUpgrade;
+import dan200.computercraft.api.upgrades.UpgradeData;
+import dan200.computercraft.core.computer.ComputerSide;
+import dan200.computercraft.shared.common.IColouredItem;
+import dan200.computercraft.shared.computer.core.ComputerFamily;
+import dan200.computercraft.shared.network.client.PocketComputerDataMessage;
+import dan200.computercraft.shared.network.server.ServerNetworking;
+import dan200.computercraft.shared.pocket.items.PocketComputerItem;
+import net.minecraft.nbt.CompoundTag;
+import net.minecraft.resources.ResourceLocation;
+import net.minecraft.world.entity.Entity;
+import net.minecraft.world.item.ItemStack;
+
+import javax.annotation.Nullable;
+import java.util.Collections;
+import java.util.Map;
+
+/**
+ * Holds additional state for a pocket computer. This includes pocket computer upgrade,
+ * {@linkplain IPocketAccess#getLight() light colour} and {@linkplain IPocketAccess#getColour() colour}.
+ * <p>
+ * This state is read when the brain is created, and written back to the holding item stack when the holding entity is
+ * ticked (see {@link #updateItem(ItemStack)}).
+ */
+public final class PocketBrain implements IPocketAccess {
+    private final PocketServerComputer computer;
+
+    private PocketHolder holder;
+
+    private boolean dirty = false;
+    private @Nullable UpgradeData<IPocketUpgrade> upgrade;
+    private int colour = -1;
+    private int lightColour = -1;
+
+    public PocketBrain(PocketHolder holder, int computerID, @Nullable String label, ComputerFamily family, @Nullable UpgradeData<IPocketUpgrade> upgrade) {
+        this.computer = new PocketServerComputer(this, holder, computerID, label, family);
+        this.holder = holder;
+        this.upgrade = UpgradeData.copyOf(upgrade);
+        invalidatePeripheral();
+    }
+
+    /**
+     * Get the corresponding pocket computer for this brain.
+     *
+     * @return The pocket computer.
+     */
+    public PocketServerComputer computer() {
+        return computer;
+    }
+
+    PocketHolder holder() {
+        return holder;
+    }
+
+    /**
+     * Update the position and holder for this computer.
+     *
+     * @param newHolder The new holder
+     */
+    public void updateHolder(PocketHolder newHolder) {
+        computer.setPosition(newHolder.level(), newHolder.blockPos());
+
+        var oldHolder = this.holder;
+        if (holder.equals(newHolder)) return;
+        holder = newHolder;
+
+        // If a new player has picked it up then rebroadcast the terminal to them
+        var oldPlayer = oldHolder instanceof PocketHolder.PlayerHolder p ? p.entity() : null;
+        if (newHolder instanceof PocketHolder.PlayerHolder player && player.entity() != oldPlayer) {
+            ServerNetworking.sendToPlayer(new PocketComputerDataMessage(computer, true), player.entity());
+        }
+    }
+
+    /**
+     * Write back properties of the pocket brain to the item.
+     *
+     * @param stack The pocket computer stack to update.
+     * @return Whether the item was changed.
+     */
+    public boolean updateItem(ItemStack stack) {
+        if (!dirty) return false;
+        this.dirty = false;
+
+        IColouredItem.setColourBasic(stack, colour);
+        PocketComputerItem.setUpgrade(stack, UpgradeData.copyOf(upgrade));
+        return true;
+    }
+
+    @Override
+    public @Nullable Entity getEntity() {
+        return holder instanceof PocketHolder.EntityHolder entity && holder.isValid(computer) ? entity.entity() : null;
+    }
+
+    @Override
+    public int getColour() {
+        return colour;
+    }
+
+    @Override
+    public void setColour(int colour) {
+        if (this.colour == colour) return;
+        dirty = true;
+        this.colour = colour;
+    }
+
+    @Override
+    public int getLight() {
+        return lightColour;
+    }
+
+    @Override
+    public void setLight(int colour) {
+        if (colour < 0 || colour > 0xFFFFFF) colour = -1;
+        lightColour = colour;
+    }
+
+    @Override
+    public CompoundTag getUpgradeNBTData() {
+        var upgrade = this.upgrade;
+        return upgrade == null ? new CompoundTag() : upgrade.data();
+    }
+
+    @Override
+    public void updateUpgradeNBTData() {
+        dirty = true;
+    }
+
+    @Override
+    public void invalidatePeripheral() {
+        var peripheral = upgrade == null ? null : upgrade.upgrade().createPeripheral(this);
+        computer.setPeripheral(ComputerSide.BACK, peripheral);
+    }
+
+    @Override
+    @Deprecated(forRemoval = true)
+    public Map<ResourceLocation, IPeripheral> getUpgrades() {
+        var upgrade = this.upgrade;
+        return upgrade == null ? Map.of() : Collections.singletonMap(upgrade.upgrade().getUpgradeID(), computer.getPeripheral(ComputerSide.BACK));
+    }
+
+    @Override
+    public @Nullable UpgradeData<IPocketUpgrade> getUpgrade() {
+        return upgrade;
+    }
+
+    /**
+     * Set the upgrade for this pocket computer, also updating the item stack.
+     * <p>
+     * Note this method is not thread safe - it must be called from the server thread.
+     *
+     * @param upgrade The new upgrade to set it to, may be {@code null}.
+     */
+    @Override
+    public void setUpgrade(@Nullable UpgradeData<IPocketUpgrade> upgrade) {
+        this.upgrade = upgrade;
+        dirty = true;
+        invalidatePeripheral();
+    }
+}

projects/common/src/main/java/dan200/computercraft/shared/pocket/core/PocketHolder.java

@@ -0,0 +1,102 @@
+// SPDX-FileCopyrightText: 2024 The CC: Tweaked Developers
+//
+// SPDX-License-Identifier: MPL-2.0
+
+package dan200.computercraft.shared.pocket.core;
+
+import dan200.computercraft.shared.computer.core.ServerComputer;
+import dan200.computercraft.shared.pocket.items.PocketComputerItem;
+import net.minecraft.core.BlockPos;
+import net.minecraft.server.level.ServerLevel;
+import net.minecraft.server.level.ServerPlayer;
+import net.minecraft.world.entity.Entity;
+import net.minecraft.world.entity.item.ItemEntity;
+
+/**
+ * An object that holds a pocket computer item.
+ */
+public sealed interface PocketHolder permits PocketHolder.EntityHolder {
+    /**
+     * The level this holder is in.
+     *
+     * @return The holder's level.
+     */
+    ServerLevel level();
+
+    /**
+     * The block position of this holder.
+     *
+     * @return The position of this holder.
+     */
+    BlockPos blockPos();
+
+    /**
+     * Determine if this holder is still valid for a particular computer.
+     *
+     * @param computer The current computer.
+     * @return Whether this holder is valid.
+     */
+    boolean isValid(ServerComputer computer);
+
+    /**
+     * Mark the pocket computer item as having changed.
+     */
+    void setChanged();
+
+    /**
+     * An {@link Entity} holding a pocket computer.
+     */
+    sealed interface EntityHolder extends PocketHolder permits PocketHolder.PlayerHolder, PocketHolder.ItemEntityHolder {
+        /**
+         * Get the entity holding this pocket computer.
+         *
+         * @return The holding entity.
+         */
+        Entity entity();
+
+        @Override
+        default ServerLevel level() {
+            return (ServerLevel) entity().level();
+        }
+
+        @Override
+        default BlockPos blockPos() {
+            return entity().blockPosition();
+        }
+    }
+
+    /**
+     * A pocket computer in a player's slot.
+     *
+     * @param entity The current player.
+     * @param slot   The slot the pocket computer is in.
+     */
+    record PlayerHolder(ServerPlayer entity, int slot) implements EntityHolder {
+        @Override
+        public boolean isValid(ServerComputer computer) {
+            return entity().isAlive() && PocketComputerItem.isServerComputer(computer, entity().getInventory().getItem(this.slot()));
+        }
+
+        @Override
+        public void setChanged() {
+            entity.getInventory().setChanged();
+        }
+    }
+
+    /**
+     * A pocket computer in an {@link ItemEntity}.
+     *
+     * @param entity The item entity.
+     */
+    record ItemEntityHolder(ItemEntity entity) implements EntityHolder {
+        @Override
+        public boolean isValid(ServerComputer computer) {
+            return entity().isAlive() && PocketComputerItem.isServerComputer(computer, this.entity().getItem());
+        }
+
+        @Override
+        public void setChanged() {
+            entity.setItem(entity.getItem().copy());
+        }
+    }
+}

projects/common/src/main/java/dan200/computercraft/shared/pocket/core/PocketServerComputer.java

@@ -4,12 +4,6 @@
 
 package dan200.computercraft.shared.pocket.core;
 
-import dan200.computercraft.api.peripheral.IPeripheral;
-import dan200.computercraft.api.pocket.IPocketAccess;
-import dan200.computercraft.api.pocket.IPocketUpgrade;
-import dan200.computercraft.api.upgrades.UpgradeData;
-import dan200.computercraft.core.computer.ComputerSide;
-import dan200.computercraft.shared.common.IColouredItem;
 import dan200.computercraft.shared.computer.core.ComputerFamily;
 import dan200.computercraft.shared.computer.core.ComputerState;
 import dan200.computercraft.shared.computer.core.ServerComputer;
@@ -18,29 +12,26 @@ import dan200.computercraft.shared.network.client.PocketComputerDataMessage;
 import dan200.computercraft.shared.network.client.PocketComputerDeletedClientMessage;
 import dan200.computercraft.shared.network.server.ServerNetworking;
 import dan200.computercraft.shared.pocket.items.PocketComputerItem;
-import net.minecraft.core.BlockPos;
-import net.minecraft.nbt.CompoundTag;
-import net.minecraft.resources.ResourceLocation;
-import net.minecraft.server.level.ServerLevel;
 import net.minecraft.server.level.ServerPlayer;
-import net.minecraft.world.entity.Entity;
-import net.minecraft.world.entity.LivingEntity;
-import net.minecraft.world.entity.item.ItemEntity;
-import net.minecraft.world.entity.player.Player;
-import net.minecraft.world.item.ItemStack;
 import net.minecraft.world.level.ChunkPos;
 
 import javax.annotation.Nullable;
-import java.util.Collections;
-import java.util.Map;
 import java.util.Set;
 
-public class PocketServerComputer extends ServerComputer implements IPocketAccess {
-    private @Nullable IPocketUpgrade upgrade;
-    private @Nullable Entity entity;
-    private ItemStack stack = ItemStack.EMPTY;
-
-    private int lightColour = -1;
+/**
+ * A {@link ServerComputer}-subclass for {@linkplain PocketComputerItem pocket computers}.
+ * <p>
+ * This extends default {@link ServerComputer} behaviour by also syncing pocket computer state to nearby players, and
+ * syncing the terminal to the current player.
+ * <p>
+ * The actual pocket computer state (upgrade, light) is maintained in {@link PocketBrain}. The two classes are tightly
+ * coupled, and maintain a reference to each other.
+ *
+ * @see PocketComputerDataMessage
+ * @see PocketComputerDeletedClientMessage
+ */
+public final class PocketServerComputer extends ServerComputer {
+    private final PocketBrain brain;
 
     // The state the previous tick, used to determine if the state needs to be sent to the client.
     private int oldLightColour = -1;
@@ -48,107 +39,13 @@ public class PocketServerComputer extends ServerComputer implements IPocketAcces
 
     private Set<ServerPlayer> tracking = Set.of();
 
-    public PocketServerComputer(ServerLevel world, BlockPos position, int computerID, @Nullable String label, ComputerFamily family) {
-        super(world, position, computerID, label, family, Config.pocketTermWidth, Config.pocketTermHeight);
+    PocketServerComputer(PocketBrain brain, PocketHolder holder, int computerID, @Nullable String label, ComputerFamily family) {
+        super(holder.level(), holder.blockPos(), computerID, label, family, Config.pocketTermWidth, Config.pocketTermHeight);
+        this.brain = brain;
     }
 
-    @Nullable
-    @Override
-    public Entity getEntity() {
-        var entity = this.entity;
-        if (entity == null || stack.isEmpty() || !entity.isAlive()) return null;
-
-        if (entity instanceof Player) {
-            var inventory = ((Player) entity).getInventory();
-            return inventory.items.contains(stack) || inventory.offhand.contains(stack) ? entity : null;
-        } else if (entity instanceof LivingEntity living) {
-            return living.getMainHandItem() == stack || living.getOffhandItem() == stack ? entity : null;
-        } else if (entity instanceof ItemEntity itemEntity) {
-            return itemEntity.getItem() == stack ? entity : null;
-        } else {
-            return null;
-        }
-    }
-
-    @Override
-    public int getColour() {
-        return IColouredItem.getColourBasic(stack);
-    }
-
-    @Override
-    public void setColour(int colour) {
-        IColouredItem.setColourBasic(stack, colour);
-        updateUpgradeNBTData();
-    }
-
-    @Override
-    public int getLight() {
-        return lightColour;
-    }
-
-    @Override
-    public void setLight(int colour) {
-        if (colour < 0 || colour > 0xFFFFFF) colour = -1;
-        lightColour = colour;
-    }
-
-    @Override
-    public CompoundTag getUpgradeNBTData() {
-        return PocketComputerItem.getUpgradeInfo(stack);
-    }
-
-    @Override
-    public void updateUpgradeNBTData() {
-        if (entity instanceof Player player) player.getInventory().setChanged();
-    }
-
-    @Override
-    public void invalidatePeripheral() {
-        var peripheral = upgrade == null ? null : upgrade.createPeripheral(this);
-        setPeripheral(ComputerSide.BACK, peripheral);
-    }
-
-    @Override
-    @Deprecated(forRemoval = true)
-    public Map<ResourceLocation, IPeripheral> getUpgrades() {
-        return upgrade == null ? Map.of() : Collections.singletonMap(upgrade.getUpgradeID(), getPeripheral(ComputerSide.BACK));
-    }
-
-    @Override
-    public @Nullable UpgradeData<IPocketUpgrade> getUpgrade() {
-        return upgrade == null ? null : UpgradeData.of(upgrade, getUpgradeNBTData());
-    }
-
-    /**
-     * Set the upgrade for this pocket computer, also updating the item stack.
-     * <p>
-     * Note this method is not thread safe - it must be called from the server thread.
-     *
-     * @param upgrade The new upgrade to set it to, may be {@code null}.
-     */
-    @Override
-    public void setUpgrade(@Nullable UpgradeData<IPocketUpgrade> upgrade) {
-        synchronized (this) {
-            PocketComputerItem.setUpgrade(stack, upgrade);
-            updateUpgradeNBTData();
-            this.upgrade = upgrade == null ? null : upgrade.upgrade();
-            invalidatePeripheral();
-        }
-    }
-
-    public synchronized void updateValues(@Nullable Entity entity, ItemStack stack, @Nullable IPocketUpgrade upgrade) {
-        if (entity != null) setPosition((ServerLevel) entity.level(), entity.blockPosition());
-
-        // If a new entity has picked it up then rebroadcast the terminal to them
-        if (entity != this.entity && entity instanceof ServerPlayer) markTerminalChanged();
-
-        this.entity = entity;
-        this.stack = stack;
-
-        if (this.upgrade != upgrade) {
-            this.upgrade = upgrade;
-            invalidatePeripheral();
-        }
+    public PocketBrain getBrain() {
+        return brain;
     }
 
     @Override
@@ -161,9 +58,10 @@ public class PocketServerComputer extends ServerComputer implements IPocketAcces
 
         // And now find any new players, add them to the tracking list, and broadcast state where appropriate.
         var state = getState();
-        if (oldLightColour != lightColour || oldComputerState != state) {
+        var light = brain.getLight();
+        if (oldLightColour != light || oldComputerState != state) {
             oldComputerState = state;
-            oldLightColour = lightColour;
+            oldLightColour = light;
 
             // Broadcast the state to all players
             ServerNetworking.sendToPlayers(new PocketComputerDataMessage(this, false), newTracking);
@@ -182,9 +80,9 @@ public class PocketServerComputer extends ServerComputer implements IPocketAcces
     protected void onTerminalChanged() {
         super.onTerminalChanged();
 
-        if (entity instanceof ServerPlayer player && entity.isAlive()) {
+        if (brain.holder() instanceof PocketHolder.PlayerHolder holder && holder.isValid(this)) {
             // Broadcast the terminal to the current player.
-            ServerNetworking.sendToPlayer(new PocketComputerDataMessage(this, true), player);
+            ServerNetworking.sendToPlayer(new PocketComputerDataMessage(this, true), holder.entity());
         }
     }
 

projects/common/src/main/java/dan200/computercraft/shared/pocket/items/PocketComputerItem.java

@@ -14,21 +14,26 @@ import dan200.computercraft.core.computer.ComputerSide;
 import dan200.computercraft.impl.PocketUpgrades;
 import dan200.computercraft.shared.common.IColouredItem;
 import dan200.computercraft.shared.computer.core.ComputerFamily;
+import dan200.computercraft.shared.computer.core.ServerComputer;
+import dan200.computercraft.shared.computer.core.ServerComputerRegistry;
 import dan200.computercraft.shared.computer.core.ServerContext;
 import dan200.computercraft.shared.computer.items.IComputerItem;
 import dan200.computercraft.shared.config.Config;
 import dan200.computercraft.shared.network.container.ComputerContainerData;
 import dan200.computercraft.shared.pocket.apis.PocketAPI;
+import dan200.computercraft.shared.pocket.core.PocketBrain;
+import dan200.computercraft.shared.pocket.core.PocketHolder;
 import dan200.computercraft.shared.pocket.core.PocketServerComputer;
 import dan200.computercraft.shared.pocket.inventory.PocketComputerMenuProvider;
 import dan200.computercraft.shared.util.IDAssigner;
+import dan200.computercraft.shared.util.InventoryUtil;
 import dan200.computercraft.shared.util.NBTUtil;
 import net.minecraft.ChatFormatting;
 import net.minecraft.nbt.CompoundTag;
 import net.minecraft.network.chat.Component;
 import net.minecraft.server.MinecraftServer;
 import net.minecraft.server.level.ServerLevel;
-import net.minecraft.world.Container;
+import net.minecraft.server.level.ServerPlayer;
 import net.minecraft.world.InteractionHand;
 import net.minecraft.world.InteractionResult;
 import net.minecraft.world.InteractionResultHolder;
@@ -72,12 +77,33 @@ public class PocketComputerItem extends Item implements IComputerItem, IMedia, I
         return result;
     }
 
-    private boolean tick(ItemStack stack, Entity entity, PocketServerComputer computer) {
-        var upgrade = getUpgrade(stack);
+    /**
+     * Tick a pocket computer.
+     *
+     * @param stack  The current pocket computer stack.
+     * @param holder The entity holding the pocket item.
+     * @param brain  The pocket computer brain.
+     */
+    private void tick(ItemStack stack, PocketHolder holder, PocketBrain brain) {
+        brain.updateHolder(holder);
 
-        computer.updateValues(entity, stack, upgrade);
+        // Update pocket upgrade
+        var upgrade = brain.getUpgrade();
+        if (upgrade != null) upgrade.upgrade().update(brain, brain.computer().getPeripheral(ComputerSide.BACK));
 
-        var changed = false;
+        if (updateItem(stack, brain)) holder.setChanged();
+    }
+
+    /**
+     * Copy properties from the brain back to the item stack.
+     *
+     * @param stack The current pocket computer stack.
+     * @param brain The current pocket brain.
+     * @return Whether the item was changed.
+     */
+    private boolean updateItem(ItemStack stack, PocketBrain brain) {
+        var changed = brain.updateItem(stack);
+        var computer = brain.computer();
 
         // Sync ID
         var id = computer.getID();
@@ -99,21 +125,20 @@ public class PocketComputerItem extends Item implements IComputerItem, IMedia, I
             stack.getOrCreateTag().putBoolean(NBT_ON, on);
         }
 
-        // Update pocket upgrade
-        if (upgrade != null) upgrade.update(computer, computer.getPeripheral(ComputerSide.BACK));
-
         return changed;
     }
 
     @Override
     public void inventoryTick(ItemStack stack, Level world, Entity entity, int slotNum, boolean selected) {
-        if (world.isClientSide) return;
-        Container inventory = entity instanceof Player player ? player.getInventory() : null;
-        var computer = createServerComputer((ServerLevel) world, entity, inventory, stack);
-        computer.keepAlive();
+        // This (in vanilla at least) is only called for players. Don't bother to handle other entities.
+        if (world.isClientSide || !(entity instanceof ServerPlayer player)) return;
 
-        var changed = tick(stack, entity, computer);
-        if (changed && inventory != null) inventory.setChanged();
+        // If we're in the inventory, create a computer and keep it alive.
+        var holder = new PocketHolder.PlayerHolder(player, slotNum);
+        var brain = getOrCreateBrain((ServerLevel) world, holder, stack);
+        brain.computer().keepAlive();
+
+        tick(stack, holder, brain);
     }
 
     @ForgeOverride
@@ -121,8 +146,11 @@ public class PocketComputerItem extends Item implements IComputerItem, IMedia, I
         var level = entity.level();
         if (level.isClientSide || level.getServer() == null) return false;
 
+        // If we're an item entity, tick an already existing computer (as to update the position), but do not keep the
+        // computer alive.
         var computer = getServerComputer(level.getServer(), stack);
-        if (computer != null && tick(stack, entity, computer)) entity.setItem(stack.copy());
+        if (computer != null) tick(stack, new PocketHolder.ItemEntityHolder(entity), computer.getBrain());
+
         return false;
     }
 
@@ -130,14 +158,18 @@ public class PocketComputerItem extends Item implements IComputerItem, IMedia, I
     public InteractionResultHolder<ItemStack> use(Level world, Player player, InteractionHand hand) {
         var stack = player.getItemInHand(hand);
         if (!world.isClientSide) {
-            var computer = createServerComputer((ServerLevel) world, player, player.getInventory(), stack);
+            var holder = new PocketHolder.PlayerHolder((ServerPlayer) player, InventoryUtil.getHandSlot(player, hand));
+            var brain = getOrCreateBrain((ServerLevel) world, holder, stack);
+            var computer = brain.computer();
             computer.turnOn();
 
             var stop = false;
             var upgrade = getUpgrade(stack);
             if (upgrade != null) {
-                computer.updateValues(player, stack, upgrade);
-                stop = upgrade.onRightClick(world, computer, computer.getPeripheral(ComputerSide.BACK));
+                brain.updateHolder(holder);
+                stop = upgrade.onRightClick(world, brain, computer.getPeripheral(ComputerSide.BACK));
+                // Sync back just in case. We don't need to setChanged, as we'll return the item anyway.
+                updateItem(stack, brain);
             }
 
             if (!stop) {
@@ -187,40 +219,51 @@ public class PocketComputerItem extends Item implements IComputerItem, IMedia, I
         return ComputerCraftAPI.MOD_ID;
     }
 
-    public PocketServerComputer createServerComputer(ServerLevel level, Entity entity, @Nullable Container inventory, ItemStack stack) {
-
+    private PocketBrain getOrCreateBrain(ServerLevel level, PocketHolder holder, ItemStack stack) {
         var registry = ServerContext.get(level.getServer()).registry();
-        var computer = (PocketServerComputer) registry.get(getSessionID(stack), getInstanceID(stack));
-        if (computer == null) {
+        {
+            var computer = getServerComputer(registry, stack);
+            if (computer != null) return computer.getBrain();
+        }
+
         var computerID = getComputerID(stack);
         if (computerID < 0) {
             computerID = ComputerCraftAPI.createUniqueNumberedSaveDir(level.getServer(), IDAssigner.COMPUTER);
             setComputerID(stack, computerID);
         }
 
-            computer = new PocketServerComputer(level, entity.blockPosition(), getComputerID(stack), getLabel(stack), getFamily());
+        var brain = new PocketBrain(holder, getComputerID(stack), getLabel(stack), getFamily(), getUpgradeWithData(stack));
+        var computer = brain.computer();
 
         var tag = stack.getOrCreateTag();
         tag.putInt(NBT_SESSION, registry.getSessionID());
         tag.putUUID(NBT_INSTANCE, computer.register());
 
-            var upgrade = getUpgrade(stack);
-
-            computer.updateValues(entity, stack, upgrade);
-            computer.addAPI(new PocketAPI(computer));
+        computer.addAPI(new PocketAPI(brain));
 
         // Only turn on when initially creating the computer, rather than each tick.
-            if (isMarkedOn(stack) && entity instanceof Player) computer.turnOn();
+        if (isMarkedOn(stack) && holder instanceof PocketHolder.PlayerHolder) computer.turnOn();
 
-            if (inventory != null) inventory.setChanged();
+        updateItem(stack, brain);
+
+        holder.setChanged();
+
+        return brain;
     }
 
-        return computer;
+    public static boolean isServerComputer(ServerComputer computer, ItemStack stack) {
+        return stack.getItem() instanceof PocketComputerItem
+            && getServerComputer(computer.getLevel().getServer(), stack) == computer;
+    }
+
+    @Nullable
+    public static PocketServerComputer getServerComputer(ServerComputerRegistry registry, ItemStack stack) {
+        return (PocketServerComputer) registry.get(getSessionID(stack), getInstanceID(stack));
     }
 
     @Nullable
     public static PocketServerComputer getServerComputer(MinecraftServer server, ItemStack stack) {
-        return (PocketServerComputer) ServerContext.get(server).registry().get(getSessionID(stack), getInstanceID(stack));
+        return getServerComputer(ServerContext.get(server).registry(), stack);
     }
 
     // IComputerItem implementation

projects/common/src/main/java/dan200/computercraft/shared/util/InventoryUtil.java

@@ -8,6 +8,9 @@ import net.minecraft.core.BlockPos;
 import net.minecraft.core.Direction;
 import net.minecraft.server.level.ServerLevel;
 import net.minecraft.world.Container;
+import net.minecraft.world.InteractionHand;
+import net.minecraft.world.entity.player.Inventory;
+import net.minecraft.world.entity.player.Player;
 import net.minecraft.world.item.ItemStack;
 import net.minecraft.world.phys.EntityHitResult;
 import net.minecraft.world.phys.Vec3;
@@ -18,6 +21,20 @@ public final class InventoryUtil {
     private InventoryUtil() {
     }
 
+    /**
+     * Get the inventory slot for a given hand.
+     *
+     * @param player The player to get the slot from.
+     * @param hand   The hand to get.
+     * @return The current slot.
+     */
+    public static int getHandSlot(Player player, InteractionHand hand) {
+        return switch (hand) {
+            case MAIN_HAND -> player.getInventory().selected;
+            case OFF_HAND -> Inventory.SLOT_OFFHAND;
+        };
+    }
+
     public static @Nullable Container getEntityContainer(ServerLevel level, BlockPos pos, Direction side) {
         var vecStart = new Vec3(
             pos.getX() + 0.5 + 0.6 * side.getStepX(),