Begin contextualization of local keys

packages/contracts-bedrock/src/cannon/MIPS.sol

@@ -141,8 +141,9 @@ contract MIPS {
     }
 
     /// @notice Handles a syscall.
+    /// @param _localContext The local key context for the preimage oracle.
     /// @return out_ The hashed MIPS state.
-    function handleSyscall() internal returns (bytes32 out_) {
+    function handleSyscall(uint256 _localContext) internal returns (bytes32 out_) {
         unchecked {
             // Load state from memory
             State memory state;
@@ -202,7 +203,7 @@ contract MIPS {
                     bytes32 preimageKey = state.preimageKey;
                     // If the preimage key is a local key, localize it in the context of the caller.
                     if (uint8(preimageKey[0]) == 1) {
-                        preimageKey = PreimageKeyLib.localize(preimageKey);
+                        preimageKey = PreimageKeyLib.localize(preimageKey, _localContext);
                     }
                     (bytes32 dat, uint256 datLen) = ORACLE.readPreimage(preimageKey, state.preimageOffset);
 
@@ -508,8 +509,8 @@ contract MIPS {
     function proofOffset(uint8 _proofIndex) internal pure returns (uint256 offset_) {
         unchecked {
             // A proof of 32 bit memory, with 32-byte leaf values, is (32-5)=27 bytes32 entries.
-            // And the leaf value itself needs to be encoded as well. And proof.offset == 388
-            offset_ = 388 + (uint256(_proofIndex) * (28 * 32));
+            // And the leaf value itself needs to be encoded as well. And proof.offset == 420
+            offset_ = 420 + (uint256(_proofIndex) * (28 * 32));
             uint256 s = 0;
             assembly {
                 s := calldatasize()
@@ -621,7 +622,11 @@ contract MIPS {
 
     /// @notice Executes a single step of the vm.
     ///         Will revert if any required input state is missing.
-    function step(bytes calldata stateData, bytes calldata proof) public returns (bytes32) {
+    /// @param _stateData The encoded state witness data.
+    /// @param _proof The encoded proof data for leaves within the MIPS VM's memory.
+    /// @param _localContext The local key context for the preimage oracle. Optional, can be set as a constant
+    ///                     if the caller only requires one set of local keys.
+    function step(bytes calldata _stateData, bytes calldata _proof, uint256 _localContext) public returns (bytes32) {
         unchecked {
             State memory state;
 
@@ -631,16 +636,16 @@ contract MIPS {
                     // expected state mem offset check
                     revert(0, 0)
                 }
-                if iszero(eq(mload(0x40), mul(32, 48))) {
+                if iszero(eq(mload(0x40), shl(5, 48))) {
                     // expected memory check
                     revert(0, 0)
                 }
-                if iszero(eq(stateData.offset, 100)) {
-                    // 32*3+4=100 expected state data offset
+                if iszero(eq(_stateData.offset, 132)) {
+                    // 32*4+4=132 expected state data offset
                     revert(0, 0)
                 }
-                if iszero(eq(proof.offset, 388)) {
-                    // 100+32+256=388 expected proof offset
+                if iszero(eq(_proof.offset, 420)) {
+                    // 132+32+256=420 expected proof offset
                     revert(0, 0)
                 }
 
@@ -653,7 +658,7 @@ contract MIPS {
                 }
 
                 // Unpack state from calldata into memory
-                let c := stateData.offset // calldata offset
+                let c := _stateData.offset // calldata offset
                 let m := 0x80 // mem offset
                 c, m := putField(c, m, 32) // memRoot
                 c, m := putField(c, m, 32) // preimageKey
@@ -764,7 +769,7 @@ contract MIPS {
 
                 // syscall (can read and write)
                 if (func == 0xC) {
-                    return handleSyscall();
+                    return handleSyscall(_localContext);
                 }
 
                 // lo and hi registers

packages/contracts-bedrock/src/cannon/PreimageKeyLib.sol

@@ -7,14 +7,15 @@ library PreimageKeyLib {
     /// @notice Generates a context-specific local key for the given local data identifier.
     /// @dev See `localize` for a description of the localization operation.
     /// @param _ident The identifier of the local data. [0, 32) bytes in size.
+    /// @param _localContext The local context for the key.
     /// @return key_ The context-specific local key.
-    function localizeIdent(uint256 _ident) internal view returns (bytes32 key_) {
+    function localizeIdent(uint256 _ident, uint256 _localContext) internal view returns (bytes32 key_) {
         assembly {
             // Set the type byte in the given identifier to `1` (Local). We only care about
             // the [1, 32) bytes in this value.
             key_ := or(shl(248, 1), and(_ident, not(shl(248, 0xFF))))
         }
-        key_ = localize(key_);
+        key_ = localize(key_, _localContext);
     }
 
     /// @notice Localizes a given local data key for the caller's context.
@@ -22,14 +23,20 @@ library PreimageKeyLib {
     ///      localize(k) = H(k .. sender) & ~(0xFF << 248) | (0x01 << 248)
     ///      where H is the Keccak-256 hash function.
     /// @param _key The local data key to localize.
+    /// @param _localContext The local context for the key.
     /// @return localizedKey_ The localized local data key.
-    function localize(bytes32 _key) internal view returns (bytes32 localizedKey_) {
+    function localize(bytes32 _key, uint256 _localContext) internal view returns (bytes32 localizedKey_) {
         assembly {
+            // Grab the current free memory pointer to restore later.
+            let ptr := mload(0x40)
             // Store the local data key and caller next to each other in memory for hashing.
             mstore(0, _key)
             mstore(0x20, caller())
+            mstore(0x40, _localContext)
             // Localize the key with the above `localize` operation.
-            localizedKey_ := or(and(keccak256(0, 0x40), not(shl(248, 0xFF))), shl(248, 1))
+            localizedKey_ := or(and(keccak256(0, 0x60), not(shl(248, 0xFF))), shl(248, 1))
+            // Restore the free memory pointer.
+            mstore(0x40, ptr)
         }
     }
 

packages/contracts-bedrock/src/cannon/PreimageOracle.sol

@@ -34,6 +34,7 @@ contract PreimageOracle is IPreimageOracle {
     /// @inheritdoc IPreimageOracle
     function loadLocalData(
         uint256 _ident,
+        uint256 _localContext,
         bytes32 _word,
         uint256 _size,
         uint256 _partOffset
@@ -42,7 +43,7 @@ contract PreimageOracle is IPreimageOracle {
         returns (bytes32 key_)
     {
         // Compute the localized key from the given local identifier.
-        key_ = PreimageKeyLib.localizeIdent(_ident);
+        key_ = PreimageKeyLib.localizeIdent(_ident, _localContext);
 
         // Revert if the given part offset is not within bounds.
         if (_partOffset > _size + 8 || _size > 32) {

packages/contracts-bedrock/src/cannon/interfaces/IPreimageOracle.sol

@@ -13,6 +13,8 @@ interface IPreimageOracle {
 
     /// @notice Loads of local data part into the preimage oracle.
     /// @param _ident The identifier of the local data.
+    /// @param _localContext The local key context for the preimage oracle. Optional, can be set as a constant
+    ///                     if the caller only requires one set of local keys.
     /// @param _word The local data word.
     /// @param _size The number of bytes in `_word` to load.
     /// @param _partOffset The offset of the local data part to write to the oracle.
@@ -32,6 +34,7 @@ interface IPreimageOracle {
     ///      └────────────┴────────────────────────┘
     function loadLocalData(
         uint256 _ident,
+        uint256 _localContext,
         bytes32 _word,
         uint256 _size,
         uint256 _partOffset

packages/contracts-bedrock/src/dispute/FaultDisputeGame.sol

@@ -278,7 +278,7 @@ contract FaultDisputeGame is IFaultDisputeGame, Clone, ISemver {
     }
 
     /// @inheritdoc IFaultDisputeGame
-    function addLocalData(uint256 _ident, uint256 _partOffset) external {
+    function addLocalData(uint256 _ident, uint256 _l2BlockNumber, uint256 _partOffset) external {
         // INVARIANT: Local data can only be added if the game is currently in progress.
         if (status != GameStatus.IN_PROGRESS) revert GameNotInProgress();
 
@@ -289,6 +289,8 @@ contract FaultDisputeGame is IFaultDisputeGame, Clone, ISemver {
             mstore(0x1C, loadLocalDataSelector)
             // Store the `_ident` argument
             mstore(0x20, _ident)
+            // Store the `_localContext` argument
+            mstore(0x40, _l2BlockNumber)
             // Store the data to load
             let data
             switch _ident
@@ -319,16 +321,16 @@ contract FaultDisputeGame is IFaultDisputeGame, Clone, ISemver {
                 // Revert with  `InvalidLocalIdent()`
                 revert(0x1C, 0x04)
             }
-            mstore(0x40, data)
+            mstore(0x60, data)
             // Store the size of the data to load
             // _ident > 3 ? 8 : 32
-            mstore(0x60, shl(sub(0x05, shl(0x01, gt(_ident, 0x03))), 0x01))
+            mstore(0x80, shl(sub(0x05, shl(0x01, gt(_ident, 0x03))), 0x01))
             // Store the part offset of the data
-            mstore(0x80, _partOffset)
+            mstore(0xA0, _partOffset)
 
             // Attempt to add the local data to the preimage oracle and bubble up the revert
             // if it fails.
-            if iszero(call(gas(), oracle, 0x00, 0x1C, 0x84, 0x00, 0x00)) {
+            if iszero(call(gas(), oracle, 0x00, 0x1C, 0xA4, 0x00, 0x00)) {
                 returndatacopy(0x00, 0x00, returndatasize())
                 revert(0x00, returndatasize())
             }

packages/contracts-bedrock/src/dispute/interfaces/IFaultDisputeGame.sol

@@ -70,8 +70,10 @@ interface IFaultDisputeGame is IDisputeGame {
 
     /// @notice Posts the requested local data to the VM's `PreimageOralce`.
     /// @param _ident The local identifier of the data to post.
+    /// @param _l2BlockNumber The L2 block number being disputed. This serves as the local context for the
+    ///                      `PreimageOracle` key.
     /// @param _partOffset The offset of the data to post.
-    function addLocalData(uint256 _ident, uint256 _partOffset) external;
+    function addLocalData(uint256 _ident, uint256 _l2BlockNumber, uint256 _partOffset) external;
 
     /// @notice Resolves the subgame rooted at the given claim index.
     /// @dev This function must be called bottom-up in the DAG

packages/contracts-bedrock/test/FaultDisputeGame.t.sol

@@ -483,12 +483,12 @@ contract FaultDisputeGame_Test is FaultDisputeGame_Init {
     }
 
     /// @dev Tests that adding local data with an out of bounds identifier reverts.
-    function testFuzz_addLocalData_oob_reverts(uint256 _ident) public {
+    function testFuzz_addLocalData_oob_reverts(uint256 _ident, uint256 _localContext) public {
         // [1, 5] are valid local data identifiers.
         if (_ident <= 5) _ident = 0;
 
         vm.expectRevert(InvalidLocalIdent.selector);
-        gameProxy.addLocalData(_ident, 0);
+        gameProxy.addLocalData(_ident, _localContext, 0);
     }
 
     /// @dev Tests that local data is loaded into the preimage oracle correctly.
@@ -508,8 +508,8 @@ contract FaultDisputeGame_Test is FaultDisputeGame_Init {
         for (uint256 i = 1; i <= 5; i++) {
             uint256 expectedLen = i > 3 ? 8 : 32;
 
-            gameProxy.addLocalData(i, 0);
-            bytes32 key = _getKey(i);
+            gameProxy.addLocalData(i, 0, 0);
+            bytes32 key = _getKey(i, 0);
             (bytes32 dat, uint256 datLen) = oracle.readPreimage(key, 0);
             assertEq(dat >> 0xC0, bytes32(expectedLen));
             // Account for the length prefix if i > 3 (the data stored
@@ -519,8 +519,8 @@ contract FaultDisputeGame_Test is FaultDisputeGame_Init {
             // total.)
             assertEq(datLen, expectedLen + (i > 3 ? 8 : 0));
 
-            gameProxy.addLocalData(i, 8);
-            key = _getKey(i);
+            gameProxy.addLocalData(i, 0, 8);
+            key = _getKey(i, 0);
             (dat, datLen) = oracle.readPreimage(key, 8);
             assertEq(dat, data[i - 1]);
             assertEq(datLen, expectedLen);
@@ -528,8 +528,8 @@ contract FaultDisputeGame_Test is FaultDisputeGame_Init {
     }
 
     /// @dev Helper to get the localized key for an identifier in the context of the game proxy.
-    function _getKey(uint256 _ident) internal view returns (bytes32) {
-        bytes32 h = keccak256(abi.encode(_ident | (1 << 248), address(gameProxy)));
+    function _getKey(uint256 _ident, uint256 _localContext) internal view returns (bytes32) {
+        bytes32 h = keccak256(abi.encode(_ident | (1 << 248), address(gameProxy), _localContext));
         return bytes32((uint256(h) & ~uint256(0xFF << 248)) | (1 << 248));
     }
 

packages/contracts-bedrock/test/PreimageOracle.t.sol

@@ -32,7 +32,7 @@ contract PreimageOracle_Test is Test {
         uint8 partOffset = 0;
 
         // Load the local data into the preimage oracle under the test contract's context.
-        bytes32 contextKey = oracle.loadLocalData(ident, word, size, partOffset);
+        bytes32 contextKey = oracle.loadLocalData(ident, 0, word, size, partOffset);
 
         // Validate that the pre-image part is set
         bool ok = oracle.preimagePartOk(contextKey, partOffset);
@@ -50,6 +50,7 @@ contract PreimageOracle_Test is Test {
     /// @notice Tests that context-specific data [0, 32] bytes in length can be loaded correctly.
     function testFuzz_loadLocalData_varyingLength_succeeds(
         uint256 ident,
+        uint256 localContext,
         bytes32 word,
         uint256 size,
         uint256 partOffset
@@ -62,7 +63,7 @@ contract PreimageOracle_Test is Test {
         partOffset = bound(partOffset, 0, size + 8);
 
         // Load the local data into the preimage oracle under the test contract's context.
-        bytes32 contextKey = oracle.loadLocalData(ident, word, uint8(size), uint8(partOffset));
+        bytes32 contextKey = oracle.loadLocalData(ident, localContext, word, uint8(size), uint8(partOffset));
 
         // Validate that the first local data part is set
         bool ok = oracle.preimagePartOk(contextKey, partOffset);
@@ -112,7 +113,7 @@ contract PreimageOracle_Test is Test {
         uint256 offset = preimage.length + 9;
 
         vm.expectRevert(PartOffsetOOB.selector);
-        oracle.loadLocalData(1, preimage, 32, offset);
+        oracle.loadLocalData(1, 0, preimage, 32, offset);
     }
 
     /// @notice Tests that a pre-image cannot be set with an out-of-bounds offset.