[DO NOT MERGE] [MEL] Changes to OSP contracts

src/challengeV2/EdgeChallengeManager.sol

@@ -395,12 +395,8 @@ contract EdgeChallengeManager is IEdgeChallengeManager, Initializable {
 
         assertionChain.validateConfig(prevAssertionHash, prevConfig);
 
-        // TODO(PR 427): OSP contracts are marked as pending work in the PR.
-        // Inbox-position-based checks no longer apply; use type(uint256).max as
-        // a stopgap until OSP is rewired against `nextParentChainBlockHash`.
         ExecutionContext memory execCtx = ExecutionContext({
-            maxInboxMessagesRead: type(uint256).max,
-            bridge: assertionChain.bridge(),
+            targetParentChainBlockHash: prevConfig.nextParentChainBlockHash,
             initialWasmModuleRoot: prevConfig.wasmModuleRoot
         });
 

src/mocks/SimpleOneStepProofEntry.sol

@@ -9,6 +9,7 @@ import "../state/Deserialize.sol";
 
 contract SimpleOneStepProofEntry is IOneStepProofEntry {
     using GlobalStateLib for GlobalState;
+    using MELStateLib for MELState;
 
     // End the batch after 2000 steps. This results in 11 blocks for an honest validator.
     // This constant must be synchronized with the one in execution/engine.go
@@ -26,23 +27,33 @@ contract SimpleOneStepProofEntry is IOneStepProofEntry {
         uint256 step,
         bytes32 beforeHash,
         bytes calldata proof
-    ) external view returns (bytes32 afterHash) {
+    ) external pure returns (bytes32 afterHash) {
         if (proof.length == 0) {
             revert("EMPTY_PROOF");
         }
         GlobalState memory globalState;
         uint256 offset;
-        (globalState.u64Vals[0], offset) = Deserialize.u64(proof, offset);
-        (globalState.u64Vals[1], offset) = Deserialize.u64(proof, offset);
-        if (step > 0 && (beforeHash[0] == 0 || globalState.getPositionInMessage() == 0)) {
+        (globalState.bytes32Vals[3], offset) = Deserialize.b32(proof, offset); // MELNextMsgHash
+        (globalState.u64Vals[2], offset) = Deserialize.u64(proof, offset); // MELMsgCount
+        (globalState.u64Vals[3], offset) = Deserialize.u64(proof, offset); // MELExecutedMsgCount
+
+        MELState memory melState;
+        (melState.parentChainBlockHash, offset) = Deserialize.b32(proof, offset);
+
+        if (step > 0 && (beforeHash[0] == 0 || globalState.getMELNextMsgHash() == bytes32(0))) {
             // We end the block when the first byte of the hash hits 0 or we advance a batch
             return beforeHash;
         }
-        if (globalState.getInboxPosition() >= execCtx.maxInboxMessagesRead) {
-            // We can't continue further because we've hit the max inbox messages read
+        if (
+            melState.parentChainBlockHash == execCtx.targetParentChainBlockHash
+                && globalState.getMELExecutedMsgCount() >= melState.msgCount
+        ) {
+            // We can't continue further because we've executed all messages up to this melState
             return beforeHash;
         }
         require(globalState.hash() == beforeHash, "BAD_PROOF");
+
+        // TODO: modify this logic once execution_engine.go is modified
         globalState.u64Vals[1]++;
         if (globalState.u64Vals[1] % STEPS_PER_BATCH == 0) {
             globalState.u64Vals[0]++;

src/osp/HashProofHelper.sol

@@ -4,32 +4,42 @@
 
 pragma solidity ^0.8.0;
 
+import "./IHashProofHelper.sol";
 import "../libraries/CryptographyPrimitives.sol";
 
-/// @dev The requested hash preimage at the given offset has not been proven yet
-error NotProven(bytes32 fullHash, uint64 offset);
-
-contract HashProofHelper {
+contract HashProofHelper is IHashProofHelper {
+    /// @dev Tracks an in-progress split preimage proof
     struct KeccakState {
+        /// @dev Offset determining which slice is extracted and stored from the preimage (up to 32 bytes)
         uint64 offset;
+        /// @dev The bytes being collected for the [offset, offset+32) slice, built up across chunks
         bytes part;
+        /// @dev The 1600-bit keccak internal state as 25 × 64-bit words
+        ///      (stored in column-major order to match CryptographyPrimitives.keccakF's layout)
         uint64[25] state;
+        /// @dev Total bytes of preimage data absorbed so far across all chunks
         uint256 length;
     }
 
+    /// @dev Stores a 32-byte (or shorter) slice extracted from a fully-proven preimage.
     struct PreimagePart {
+        /// @dev Whether this entry has been set by a completed proof.
         bool proven;
+        /// @dev The extracted slice at this offset. Empty if offset >= preimage length.
         bytes part;
     }
 
+    /// @dev Completed proofs, keyed by (keccak256 hash of the full preimage, byte offset)
     mapping(bytes32 => mapping(uint64 => PreimagePart)) private preimageParts;
+    /// @dev In-progress split proofs, keyed by msg.sender
     mapping(address => KeccakState) public keccakStates;
 
-    event PreimagePartProven(bytes32 indexed fullHash, uint64 indexed offset, bytes part);
-
+    /// @dev Maximum bytes stored per part — matches one EVM word and one WASM ReadPreImage result
     uint256 private constant MAX_PART_LENGTH = 32;
+    /// @dev Number of bytes absorbed into the keccak sponge state per round — matches the keccak256 rate for 1600-bit state
     uint256 private constant KECCAK_ROUND_INPUT = 136;
 
+    /// @inheritdoc IHashProofHelper
     function proveWithFullPreimage(
         bytes calldata data,
         uint64 offset
@@ -38,18 +48,16 @@ contract HashProofHelper {
         bytes memory part;
         if (data.length > offset) {
             uint256 partLength = data.length - offset;
-            if (partLength > 32) {
-                partLength = 32;
+            if (partLength > MAX_PART_LENGTH) {
+                partLength = MAX_PART_LENGTH;
             }
             part = data[offset:(offset + partLength)];
         }
         preimageParts[fullHash][offset] = PreimagePart({proven: true, part: part});
         emit PreimagePartProven(fullHash, offset, part);
     }
 
-    // Flags: a bitset signaling various things about the proof, ordered from least to most significant bits.
-    //   0th bit: indicates that this data is the final chunk of preimage data.
-    //   1st bit: indicates that the preimage part currently being built should be cleared before this.
+    /// @inheritdoc IHashProofHelper
     function proveWithSplitPreimage(
         bytes calldata data,
         uint64 offset,
@@ -67,7 +75,12 @@ contract HashProofHelper {
         } else {
             require(state.offset == offset, "DIFF_OFFSET");
         }
+
+        // Update the keccak state with the new data
+        // (updates state.state and state.length)
         keccakUpdate(state, data, isFinal);
+
+        // Obtain the `part`
         if (uint256(offset) + MAX_PART_LENGTH > startLength && offset < state.length) {
             uint256 startIdx = 0;
             if (offset > startLength) {
@@ -81,9 +94,14 @@ contract HashProofHelper {
                 state.part.push(data[i]);
             }
         }
+
+        // If this is not the final chunk, we can't yet determine the full hash, so we return early
         if (!isFinal) {
             return bytes32(0);
         }
+
+        // Obtain the full hash from the keccak state
+        // (the first 32 bytes)
         for (uint256 i = 0; i < 32; i++) {
             uint256 stateIdx = i / 8;
             // work around our weird keccakF function state ordering
@@ -96,21 +114,33 @@ contract HashProofHelper {
         delete keccakStates[msg.sender];
     }
 
+    /**
+     * @notice Absorbs data into the keccak sponge state, one 136-byte round at a time.
+     *         On the final call, applies keccak padding.
+     * @param state The in-progress keccak state to update (modified in place)
+     * @param data The next chunk of preimage bytes to absorb
+     * @param isFinal If true, pads and processes the final block
+     */
     function keccakUpdate(KeccakState storage state, bytes calldata data, bool isFinal) internal {
         state.length += data.length;
         while (true) {
             if (data.length == 0 && !isFinal) {
                 break;
             }
+
+            // XOR in the next chunk of data, padding if necessary
+            // (1 byte per iteration)
             for (uint256 i = 0; i < KECCAK_ROUND_INPUT; i++) {
                 uint8 b = 0;
                 if (i < data.length) {
                     b = uint8(data[i]);
                 } else {
-                    // Padding
+                    // Padding added in the final chunk or in a chunk on its own if the final chunk is exactly round-aligned
+                    // 1st bit (LSB) is set if this is the first byte after the data
                     if (i == data.length) {
                         b |= uint8(0x01);
                     }
+                    // Last bit (MSB) is always set in the final chunk
                     if (i == KECCAK_ROUND_INPUT - 1) {
                         b |= uint8(0x80);
                     }
@@ -124,22 +154,29 @@ contract HashProofHelper {
             for (uint256 i = 0; i < 25; i++) {
                 state256[i] = state.state[i];
             }
+
+            // Scramble the state with keccakF
             state256 = CryptographyPrimitives.keccakF(state256);
+
+            // Write the new state back to storage
             for (uint256 i = 0; i < 25; i++) {
                 state.state[i] = uint64(state256[i]);
             }
+
+            // Strict inequality, because if data is an exact multiple of the round size, keccak still adds a padding chunk
             if (data.length < KECCAK_ROUND_INPUT) {
                 break;
             }
             data = data[KECCAK_ROUND_INPUT:];
         }
     }
 
+    /// @notice Deletes the caller's in-progress split proof state
     function clearSplitProof() external {
         delete keccakStates[msg.sender];
     }
 
-    /// Retrieves up to 32 bytes of the preimage of fullHash at the given offset, reverting if it hasn't been proven yet.
+    /// @inheritdoc IHashProofHelper
     function getPreimagePart(
         bytes32 fullHash,
         uint64 offset

src/osp/IHashProofHelper.sol

@@ -0,0 +1,64 @@
+// Copyright 2021-2022, Offchain Labs, Inc.
+// For license information, see https://github.com/OffchainLabs/nitro-contracts/blob/main/LICENSE
+// SPDX-License-Identifier: BUSL-1.1
+
+pragma solidity ^0.8.0;
+
+/**
+ * @title IHashProofHelper
+ * @notice Proves keccak256 preimages and stores slices of up to 32 bytes for later retrieval.
+ *         Used by the OneStepProverHostIo to verify large preimages that don't fit
+ *         in a single transaction's calldata. The preimage is uploaded in advance
+ *         (in one or more transactions), and later retrieved during the one-step proof.
+ */
+interface IHashProofHelper {
+    /// @notice Emitted when a preimage part is proven and stored
+    event PreimagePartProven(bytes32 indexed fullHash, uint64 indexed offset, bytes part);
+
+    /// @notice The requested preimage part has not been proven yet
+    error NotProven(bytes32 fullHash, uint64 offset);
+
+    /**
+     * @notice Proves a preimage in a single transaction by providing all data at once
+     * @param data The full preimage
+     * @param offset Byte offset into the preimage. Up to 32 bytes at this offset are stored
+     *        as the proven part. If offset is past the end of data, an empty part is stored.
+     * @return fullHash keccak256(data)
+     */
+    function proveWithFullPreimage(
+        bytes calldata data,
+        uint64 offset
+    ) external returns (bytes32 fullHash);
+
+    /**
+     * @notice Proves a preimage across multiple transactions by uploading data in chunks.
+     *         Each chunk is absorbed into an incremental keccak256 computation stored
+     *         per sender. On the final chunk, the hash is finalized and the proven part
+     *         is stored.
+     * @param data The next chunk of preimage data. Must be a multiple of 136 bytes
+     *        (the keccak round size) unless this is the final chunk
+     * @param offset Byte offset into the full preimage for the 32-byte part to prove.
+     *        Must be the same value across all chunks of a single proof.
+     * @param flags a bitset signaling various things about the proof, ordered from least to most significant bits:
+     *        - bit 0: indicates that this data is the final chunk of preimage data, which triggers padding, hash finalization, and storage
+     *        - bit 1: indicates that the unfinished preimage currently being built should be cleared before this
+     * @return fullHash if the final chunk is passed, keccak256 of the full preimage; otherwise, bytes32(0)
+     */
+    function proveWithSplitPreimage(
+        bytes calldata data,
+        uint64 offset,
+        uint256 flags
+    ) external returns (bytes32 fullHash);
+
+    /**
+     * @notice Retrieves a previously proven preimage part
+     * @param fullHash The keccak256 hash of the preimage
+     * @param offset The byte offset that was used when the part was proven
+     * @return The proven bytes (up to 32). Reverts with NotProven if the part
+     *         at this hash and offset hasn't been proven yet.
+     */
+    function getPreimagePart(
+        bytes32 fullHash,
+        uint64 offset
+    ) external view returns (bytes memory);
+}

src/osp/IOneStepProver.sol

@@ -12,8 +12,7 @@ import "../bridge/ISequencerInbox.sol";
 import "../bridge/IBridge.sol";
 
 struct ExecutionContext {
-    uint256 maxInboxMessagesRead;
-    IBridge bridge;
+    bytes32 targetParentChainBlockHash;
     bytes32 initialWasmModuleRoot;
 }
 

src/osp/OneStepProofEntry.sol

@@ -15,6 +15,7 @@ contract OneStepProofEntry is IOneStepProofEntry {
     using MerkleProofLib for MerkleProof;
     using MachineLib for Machine;
     using GlobalStateLib for GlobalState;
+    using MELStateLib for MELState;
     using MultiStackLib for MultiStack;
 
     using ValueStackLib for ValueStack;
@@ -104,9 +105,20 @@ contract OneStepProofEntry is IOneStepProofEntry {
                 GlobalState memory globalState;
                 (globalState, offset) = Deserialize.globalState(proof, offset);
                 require(globalState.hash() == mach.globalStateHash, "BAD_GLOBAL_STATE");
+
+                MELState memory melState;
+                (melState, offset) = Deserialize.melState(proof, offset);
+                require(melState.hash() == globalState.getMELStateHash(), "BAD_MEL_STATE");
+
+                // The machine has finished processing a message and we're at the start of the next execution segment (machineStep == 0).
+                // If the MELState is not at its target (meaning that it hasn't finished extracting messages, which should only happen before the extraction process is started),
+                // or if all messages were extracted, but there are still messages to be executed in MEL, we kickstart the machine.
                 if (
                     mach.status == MachineStatus.FINISHED && machineStep == 0
-                        && globalState.getInboxPosition() < execCtx.maxInboxMessagesRead
+                        && (
+                            melState.parentChainBlockHash != execCtx.targetParentChainBlockHash
+                                || globalState.getMELExecutedMsgCount() < melState.msgCount
+                        )
                 ) {
                     // Kickstart the machine
                     return getStartMachineHash(mach.globalStateHash, execCtx.initialWasmModuleRoot);
@@ -187,6 +199,7 @@ contract OneStepProofEntry is IOneStepProofEntry {
             )
                 || (opcode >= Instructions.VALIDATE_CERTIFICATE && opcode <= Instructions.UNLINK_MODULE)
                 || (opcode >= Instructions.NEW_COTHREAD && opcode <= Instructions.SWITCH_COTHREAD)
+                || (opcode == Instructions.GET_END_PARENT_CHAIN_BLOCK_HASH)
         ) {
             prover = proverHostIo;
         } else {