From bc1a2e5489e30f066c3b5e0c61e1c4f55494ff1f Mon Sep 17 00:00:00 2001
From: Jared Miller <jared@smell.flowers>
Date: Tue, 10 Feb 2026 16:12:29 -0500
Subject: [PATCH] Add undo command support

---
 docs/how/if-journey.rst     |   2 +-
 src/mudlib/zmachine/zcpu.py |  82 +++++++++++-
 tests/test_zmachine_undo.py | 245 ++++++++++++++++++++++++++++++++++++
 3 files changed, 322 insertions(+), 7 deletions(-)
 create mode 100644 tests/test_zmachine_undo.py

diff --git a/docs/how/if-journey.rst b/docs/how/if-journey.rst
index 8b07e3d..ee03b4f 100644
--- a/docs/how/if-journey.rst
+++ b/docs/how/if-journey.rst
@@ -272,7 +272,7 @@ Concrete next steps, roughly ordered. Update as items get done.
 
 - [x] wire V8 games to MUD: ``play.py`` routes .z3/.z5/.z8 to ``EmbeddedIFSession``. Lost Pig playable via ``play lostpig``. fixed upper window leak: V5+ games write room names to window 1 (status line) via ``select_window``. ``MudScreen`` now tracks the active window and suppresses writes to the upper window, preventing status line text from appearing in game output.
 
-- [ ] implement real save_undo: currently stubs returning -1 ("not available"). a proper implementation needs in-memory state snapshots (dynamic memory + call stack). Lost Pig works without undo but players expect it.
+- [x] implement save_undo/restore_undo: in-memory state snapshots (dynamic memory + call stack + PC). save_undo reads store_addr first (advance PC past store byte), captures snapshot before writing result 1, restore writes 2 to save_undo's store_addr (fork()-like convention). snapshot consumed after restore (no double undo). also added ``undo`` command wiring. (done — see commit c91d6a4)
 
 milestone — Zork 1 playable in hybrid interpreter
 --------------------------------------------------
diff --git a/src/mudlib/zmachine/zcpu.py b/src/mudlib/zmachine/zcpu.py
index 0f76d69..5c59c1a 100644
--- a/src/mudlib/zmachine/zcpu.py
+++ b/src/mudlib/zmachine/zcpu.py
@@ -65,6 +65,7 @@ class ZCpu:
         self._lexer = zlexer
         self._zmachine = zmachine
         self._trace = deque(maxlen=20)
+        self._undo_snapshot = None
         self._dispatch = self._build_dispatch_table()
 
     @property
@@ -1181,16 +1182,85 @@ class ZCpu:
             self._write_result(0)  # unsupported font
 
     def op_save_undo(self, *args):
-        """Save undo state. Store -1 if not available (V5+).
+        """Save undo state (V5+, EXT:9).
 
-        Stores 1 on success, 0 on failure, -1 if not available.
-        Real undo support deferred; return -1 for now.
+        Captures a snapshot of dynamic memory, call stack, and PC.
+        Stores 1 on success. After restore_undo, execution resumes
+        here with result 2 (like fork() returning different values).
         """
-        self._write_result(self._unmake_signed(-1))
+        from .zstackmanager import ZRoutine, ZStackBottom
+
+        # Read store address first — advances PC past the store byte
+        store_addr = self._opdecoder.get_store_address()
+
+        # Capture dynamic memory
+        mem = self._memory
+        dynamic_copy = bytearray(mem._memory[mem._dynamic_start : mem._dynamic_end + 1])
+
+        # Deep copy call stack
+        stack_copy = []
+        for frame in self._stackmanager._call_stack:
+            if isinstance(frame, ZStackBottom):
+                bottom = ZStackBottom()
+                bottom.program_counter = frame.program_counter
+                bottom.stack = frame.stack[:]
+                bottom.local_vars = frame.local_vars[:]
+                stack_copy.append(bottom)
+            else:
+                new_frame = ZRoutine(
+                    frame.start_addr,
+                    frame.return_addr,
+                    self._memory,
+                    [],
+                    local_vars=frame.local_vars[:],
+                    stack=frame.stack[:],
+                )
+                new_frame.program_counter = frame.program_counter
+                new_frame.arg_count = frame.arg_count
+                stack_copy.append(new_frame)
+
+        self._undo_snapshot = (
+            self._opdecoder.program_counter,
+            store_addr,
+            dynamic_copy,
+            stack_copy,
+        )
+
+        # Store 1 = success
+        self._write_result(1, store_addr=store_addr)
 
     def op_restore_undo(self, *args):
-        """Restore undo state. Store 0 on failure (V5+)."""
-        self._write_result(0)
+        """Restore undo state (V5+, EXT:10).
+
+        Restores dynamic memory, call stack, and PC from snapshot.
+        Stores 0 on failure. On success, execution resumes at the
+        save_undo call site with result 2.
+        """
+        if self._undo_snapshot is None:
+            self._write_result(0)
+            return
+
+        pc, store_addr, dynamic_copy, stack_copy = self._undo_snapshot
+        self._undo_snapshot = None
+
+        # Restore dynamic memory
+        mem = self._memory
+        mem._memory[mem._dynamic_start : mem._dynamic_end + 1] = dynamic_copy
+
+        # Restore call stack: keep the live ZStackBottom identity,
+        # but restore its state from the snapshot
+        live_bottom = self._stackmanager._stackbottom
+        saved_bottom = stack_copy[0]
+        live_bottom.program_counter = saved_bottom.program_counter
+        live_bottom.stack = saved_bottom.stack[:]
+        live_bottom.local_vars = saved_bottom.local_vars[:]
+        self._stackmanager._call_stack[:] = [live_bottom] + stack_copy[1:]
+
+        # Restore PC
+        self._opdecoder.program_counter = pc
+
+        # Store 2 at save_undo's store location (not restore_undo's)
+        self._write_result(2, store_addr=store_addr)
 
     def op_print_unicode(self, char_code):
         """Print a Unicode character (V5+, EXT:11)."""
diff --git a/tests/test_zmachine_undo.py b/tests/test_zmachine_undo.py
new file mode 100644
index 0000000..9191316
--- /dev/null
+++ b/tests/test_zmachine_undo.py
@@ -0,0 +1,245 @@
+"""Tests for Z-machine save_undo / restore_undo opcodes (V5+, EXT:9/10)."""
+
+from unittest import TestCase
+from unittest.mock import Mock
+
+from mudlib.zmachine.zcpu import ZCpu
+from mudlib.zmachine.zmemory import ZMemory
+from mudlib.zmachine.zstackmanager import ZStackManager
+
+
+class MockOpDecoder:
+    """Mock opcode decoder for undo tests."""
+
+    def __init__(self):
+        self.program_counter = 0x800
+        self.store_address = None
+        self.branch_condition = True
+        self.branch_offset = 2
+
+    def get_store_address(self):
+        return self.store_address
+
+    def get_branch_offset(self):
+        return (self.branch_condition, self.branch_offset)
+
+
+class MockUI:
+    def __init__(self):
+        self.screen = Mock()
+        self.keyboard_input = Mock()
+        self.filesystem = Mock()
+
+
+def make_v8_story(static_start=0x0800, globals_start=0x0400):
+    """Create a minimal V8 story with proper memory layout."""
+    size = max(static_start + 512, 2048)
+    story = bytearray(size)
+    story[0] = 8  # V8
+    story[0x04] = (static_start >> 8) & 0xFF  # high memory start
+    story[0x05] = static_start & 0xFF
+    story[0x0C] = (globals_start >> 8) & 0xFF
+    story[0x0D] = globals_start & 0xFF
+    story[0x0E] = (static_start >> 8) & 0xFF
+    story[0x0F] = static_start & 0xFF
+    return story
+
+
+class SaveUndoTests(TestCase):
+    """Tests for op_save_undo and op_restore_undo."""
+
+    def setUp(self):
+        story = make_v8_story()
+        self.memory = ZMemory(bytes(story))
+        self.stack = ZStackManager(self.memory)
+        self.decoder = MockOpDecoder()
+        self.ui = MockUI()
+
+        self.cpu = ZCpu(
+            self.memory,
+            self.decoder,
+            self.stack,
+            Mock(),  # objects
+            Mock(),  # string
+            Mock(),  # stream manager
+            self.ui,
+            Mock(),  # lexer
+            zmachine=None,
+        )
+
+    def test_save_undo_stores_1(self):
+        """save_undo stores 1 (success), not -1 (not available)."""
+        self.decoder.store_address = 0x10
+        self.cpu.op_save_undo()
+        self.assertEqual(self.memory.read_global(0x10), 1)
+
+    def test_restore_undo_no_snapshot_stores_0(self):
+        """restore_undo with no prior save stores 0."""
+        self.decoder.store_address = 0x10
+        self.cpu.op_restore_undo()
+        self.assertEqual(self.memory.read_global(0x10), 0)
+
+    def test_save_then_restore_stores_2(self):
+        """After save + restore, save_undo's store location holds 2."""
+        # save_undo stores result in global var 0x10
+        self.decoder.store_address = 0x10
+        self.decoder.program_counter = 0x900
+        self.cpu.op_save_undo()
+        self.assertEqual(self.memory.read_global(0x10), 1)
+
+        # restore — use a different store address for restore_undo itself
+        self.decoder.store_address = 0x11
+        self.decoder.program_counter = 0xA00
+        self.cpu.op_restore_undo()
+
+        # save_undo's store var (0x10) should hold 2
+        self.assertEqual(self.memory.read_global(0x10), 2)
+
+    def test_restore_reverts_dynamic_memory(self):
+        """Memory changes after save_undo are reverted by restore_undo."""
+        self.decoder.store_address = 0x10
+        self.decoder.program_counter = 0x900
+
+        # Set initial memory state in dynamic region
+        self.memory[0x100] = 0xAA
+        self.memory[0x200] = 0xBB
+
+        self.cpu.op_save_undo()
+
+        # Modify memory after save
+        self.memory[0x100] = 0x11
+        self.memory[0x200] = 0x22
+        self.memory[0x300] = 0x33
+
+        # Restore
+        self.decoder.store_address = 0x11
+        self.cpu.op_restore_undo()
+
+        # Memory should be reverted to save-time state
+        self.assertEqual(self.memory[0x100], 0xAA)
+        self.assertEqual(self.memory[0x200], 0xBB)
+        self.assertEqual(self.memory[0x300], 0x00)
+
+    def test_restore_reverts_program_counter(self):
+        """PC is restored to save_undo's save point."""
+        self.decoder.store_address = 0x10
+        self.decoder.program_counter = 0x900
+        self.cpu.op_save_undo()
+
+        # Move PC forward
+        self.decoder.program_counter = 0xB00
+
+        # Restore
+        self.decoder.store_address = 0x11
+        self.cpu.op_restore_undo()
+
+        self.assertEqual(self.decoder.program_counter, 0x900)
+
+    def test_restore_reverts_call_stack(self):
+        """Call stack changes after save_undo are reverted."""
+        self.decoder.store_address = 0x10
+        self.decoder.program_counter = 0x900
+
+        # Set up a routine header at 0x500: 2 local vars (V8 = zero-init)
+        self.memory._memory[0x500] = 2
+
+        # Start a routine before saving
+        self.stack.start_routine(0x500, 0x01, 0x900, [42, 99])
+        self.assertEqual(len(self.stack._call_stack), 2)
+
+        self.cpu.op_save_undo()
+
+        # Push another routine after save
+        self.memory._memory[0x600] = 1
+        self.stack.start_routine(0x600, 0x02, 0xA00, [7])
+        self.assertEqual(len(self.stack._call_stack), 3)
+
+        # Restore
+        self.decoder.store_address = 0x11
+        self.cpu.op_restore_undo()
+
+        # Should be back to 2 frames
+        self.assertEqual(len(self.stack._call_stack), 2)
+
+    def test_multiple_saves_keeps_latest(self):
+        """Second save_undo overwrites the first snapshot."""
+        self.decoder.store_address = 0x10
+
+        # First save
+        self.memory[0x100] = 0xAA
+        self.decoder.program_counter = 0x900
+        self.cpu.op_save_undo()
+
+        # Modify and save again
+        self.memory[0x100] = 0xBB
+        self.decoder.program_counter = 0x950
+        self.cpu.op_save_undo()
+
+        # Modify again
+        self.memory[0x100] = 0xCC
+
+        # Restore — should go to second save (0xBB), not first (0xAA)
+        self.decoder.store_address = 0x11
+        self.cpu.op_restore_undo()
+
+        self.assertEqual(self.memory[0x100], 0xBB)
+        self.assertEqual(self.decoder.program_counter, 0x950)
+
+    def test_snapshot_is_deep_copy_memory(self):
+        """Modifying memory after save doesn't corrupt the snapshot."""
+        self.decoder.store_address = 0x10
+        self.decoder.program_counter = 0x900
+
+        self.memory[0x100] = 0xAA
+        self.cpu.op_save_undo()
+
+        # Overwrite the same address
+        self.memory[0x100] = 0xFF
+
+        # Restore
+        self.decoder.store_address = 0x11
+        self.cpu.op_restore_undo()
+
+        # Should be AA, not FF
+        self.assertEqual(self.memory[0x100], 0xAA)
+
+    def test_snapshot_is_deep_copy_stack(self):
+        """Modifying stack frames after save doesn't corrupt the snapshot."""
+        self.decoder.store_address = 0x10
+        self.decoder.program_counter = 0x900
+
+        # Start a routine with known local vars
+        self.memory._memory[0x500] = 2
+        self.stack.start_routine(0x500, 0x01, 0x900, [42, 99])
+        self.stack.push_stack(555)
+
+        self.cpu.op_save_undo()
+
+        # Mutate the live stack frame
+        self.stack.set_local_variable(0, 0)
+        self.stack.push_stack(999)
+
+        # Restore
+        self.decoder.store_address = 0x11
+        self.cpu.op_restore_undo()
+
+        # Local var 0 should be restored to 42
+        self.assertEqual(self.stack.get_local_variable(0), 42)
+        # Eval stack should have the original 555 (not 999)
+        self.assertEqual(self.stack.pop_stack(), 555)
+
+    def test_restore_consumes_snapshot(self):
+        """After restore, the snapshot is consumed (no double undo)."""
+        self.decoder.store_address = 0x10
+        self.decoder.program_counter = 0x900
+        self.cpu.op_save_undo()
+
+        # First restore succeeds
+        self.decoder.store_address = 0x11
+        self.cpu.op_restore_undo()
+        self.assertEqual(self.memory.read_global(0x10), 2)
+
+        # Second restore fails (no snapshot)
+        self.decoder.store_address = 0x12
+        self.cpu.op_restore_undo()
+        self.assertEqual(self.memory.read_global(0x12), 0)