Modify `loadLocalData` to work with the program

packages/contracts-bedrock/.gas-snapshot

@@ -85,36 +85,36 @@ FaucetTest:test_nonAdmin_drip_fails() (gas: 262520)
 FaucetTest:test_receive_succeeds() (gas: 17401)
 FaucetTest:test_withdraw_nonAdmin_reverts() (gas: 13145)
 FaucetTest:test_withdraw_succeeds() (gas: 78359)
-FaultDisputeGame_ResolvesCorrectly_CorrectRoot1:test_resolvesCorrectly_succeeds() (gas: 502209)
-FaultDisputeGame_ResolvesCorrectly_CorrectRoot2:test_resolvesCorrectly_succeeds() (gas: 508989)
-FaultDisputeGame_ResolvesCorrectly_CorrectRoot3:test_resolvesCorrectly_succeeds() (gas: 505730)
-FaultDisputeGame_ResolvesCorrectly_CorrectRoot4:test_resolvesCorrectly_succeeds() (gas: 508887)
-FaultDisputeGame_ResolvesCorrectly_CorrectRoot5:test_resolvesCorrectly_succeeds() (gas: 508226)
-FaultDisputeGame_ResolvesCorrectly_IncorrectRoot1:test_resolvesCorrectly_succeeds() (gas: 500974)
-FaultDisputeGame_ResolvesCorrectly_IncorrectRoot2:test_resolvesCorrectly_succeeds() (gas: 507754)
-FaultDisputeGame_ResolvesCorrectly_IncorrectRoot3:test_resolvesCorrectly_succeeds() (gas: 504495)
-FaultDisputeGame_ResolvesCorrectly_IncorrectRoot4:test_resolvesCorrectly_succeeds() (gas: 505652)
-FaultDisputeGame_ResolvesCorrectly_IncorrectRoot5:test_resolvesCorrectly_succeeds() (gas: 504991)
-FaultDisputeGame_Test:test_extraData_succeeds() (gas: 17449)
+FaultDisputeGame_ResolvesCorrectly_CorrectRoot1:test_resolvesCorrectly_succeeds() (gas: 501957)
+FaultDisputeGame_ResolvesCorrectly_CorrectRoot2:test_resolvesCorrectly_succeeds() (gas: 508759)
+FaultDisputeGame_ResolvesCorrectly_CorrectRoot3:test_resolvesCorrectly_succeeds() (gas: 505478)
+FaultDisputeGame_ResolvesCorrectly_CorrectRoot4:test_resolvesCorrectly_succeeds() (gas: 508657)
+FaultDisputeGame_ResolvesCorrectly_CorrectRoot5:test_resolvesCorrectly_succeeds() (gas: 507974)
+FaultDisputeGame_ResolvesCorrectly_IncorrectRoot1:test_resolvesCorrectly_succeeds() (gas: 500722)
+FaultDisputeGame_ResolvesCorrectly_IncorrectRoot2:test_resolvesCorrectly_succeeds() (gas: 507524)
+FaultDisputeGame_ResolvesCorrectly_IncorrectRoot3:test_resolvesCorrectly_succeeds() (gas: 504243)
+FaultDisputeGame_ResolvesCorrectly_IncorrectRoot4:test_resolvesCorrectly_succeeds() (gas: 505422)
+FaultDisputeGame_ResolvesCorrectly_IncorrectRoot5:test_resolvesCorrectly_succeeds() (gas: 504739)
+FaultDisputeGame_Test:test_extraData_succeeds() (gas: 17404)
 FaultDisputeGame_Test:test_gameData_succeeds() (gas: 17939)
-FaultDisputeGame_Test:test_gameStart_succeeds() (gas: 10293)
-FaultDisputeGame_Test:test_gameType_succeeds() (gas: 8238)
-FaultDisputeGame_Test:test_initialRootClaimData_succeeds() (gas: 17669)
-FaultDisputeGame_Test:test_move_clockCorrectness_succeeds() (gas: 419317)
-FaultDisputeGame_Test:test_move_clockTimeExceeded_reverts() (gas: 26399)
-FaultDisputeGame_Test:test_move_defendRoot_reverts() (gas: 13338)
-FaultDisputeGame_Test:test_move_duplicateClaim_reverts() (gas: 104076)
-FaultDisputeGame_Test:test_move_gameDepthExceeded_reverts() (gas: 411414)
-FaultDisputeGame_Test:test_move_gameNotInProgress_reverts() (gas: 10968)
-FaultDisputeGame_Test:test_move_nonExistentParent_reverts() (gas: 24633)
-FaultDisputeGame_Test:test_move_simpleAttack_succeeds() (gas: 108178)
-FaultDisputeGame_Test:test_resolve_challengeContested_succeeds() (gas: 226490)
-FaultDisputeGame_Test:test_resolve_notInProgress_reverts() (gas: 9702)
-FaultDisputeGame_Test:test_resolve_rootContested_succeeds() (gas: 110643)
-FaultDisputeGame_Test:test_resolve_rootUncontestedClockNotExpired_succeeds() (gas: 21482)
-FaultDisputeGame_Test:test_resolve_rootUncontested_succeeds() (gas: 27311)
-FaultDisputeGame_Test:test_resolve_teamDeathmatch_succeeds() (gas: 398794)
-FaultDisputeGame_Test:test_rootClaim_succeeds() (gas: 8203)
+FaultDisputeGame_Test:test_gameStart_succeeds() (gas: 10315)
+FaultDisputeGame_Test:test_gameType_succeeds() (gas: 8260)
+FaultDisputeGame_Test:test_initialRootClaimData_succeeds() (gas: 17624)
+FaultDisputeGame_Test:test_move_clockCorrectness_succeeds() (gas: 419180)
+FaultDisputeGame_Test:test_move_clockTimeExceeded_reverts() (gas: 26421)
+FaultDisputeGame_Test:test_move_defendRoot_reverts() (gas: 13360)
+FaultDisputeGame_Test:test_move_duplicateClaim_reverts() (gas: 104120)
+FaultDisputeGame_Test:test_move_gameDepthExceeded_reverts() (gas: 411546)
+FaultDisputeGame_Test:test_move_gameNotInProgress_reverts() (gas: 11012)
+FaultDisputeGame_Test:test_move_nonExistentParent_reverts() (gas: 24677)
+FaultDisputeGame_Test:test_move_simpleAttack_succeeds() (gas: 108110)
+FaultDisputeGame_Test:test_resolve_challengeContested_succeeds() (gas: 226511)
+FaultDisputeGame_Test:test_resolve_notInProgress_reverts() (gas: 9657)
+FaultDisputeGame_Test:test_resolve_rootContested_succeeds() (gas: 110642)
+FaultDisputeGame_Test:test_resolve_rootUncontestedClockNotExpired_succeeds() (gas: 21437)
+FaultDisputeGame_Test:test_resolve_rootUncontested_succeeds() (gas: 27288)
+FaultDisputeGame_Test:test_resolve_teamDeathmatch_succeeds() (gas: 398859)
+FaultDisputeGame_Test:test_rootClaim_succeeds() (gas: 8225)
 FeeVault_Test:test_constructor_succeeds() (gas: 18185)
 GasBenchMark_L1CrossDomainMessenger:test_sendMessage_benchmark_0() (gas: 352113)
 GasBenchMark_L1CrossDomainMessenger:test_sendMessage_benchmark_1() (gas: 2950320)
@@ -436,9 +436,9 @@ OptimistTest:test_tokenIdOfAddress_returnsOwnerID_succeeds() (gas: 63730)
 OptimistTest:test_tokenURI_returnsCorrectTokenURI_succeeds() (gas: 195908)
 OptimistTest:test_transferFrom_soulbound_reverts() (gas: 75512)
 PreimageOracle_Test:test_keccak256PreimageKey_succeeds() (gas: 420)
-PreimageOracle_Test:test_loadKeccak256PreimagePart_outOfBoundsOffset_reverts() (gas: 8842)
-PreimageOracle_Test:test_loadKeccak256PreimagePart_succeeds() (gas: 76496)
-PreimageOracle_Test:test_loadLocalData_onePart_succeeds() (gas: 76266)
+PreimageOracle_Test:test_loadKeccak256PreimagePart_outOfBoundsOffset_reverts() (gas: 8864)
+PreimageOracle_Test:test_loadKeccak256PreimagePart_succeeds() (gas: 76518)
+PreimageOracle_Test:test_loadLocalData_onePart_succeeds() (gas: 76339)
 ProxyAdmin_Test:test_chugsplashChangeProxyAdmin_succeeds() (gas: 35586)
 ProxyAdmin_Test:test_chugsplashGetProxyAdmin_succeeds() (gas: 15675)
 ProxyAdmin_Test:test_chugsplashGetProxyImplementation_succeeds() (gas: 51084)

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

@@ -74,31 +74,34 @@ contract PreimageOracle {
     function loadLocalData(
         uint256 _ident,
         bytes32 _word,
-        uint8 _size
+        uint256 _size,
+        uint256 _partOffset
     ) external returns (bytes32 key_) {
         // Compute the localized key from the given local identifier.
         key_ = PreimageKeyLib.localizeIdent(_ident);
 
-        // Load both parts of the local data word into storage for future
-        // reads.
-        bytes32 part1;
-        assembly {
-            // The first part is prepended with an 8 byte length prefix and contains
-            // the first 24 bytes of the passed word.
-            part1 := or(shl(192, _size), shr(64, _word))
+        // Revert if the given part offset is not within bounds.
+        if (_partOffset > _size + 8 || _size > 32) {
+            revert("part offset must be within bounds");
         }
 
-        // Store the first part with offset 0.
-        preimagePartOk[key_][0] = true;
-        preimageParts[key_][0] = part1;
+        // Prepare the local data part at the given offset
+        bytes32 part;
+        assembly {
+            // Clean the memory in [0x20, 0x40)
+            mstore(0x20, 0x00)
+
+            // Store the full local data in scratch space.
+            mstore(0x00, shl(192, _size))
+            mstore(0x08, _word)
 
-        // If the size is greater than 24, we need to store a second part as well.
-        if (_size > 24) {
-            bytes32 part2 = _word << 192;
-            preimagePartOk[key_][32] = true;
-            preimageParts[key_][32] = part2;
+            // Prepare the local data part at the requested offset.
+            part := mload(_partOffset)
         }
 
+        // Store the first part with `_partOffset`.
+        preimagePartOk[key_][_partOffset] = true;
+        preimageParts[key_][_partOffset] = part;
         // Assign the length of the preimage at the localized key.
         preimageLengths[key_] = _size;
     }

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

@@ -252,34 +252,34 @@ contract FaultDisputeGame is IFaultDisputeGame, Clone, Semver {
     }
 
     /// @inheritdoc IFaultDisputeGame
-    function addLocalData(uint256 _ident) external {
+    function addLocalData(uint256 _ident, uint256 _partOffset) external {
         // INVARIANT: Local data can only be added if the game is currently in progress.
         if (status != GameStatus.IN_PROGRESS) revert GameNotInProgress();
 
         IPreimageOracle oracle = VM.oracle();
         if (_ident == 1) {
             // Load the L1 head hash into the game's local context in the preimage oracle.
-            oracle.loadLocalData(_ident, Hash.unwrap(l1Head), 32);
+            oracle.loadLocalData(_ident, Hash.unwrap(l1Head), 32, _partOffset);
         } else if (_ident == 2) {
             // Load the earliest output root that commits to the passed L2 block number
             // into the game's local context in the preimage oracle.
             Types.OutputProposal memory proposal = L2_OUTPUT_ORACLE.getL2OutputAfter(
                 l2BlockNumber()
             );
-            oracle.loadLocalData(_ident, proposal.outputRoot, 32);
+            oracle.loadLocalData(_ident, proposal.outputRoot, 32, _partOffset);
         } else if (_ident == 3) {
             // Load the root claim into the game's local context in the preimage oracle.
-            oracle.loadLocalData(_ident, Claim.unwrap(rootClaim()), 32);
+            oracle.loadLocalData(_ident, Claim.unwrap(rootClaim()), 32, _partOffset);
         } else if (_ident == 4) {
             // Load the L2 block number into the game's local context in the preimage oracle.
             // The L2 block number is stored as a big-endian uint64 in the upper 8 bytes of the
             // passed word.
-            oracle.loadLocalData(_ident, bytes32(l2BlockNumber() << 192), 8);
+            oracle.loadLocalData(_ident, bytes32(l2BlockNumber() << 192), 8, _partOffset);
         } else if (_ident == 5) {
             // Load the chain ID into the game's local context in the preimage oracle.
             // The chain ID is stored as a big-endian uint64 in the upper 8 bytes of the
             // passed word.
-            oracle.loadLocalData(_ident, bytes32(block.chainid << 192), 8);
+            oracle.loadLocalData(_ident, bytes32(block.chainid << 192), 8, _partOffset);
         }
     }
 

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

@@ -42,6 +42,7 @@ interface IPreimageOracle {
     function loadLocalData(
         uint256 _ident,
         bytes32 _word,
-        uint8 _size
+        uint256 _size,
+        uint256 _partOffset
     ) external returns (bytes32 key_);
 }

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

@@ -55,7 +55,8 @@ 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.
-    function addLocalData(uint256 _ident) external;
+    /// @param _partOffset The offset of the data to post.
+    function addLocalData(uint256 _ident, uint256 _partOffset) external;
 
     /// @notice Returns the L1 block hash at the time of the game's creation.
     function l1Head() external view returns (Hash l1Head_);

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

@@ -29,17 +29,18 @@ contract PreimageOracle_Test is Test {
         uint256 ident = 1;
         bytes32 word = bytes32(uint256(0xdeadbeef) << 224);
         uint8 size = 4;
+        uint8 partOffset = 0;
 
         // Load the local data into the preimage oracle under the test contract's context.
-        bytes32 contextKey = oracle.loadLocalData(ident, word, size);
+        bytes32 contextKey = oracle.loadLocalData(ident, word, size, partOffset);
 
         // Validate that the pre-image part is set
-        bool ok = oracle.preimagePartOk(contextKey, 0);
+        bool ok = oracle.preimagePartOk(contextKey, partOffset);
         assertTrue(ok);
 
         // Validate the local data part
         bytes32 expectedPart = 0x0000000000000004deadbeef0000000000000000000000000000000000000000;
-        assertEq(oracle.preimageParts(contextKey, 0), expectedPart);
+        assertEq(oracle.preimageParts(contextKey, partOffset), expectedPart);
 
         // Validate the local data length
         uint256 length = oracle.preimageLengths(contextKey);
@@ -50,32 +51,31 @@ contract PreimageOracle_Test is Test {
     function testFuzz_loadLocalData_varyingLength_succeeds(
         uint256 ident,
         bytes32 word,
-        uint256 size
+        uint256 size,
+        uint256 partOffset
     ) public {
         // Bound the size to [0, 32]
         size = bound(size, 0, 32);
+        // Bound the part offset to [0, size + 8]
+        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));
+        bytes32 contextKey = oracle.loadLocalData(ident, word, uint8(size), uint8(partOffset));
 
         // Validate that the first local data part is set
-        bool ok = oracle.preimagePartOk(contextKey, 0);
+        bool ok = oracle.preimagePartOk(contextKey, partOffset);
         assertTrue(ok);
         // Validate the first local data part
-        bytes32 expectedPart1 = bytes32((size << 192) | (uint256(word) >> 64));
-        assertEq(oracle.preimageParts(contextKey, 0), expectedPart1);
-
-        // If the size is > 24, validate the second part. Otherwise, ensure
-        // that the second part is not set.
-        ok = oracle.preimagePartOk(contextKey, 32);
-        if (size > 24) {
-            assertTrue(ok);
-            // Validate the second local data part
-            bytes32 expectedPart2 = word << 192;
-            assertEq(oracle.preimageParts(contextKey, 32), expectedPart2);
-        } else {
-            assertFalse(ok);
+        bytes32 expectedPart;
+        assembly {
+            mstore(0x20, 0x00)
+
+            mstore(0x00, shl(192, size))
+            mstore(0x08, word)
+
+            expectedPart := mload(partOffset)
         }
+        assertEq(oracle.preimageParts(contextKey, partOffset), expectedPart);
 
         // Validate the local data length
         uint256 length = oracle.preimageLengths(contextKey);