From 8c22b28b1bec3eb5a6c143ff6bdeb4322e25f9bf Mon Sep 17 00:00:00 2001
From: gzeon <im@gzeon.dev>
Date: Thu, 25 Apr 2024 00:11:32 +0800
Subject: [PATCH] refactor: move machine hash logic to osp

---
 src/challenge/ChallengeLib.sol     | 45 ----------------------------
 src/challenge/ChallengeManager.sol |  7 ++---
 src/osp/IOneStepProofEntry.sol     | 10 +++++++
 src/osp/OneStepProofEntry.sol      | 47 ++++++++++++++++++++++++++++++
 4 files changed, 59 insertions(+), 50 deletions(-)

diff --git a/src/challenge/ChallengeLib.sol b/src/challenge/ChallengeLib.sol
index 25ff894d..6aaa3579 100644
--- a/src/challenge/ChallengeLib.sol
+++ b/src/challenge/ChallengeLib.sol
@@ -48,51 +48,6 @@ library ChallengeLib {
         return challenge.timeUsedSinceLastMove() > challenge.current.timeLeft;
     }
 
-    function getStartMachineHash(bytes32 globalStateHash, bytes32 wasmModuleRoot)
-        internal
-        pure
-        returns (bytes32)
-    {
-        // Start the value stack with the function call ABI for the entrypoint
-        Value[] memory startingValues = new Value[](3);
-        startingValues[0] = ValueLib.newRefNull();
-        startingValues[1] = ValueLib.newI32(0);
-        startingValues[2] = ValueLib.newI32(0);
-        ValueArray memory valuesArray = ValueArray({inner: startingValues});
-        ValueStack memory values = ValueStack({proved: valuesArray, remainingHash: 0});
-        ValueStack memory internalStack;
-        StackFrameWindow memory frameStack;
-
-        Machine memory mach = Machine({
-            status: MachineStatus.RUNNING,
-            valueStack: values,
-            internalStack: internalStack,
-            frameStack: frameStack,
-            globalStateHash: globalStateHash,
-            moduleIdx: 0,
-            functionIdx: 0,
-            functionPc: 0,
-            modulesRoot: wasmModuleRoot
-        });
-        return mach.hash();
-    }
-
-    function getEndMachineHash(MachineStatus status, bytes32 globalStateHash)
-        internal
-        pure
-        returns (bytes32)
-    {
-        if (status == MachineStatus.FINISHED) {
-            return keccak256(abi.encodePacked("Machine finished:", globalStateHash));
-        } else if (status == MachineStatus.ERRORED) {
-            return keccak256(abi.encodePacked("Machine errored:"));
-        } else if (status == MachineStatus.TOO_FAR) {
-            return keccak256(abi.encodePacked("Machine too far:"));
-        } else {
-            revert("BAD_BLOCK_STATUS");
-        }
-    }
-
     function extractChallengeSegment(SegmentSelection calldata selection)
         internal
         pure
diff --git a/src/challenge/ChallengeManager.sol b/src/challenge/ChallengeManager.sol
index c5427e7e..99a81f54 100644
--- a/src/challenge/ChallengeManager.sol
+++ b/src/challenge/ChallengeManager.sol
@@ -233,11 +233,8 @@ contract ChallengeManager is DelegateCallAware, IChallengeManager {
         }
 
         bytes32[] memory segments = new bytes32[](2);
-        segments[0] = ChallengeLib.getStartMachineHash(
-            globalStateHashes[0],
-            challenge.wasmModuleRoot
-        );
-        segments[1] = ChallengeLib.getEndMachineHash(machineStatuses[1], globalStateHashes[1]);
+        segments[0] = osp.getStartMachineHash(globalStateHashes[0], challenge.wasmModuleRoot);
+        segments[1] = osp.getEndMachineHash(machineStatuses[1], globalStateHashes[1]);
 
         challenge.mode = ChallengeLib.ChallengeMode.EXECUTION;
 
diff --git a/src/osp/IOneStepProofEntry.sol b/src/osp/IOneStepProofEntry.sol
index fb00b74f..552819cf 100644
--- a/src/osp/IOneStepProofEntry.sol
+++ b/src/osp/IOneStepProofEntry.sol
@@ -11,6 +11,16 @@ library OneStepProofEntryLib {
 }
 
 interface IOneStepProofEntry {
+    function getStartMachineHash(bytes32 globalStateHash, bytes32 wasmModuleRoot)
+        external
+        pure
+        returns (bytes32);
+
+    function getEndMachineHash(MachineStatus status, bytes32 globalStateHash)
+        external
+        pure
+        returns (bytes32);
+
     function proveOneStep(
         ExecutionContext calldata execCtx,
         uint256 machineStep,
diff --git a/src/osp/OneStepProofEntry.sol b/src/osp/OneStepProofEntry.sol
index 390727c3..4ccc438c 100644
--- a/src/osp/OneStepProofEntry.sol
+++ b/src/osp/OneStepProofEntry.sol
@@ -31,6 +31,53 @@ contract OneStepProofEntry is IOneStepProofEntry {
         proverHostIo = proverHostIo_;
     }
 
+    // Copied from ChallengeLib.sol
+    function getStartMachineHash(bytes32 globalStateHash, bytes32 wasmModuleRoot)
+        external
+        pure
+        returns (bytes32)
+    {
+        // Start the value stack with the function call ABI for the entrypoint
+        Value[] memory startingValues = new Value[](3);
+        startingValues[0] = ValueLib.newRefNull();
+        startingValues[1] = ValueLib.newI32(0);
+        startingValues[2] = ValueLib.newI32(0);
+        ValueArray memory valuesArray = ValueArray({inner: startingValues});
+        ValueStack memory values = ValueStack({proved: valuesArray, remainingHash: 0});
+        ValueStack memory internalStack;
+        StackFrameWindow memory frameStack;
+
+        Machine memory mach = Machine({
+            status: MachineStatus.RUNNING,
+            valueStack: values,
+            internalStack: internalStack,
+            frameStack: frameStack,
+            globalStateHash: globalStateHash,
+            moduleIdx: 0,
+            functionIdx: 0,
+            functionPc: 0,
+            modulesRoot: wasmModuleRoot
+        });
+        return mach.hash();
+    }
+
+    // Copied from ChallengeLib.sol
+    function getEndMachineHash(MachineStatus status, bytes32 globalStateHash)
+        external
+        pure
+        returns (bytes32)
+    {
+        if (status == MachineStatus.FINISHED) {
+            return keccak256(abi.encodePacked("Machine finished:", globalStateHash));
+        } else if (status == MachineStatus.ERRORED) {
+            return keccak256(abi.encodePacked("Machine errored:"));
+        } else if (status == MachineStatus.TOO_FAR) {
+            return keccak256(abi.encodePacked("Machine too far:"));
+        } else {
+            revert("BAD_BLOCK_STATUS");
+        }
+    }
+
     function proveOneStep(
         ExecutionContext calldata execCtx,
         uint256 machineStep,