Add local context for rvsol

testinprod-io · Feb 1, 2024 · 5e6db03 · 5e6db03
1 parent 110e5ff
commit 5e6db03
Showing 1 changed file with 24 additions and 6 deletions.
diff --git a/rvsol/src/Step.sol b/rvsol/src/Step.sol
@@ -11,7 +11,7 @@ contract Step {
     }
 
     // Executes a single RISC-V instruction, starting from
-    function step(bytes calldata stateData, bytes calldata proof) public returns (bytes32) {
+    function step(bytes calldata stateData, bytes calldata proof, bytes32 localContext) public returns (bytes32) {
         assembly {
             function revertWithCode(code) {
                 mstore(0, code)
@@ -287,6 +287,7 @@ contract Step {
                 revert(0, 0)
             }
             function proofContentOffset() -> out { // since we can't reference proof.offset in functions, blame Yul
+                // 100+362+(32-362%32)+32=516
                 out := 516
             }
             if iszero(eq(proof.offset, proofContentOffset())) {
@@ -772,10 +773,27 @@ contract Step {
                 revertWithCode(0xbadf00d0)
             }
 
-            function readPreimageValue(addr, count) -> out {
+            function localize(preImageKey, localContext_) -> localizedKey {
+                // Grab the current free memory pointer to restore later.
+                let memPtr := mload(0x40)
+                // Store the local data key and caller next to each other in memory for hashing.
+                mstore(0, preImageKey)
+                mstore(0x20, caller())
+                mstore(0x40, localContext_)
+                // Localize the key
+                localizedKey := or(and(keccak256(0, 0x60), not(shl(248, 0xFF))), shl(248, 1))
+                // Restore the free memory pointer
+                mstore(0x40, memPtr)
+            }
+
+            function readPreimageValue(addr, count, localContext_) -> out {
                 let preImageKey := getPreimageKey()
                 let offset := getPreimageOffset()
-
+                // If the preimage key is a local key, localize it in the context of the caller.
+                let preImageKeyPrefix := shr(248, preImageKey) // 256-8=248
+                if eq(preImageKeyPrefix, 1) {
+                    preImageKey := localize(preImageKey, localContext_)
+                }
                 // make call to pre-image oracle contract
                 let pdatB32, pdatlen := readPreimagePart(preImageKey, offset)
                 if iszero64(pdatlen) { // EOF
@@ -811,7 +829,7 @@ contract Step {
             //
             // Syscall handling
             //
-            function sysCall() {
+            function sysCall(localContext_) {
                 let a7 := getRegister(toU64(17))
                 switch a7
                 case 93 { // exit the calling thread. No multi-thread support yet, so just exit.
@@ -872,7 +890,7 @@ contract Step {
                         n := count
                         errCode := toU64(0)
                     } case 5 { // preimage read
-                        n := readPreimageValue(addr, count)
+                        n := readPreimageValue(addr, count, localContext_)
                         errCode := toU64(0)
                     } default {
                         n := u64Mask()         //  -1 (reading error)
@@ -1275,7 +1293,7 @@ contract Step {
                 case 0 { // 000 = ECALL/EBREAK
                     switch shr64(toU64(20), instr) // I-type, top 12 bits
                     case 0 { // imm12 = 000000000000 ECALL
-                        sysCall()
+                        sysCall(localContext)
                         setPC(add64(_pc, toU64(4)))
                     } default { // imm12 = 000000000001 EBREAK
                         setPC(add64(_pc, toU64(4))) // ignore breakpoint