Fix Segfault in Montezuma's Revenge

src/ale/emucore/TIA.cxx

@@ -465,6 +465,14 @@ bool TIA::load(Deserializer& in)
     myPOSM1 = (int16_t) in.getInt();
     myPOSBL = (int16_t) in.getInt();
 
+    // Ensure restored positions are in valid range [0, 159]
+    // Only called on state restore, so expensive modulo operations aren't a problem.
+    myPOSP0 = ((myPOSP0 % 160) + 160) % 160;
+    myPOSP1 = ((myPOSP1 % 160) + 160) % 160;
+    myPOSM0 = ((myPOSM0 % 160) + 160) % 160;
+    myPOSM1 = ((myPOSM1 % 160) + 160) % 160;
+    myPOSBL = ((myPOSBL % 160) + 160) % 160;
+
     myCurrentGRP0 = (uint8_t) in.getInt();
     myCurrentGRP1 = (uint8_t) in.getInt();
 
@@ -2293,6 +2301,11 @@ void TIA::poke(uint16_t addr, uint8_t value)
       int newx = hpos < HBLANK ? 3 : (((hpos - HBLANK) + 5) % 160);
 
       // Find out under what condition the player is being reset
+      // Clamp position to valid range to prevent out-of-bounds table access
+      // See: https://github.com/Farama-Foundation/Arcade-Learning-Environment/issues/11
+      // Only fires when RESP0/RESP1 is used in already out of bounds - should be rare.
+      if(myPOSP0 < 0 || myPOSP0 >= 160)
+        myPOSP0 = ((myPOSP0 % 160) + 160) % 160;
       int8_t when = ourPlayerPositionResetWhenTable[myNUSIZ0 & 7][myPOSP0][newx];
 
       // Player is being reset during the display of one of its copies
@@ -2337,6 +2350,11 @@ void TIA::poke(uint16_t addr, uint8_t value)
       int newx = hpos < HBLANK ? 3 : (((hpos - HBLANK) + 5) % 160);
 
       // Find out under what condition the player is being reset
+      // Clamp position to valid range to prevent out-of-bounds table access
+      // See: https://github.com/Farama-Foundation/Arcade-Learning-Environment/issues/11
+      // Only files when RESP0/RESP1 is used in already out of bounds - should be rare.
+      if(myPOSP1 < 0 || myPOSP1 >= 160)
+        myPOSP1 = ((myPOSP1 % 160) + 160) % 160;
       int8_t when = ourPlayerPositionResetWhenTable[myNUSIZ1 & 7][myPOSP1][newx];
 
       // Player is being reset during the display of one of its copies
@@ -2746,6 +2764,10 @@ void TIA::poke(uint16_t addr, uint8_t value)
       myPOSM1 += ourCompleteMotionTable[x][myHMM1];
       myPOSBL += ourCompleteMotionTable[x][myHMBL];
 
+      // Single-step wrapping is safe here: motion table values are in
+      // [-15, +8], so positions starting in [0, 159] land in [-15, 167]
+      // — always within one correction of the valid range. The RESP and
+      // load() clamps elsewhere guarantee we enter this path in-bounds.
       if(myPOSP0 >= 160)
         myPOSP0 -= 160;
       else if(myPOSP0 < 0)

tests/python/test_python_interface.py

@@ -422,6 +422,173 @@ def test_display_screen(ale, tetris_rom_path):
     assert True
 
 
+def _find_tia_pos_offsets(data):
+    """Find byte offsets of TIA position registers in serialized ALE state.
+
+    The TIA section starts with a putString("TIA") = 4-byte length + "TIA".
+    Position registers (myPOSP0..myPOSBL) are 5 consecutive putInt values
+    at a fixed offset from the section start. Each putInt/putBool is 4 bytes.
+    """
+    # Find the "TIA" device marker: putString writes 4-byte LE length then chars
+    tia_len = int.to_bytes(3, 4, "little")  # length prefix for "TIA"
+    marker = tia_len + b"TIA"
+    idx = data.find(marker)
+    assert idx >= 0, "TIA section not found in serialized state"
+
+    # Count bytes from section start to myPOSP0.
+    # Order in TIA::save: putString("TIA"), then 8 ints, 1 int, 4 ints,
+    # 4 ints, 2 ints, 2 bools, 5 ints, 4 bools, 5 ints, 5 bools, 1 int,
+    # then myPOSP0..myPOSBL (5 ints).
+    # Total fields before myPOSP0: 8+1+4+4+2+2+5+4+5+5+1 = 41 putInt/putBool
+    # Each is 4 bytes. Plus the 7-byte string header.
+    offset = idx + 7 + 41 * 4  # = idx + 171
+    return offset
+
+
+def _write_le_int(data, offset, value):
+    """Write a signed 32-bit little-endian int into a bytearray."""
+    data[offset : offset + 4] = value.to_bytes(4, "little", signed=True)
+
+
+def _read_le_int(data, offset):
+    """Read a signed 32-bit little-endian int from bytes."""
+    return int.from_bytes(data[offset : offset + 4], "little", signed=True)
+
+
+def test_restore_state_with_corrupted_positions(tetris):
+    """Restoring a state with out-of-bounds TIA positions must not crash.
+
+    The TIA position clamping has three layers:
+    1. TIA::load — clamps positions on state restore via ((x % 160) + 160) % 160
+    2. RESP0/RESP1 — clamps myPOSP0/myPOSP1 before table lookup in TIA::poke
+    3. HMOVE — single-step wrapping after applying bounded motion deltas
+
+    This test exercises layer 1 by injecting out-of-bounds positions into
+    serialized state and verifying they are correctly wrapped on restore.
+    Subsequent gameplay exercises layers 2 and 3 to confirm no crash.
+    See: https://github.com/Farama-Foundation/Arcade-Learning-Environment/issues/11
+    """
+    for _ in range(50):
+        tetris.act(0)
+
+    state = tetris.cloneState()
+    raw = bytearray(state.__getstate__())
+    pos_offset = _find_tia_pos_offsets(raw)
+
+    # Verify baseline: positions should be in valid range
+    for i in range(5):
+        val = _read_le_int(raw, pos_offset + i * 4)
+        assert 0 <= val < 160, f"Position register {i} has unexpected value {val}"
+
+    # Comprehensive edge cases: (input, expected) where
+    # expected = ((x % 160) + 160) % 160 using C++ truncation semantics
+    test_cases = [
+        (-1, 159),  # just below valid range
+        (160, 0),  # just above valid range
+        (-500, 140),  # large negative
+        (9999, 79),  # large positive
+        (-160, 0),  # exact negative multiple
+        (320, 0),  # exact positive multiple
+        (32000, 0),  # large exact multiple (200 * 160)
+        (-32768, 32),  # int16_t minimum
+        (32767, 127),  # int16_t maximum
+        (0, 0),  # valid boundary (min, unchanged)
+        (159, 159),  # valid boundary (max, unchanged)
+    ]
+
+    for bad_val, expected in test_cases:
+        test_raw = bytearray(raw)
+        for i in range(5):
+            _write_le_int(test_raw, pos_offset + i * 4, bad_val)
+
+        # Restore the corrupted state — exercises TIA::load validation
+        corrupted = ale_py.ALEState.__new__(ale_py.ALEState)
+        corrupted.__setstate__(bytes(test_raw))
+        tetris.restoreState(corrupted)
+
+        # Verify TIA::load clamped all 5 positions correctly
+        restored_raw = tetris.cloneState().__getstate__()
+        restored_offset = _find_tia_pos_offsets(restored_raw)
+        for i in range(5):
+            val = _read_le_int(restored_raw, restored_offset + i * 4)
+            assert val == expected, (
+                f"Position register {i}: input {bad_val} should wrap to "
+                f"{expected}, got {val}"
+            )
+
+        # Play frames to exercise RESP0/RESP1 and HMOVE paths —
+        # without the clamps this would segfault
+        for _ in range(50):
+            tetris.act(0)
+            if tetris.game_over():
+                tetris.reset_game()
+
+
+def test_hmove_single_step_wrapping_bounds():
+    """Verify the bounds of single-step HMOVE position wrapping.
+
+    The HMOVE path wraps positions with:
+        if (x >= 160) x -= 160;
+        else if (x < 0) x += 160;
+
+    This is correct when positions enter HMOVE in [0, 159] and motion
+    values are in [-15, +8] (the range of ourCompleteMotionTable), giving
+    post-motion values in [-15, 167]. The RESP and load() clamps guarantee
+    the [0, 159] precondition.
+
+    This test verifies correctness across the full expected range and
+    asserts that values outside this range would NOT be corrected —
+    documenting the limits of the current implementation.
+    """
+
+    def single_step_wrap(x):
+        """Replicate the C++ HMOVE wrapping logic."""
+        if x >= 160:
+            x -= 160
+        elif x < 0:
+            x += 160
+        return x
+
+    def modulo_wrap(x):
+        """Full modular wrapping used by load() and RESP paths.
+
+        Equivalent to C++ ((x % 160) + 160) % 160 since Python's modulo
+        is non-negative for positive divisors.
+        """
+        return x % 160
+
+    # All valid (position, motion) pairs produce correct results.
+    # Positions in [0, 159], motion in [-15, +8] → post-motion in [-15, 167].
+    for pos in range(160):
+        for motion in range(-15, 9):
+            raw = pos + motion
+            result = single_step_wrap(raw)
+            assert 0 <= result < 160, (
+                f"single_step_wrap({raw}) = {result}, expected [0, 159] "
+                f"(pos={pos}, motion={motion})"
+            )
+            assert result == modulo_wrap(raw), (
+                f"single_step_wrap({raw}) = {result} != "
+                f"modulo_wrap({raw}) = {modulo_wrap(raw)}"
+            )
+
+    # Values outside the valid post-motion range that require more than
+    # one correction are NOT handled by single-step wrapping. Assert they
+    # remain out of [0, 159] — documenting the limits of the implementation
+    # and why the RESP/load() clamps are essential.
+    #
+    # These values can never occur in practice because the clamps guarantee
+    # positions enter HMOVE in [0, 159].
+    out = single_step_wrap(-161)  # -161 + 160 = -1
+    assert not (
+        0 <= out < 160
+    ), f"Expected -161 to remain invalid after single-step wrap, got {out}"
+    out = single_step_wrap(320)  # 320 - 160 = 160
+    assert not (
+        0 <= out < 160
+    ), f"Expected 320 to remain invalid after single-step wrap, got {out}"
+
+
 def test_set_logger(ale):
     ale.setLoggerMode(ale_py.LoggerMode.Info)
     ale.setLoggerMode(ale_py.LoggerMode.Warning)