Merge pull request #9 from ethereum-optimism/tip/pcw109550/local-context

pcw109550 · web-flow · commit 8af5ae93f7f0 · 2024-02-13T08:01:02.000Z
feat: Local Context
diff --git a/rvgo/fast/evm.go b/rvgo/fast/evm.go
@@ -3,7 +3,8 @@ package fast
 import "github.com/ethereum/go-ethereum/crypto"
 
 var (
-	StepBytes4                      = crypto.Keccak256([]byte("step(bytes,bytes)"))[:4]
+	StepBytes4                      = crypto.Keccak256([]byte("step(bytes,bytes,bytes32)"))[:4]
 	CheatBytes4                     = crypto.Keccak256([]byte("cheat(uint256,bytes32,bytes32,uint256)"))[:4]
+	CheatLocalKeyBytes4             = crypto.Keccak256([]byte("cheatLocalKey(uint256,bytes32,bytes32,uint256,bytes32)"))[:4]
 	LoadKeccak256PreimagePartBytes4 = crypto.Keccak256([]byte("loadKeccak256PreimagePart(uint256,bytes)"))[:4]
 )
diff --git a/rvgo/fast/witness.go b/rvgo/fast/witness.go
@@ -5,9 +5,12 @@ import (
 	"errors"
 	"fmt"
 
-	"github.com/ethereum-optimism/optimism/op-preimage"
+	preimage "github.com/ethereum-optimism/optimism/op-preimage"
+	"github.com/ethereum/go-ethereum/common"
 )
 
+type LocalContext common.Hash
+
 type StepWitness struct {
 	// encoded state witness
 	State []byte
@@ -25,29 +28,35 @@ func uint64ToBytes32(v uint64) []byte {
 	return out[:]
 }
 
-func (wit *StepWitness) EncodeStepInput() []byte {
+func (wit *StepWitness) EncodeStepInput(localContext LocalContext) []byte {
 	abiStatePadding := (32 - (uint64(len(wit.State)) % 32)) % 32
+	abiProofPadding := (32 - (uint64(len(wit.MemProof)) % 32)) % 32
 
 	var input []byte
 	input = append(input, StepBytes4...)
-	input = append(input, uint64ToBytes32(32*2)...)                                           // state data offset in bytes
-	input = append(input, uint64ToBytes32(32*2+32+uint64(len(wit.State))+abiStatePadding)...) // proof data offset in bytes
-	// TODO pad state data to multiple of 32 bytes
-	// TODO also pad proof data
-
-	input = append(input, uint64ToBytes32(uint64(len(wit.State)))...) // state data length in bytes
+	// state data offset in bytes
+	input = append(input, uint64ToBytes32(32*3)...)
+	// proof data offset in bytes
+	input = append(input, uint64ToBytes32(32*3+32+uint64(len(wit.State))+abiStatePadding)...)
+	// local context in bytes
+	input = append(input, common.Hash(localContext).Bytes()...)
+	
+	// state data length in bytes
+	input = append(input, uint64ToBytes32(uint64(len(wit.State)))...)
 	input = append(input, wit.State[:]...)
 	input = append(input, make([]byte, abiStatePadding)...)
-	input = append(input, uint64ToBytes32(uint64(len(wit.MemProof)))...) // proof data length in bytes
+	// proof data length in bytes
+	input = append(input, uint64ToBytes32(uint64(len(wit.MemProof)))...)
 	input = append(input, wit.MemProof[:]...)
+	input = append(input, make([]byte, abiProofPadding)...)
 	return input
 }
 
 func (wit *StepWitness) HasPreimage() bool {
 	return wit.PreimageKey != ([32]byte{})
 }
 
-func (wit *StepWitness) EncodePreimageOracleInput() ([]byte, error) {
+func (wit *StepWitness) EncodePreimageOracleInput(localContext LocalContext) ([]byte, error) {
 	if wit.PreimageKey == ([32]byte{}) {
 		return nil, errors.New("cannot encode pre-image oracle input, witness has no pre-image to proof")
 	}
@@ -59,13 +68,14 @@ func (wit *StepWitness) EncodePreimageOracleInput() ([]byte, error) {
 		// In production usage there should be an on-chain contract that exposes this,
 		// rather than going through the global keccak256 oracle.
 		var input []byte
-		input = append(input, CheatBytes4...)
+		input = append(input, CheatLocalKeyBytes4...)
 		input = append(input, uint64ToBytes32(wit.PreimageOffset)...)
 		input = append(input, wit.PreimageKey[:]...)
 		var tmp [32]byte
 		copy(tmp[:], wit.PreimageValue[wit.PreimageOffset:])
 		input = append(input, tmp[:]...)
 		input = append(input, uint64ToBytes32(uint64(len(wit.PreimageValue))-8)...)
+		input = append(input, common.Hash(localContext).Bytes()...)
 		// Note: we can pad calldata to 32 byte multiple, but don't strictly have to
 		return input, nil
 	case preimage.Keccak256KeyType:
diff --git a/rvgo/slow/vm.go b/rvgo/slow/vm.go
@@ -3,6 +3,7 @@ package slow
 import (
 	"encoding/binary"
 	"fmt"
+
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
 )
@@ -103,8 +104,8 @@ func Step(calldata []byte, po PreimageOracle) (stateHash common.Hash, outErr err
 
 	// TODO: validate abi offset values?
 
-	stateContentOffset := uint8(4 + 32 + 32 + 32)
-	if iszero(eq(b32asBEWord(calldataload(toU64(4+32*2))), shortToU256(stateSize))) {
+	stateContentOffset := uint8(4 + 32 + 32 + 32 + 32)
+	if iszero(eq(b32asBEWord(calldataload(toU64(4+32*3))), shortToU256(stateSize))) {
 		// user-provided state size must match expected state size
 		panic("invalid state size input")
 	}
diff --git a/rvgo/test/evm_test.go b/rvgo/test/evm_test.go
@@ -149,13 +149,13 @@ func stepEVM(t *testing.T, env *vm.EVM, wit *fast.StepWitness, addrs *Addresses,
 	snap := env.StateDB.Snapshot()
 
 	if wit.HasPreimage() {
-		input, err := wit.EncodePreimageOracleInput()
+		input, err := wit.EncodePreimageOracleInput(fast.LocalContext{})
 		require.NoError(t, err)
 		ret, leftOverGas, err := env.Call(vm.AccountRef(addrs.Sender), addrs.Oracle, input, startingGas, big.NewInt(0))
 		require.NoError(t, err, "evm must not fail (ret: %x, gas: %d)", ret, startingGas-leftOverGas)
 	}
 
-	input := wit.EncodeStepInput()
+	input := wit.EncodeStepInput(fast.LocalContext{})
 
 	ret, leftOverGas, err := env.Call(vm.AccountRef(addrs.Sender), addrs.RISCV, input, startingGas, big.NewInt(0))
 	require.NoError(t, err, "evm must not fail (ret: %x), at step %d", ret, step)
diff --git a/rvgo/test/vm_go_test.go b/rvgo/test/vm_go_test.go
@@ -108,7 +108,7 @@ func fullTest(t *testing.T, vmState *fast.VMState, po *testOracle, symbols fast.
 			require.NoError(t, err)
 
 			if runSlow {
-				slowPostHash, err := slow.Step(wit.EncodeStepInput(), po)
+				slowPostHash, err := slow.Step(wit.EncodeStepInput(fast.LocalContext{}), po)
 				require.NoErrorf(t, err, "slow VM err at step %d, PC %08x: %v", i, vmState.PC, err)
 				require.Equal(t, fastStateHash, slowPostHash, "fast post-state must match slow post-state")
 			}
diff --git a/rvgo/test/vm_test.go b/rvgo/test/vm_test.go
@@ -75,7 +75,7 @@ func runSlowTestSuite(t *testing.T, path string) {
 		require.NoError(t, err)
 
 		// Now run the same in slow mode
-		input := wit.EncodeStepInput()
+		input := wit.EncodeStepInput(fast.LocalContext{})
 		post, err := slow.Step(input, nil)
 		require.NoErrorf(t, err, "slow VM err at step %d, PC %08x: %v", i, vmState.PC, err)
 
diff --git a/rvsol/src/PreimageOracle.sol b/rvsol/src/PreimageOracle.sol
@@ -32,6 +32,33 @@ contract PreimageOracle {
         preimageLengths[key] = size;
     }
 
+    // temporary method for localization. Will be removed to PreimageKeyLib.sol
+    function localize(bytes32 _key, bytes32 _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, 0x60), not(shl(248, 0xFF))), shl(248, 1))
+            // Restore the free memory pointer.
+            mstore(0x40, ptr)
+        }
+    }
+
+    // temporary method for localization. Will be removed to PreimageKeyLib.sol
+    function cheatLocalKey(uint256 partOffset, bytes32 key, bytes32 part, uint256 size, bytes32 localContext) external {
+        // sanity check key is local key using prefix
+        require(uint8(key[0]) == 1, "must be used for local key");
+        
+        bytes32 localizedKey = localize(key, localContext);
+        preimagePartOk[localizedKey][partOffset] = true;
+        preimageParts[localizedKey][partOffset] = part;
+        preimageLengths[localizedKey] = size;
+    }
+
     // loadKeccak256PreimagePart prepares the pre-image to be read by keccak256 key,
     // starting at the given offset, up to 32 bytes (clipped at preimage length, if out of data).
     function loadKeccak256PreimagePart(uint256 _partOffset, bytes calldata _preimage) external {
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)
@@ -270,8 +270,8 @@ contract Step {
                 // expected memory check: no allocated memory (start after scratch + free-mem-ptr + zero slot = 0x80)
                 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(calldataload(sub(stateData.offset, 32)), stateSize())) {
@@ -283,11 +283,12 @@ contract Step {
                 out := add(v, padding)
             }
             if iszero(eq(proof.offset, add(add(stateData.offset, paddedLen(stateSize())), 32))) {
-                // 100+stateSize+padding+32 = expected proof offset
+                // 132+stateSize+padding+32 = expected proof offset
                 revert(0, 0)
             }
             function proofContentOffset() -> out { // since we can't reference proof.offset in functions, blame Yul
-                out := 516
+                // 132+362+(32-362%32)+32=548
+                out := 548
             }
             if iszero(eq(proof.offset, proofContentOffset())) {
                 revert(0, 0)
@@ -772,10 +773,28 @@ contract Step {
                 revertWithCode(0xbadf00d0)
             }
 
-            function readPreimageValue(addr, count) -> out {
+            function localize(preImageKey, localContext_) -> localizedKey {
+                // TODO: deduplicate definition of localize using lib
+                // 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, preImageKey)
+                mstore(0x20, caller())
+                mstore(0x40, localContext_)
+                // Localize the key with the above `localize` operation.
+                localizedKey := or(and(keccak256(0, 0x60), not(shl(248, 0xFF))), shl(248, 1))
+                // Restore the free memory pointer.
+                mstore(0x40, ptr)
+            }
+
+            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 +830,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 +891,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 +1294,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

Original file line number	Diff line number	Diff line change
`@@ -3,7 +3,8 @@ package fast`
`3`	`3`	`import "github.com/ethereum/go-ethereum/crypto"`
`4`	`4`
`5`	`5`	`var (`
`6`		`- StepBytes4 = crypto.Keccak256([]byte("step(bytes,bytes)"))[:4]`
	`6`	`+ StepBytes4 = crypto.Keccak256([]byte("step(bytes,bytes,bytes32)"))[:4]`
`7`	`7`	`CheatBytes4 = crypto.Keccak256([]byte("cheat(uint256,bytes32,bytes32,uint256)"))[:4]`
	`8`	`+ CheatLocalKeyBytes4 = crypto.Keccak256([]byte("cheatLocalKey(uint256,bytes32,bytes32,uint256,bytes32)"))[:4]`
`8`	`9`	`LoadKeccak256PreimagePartBytes4 = crypto.Keccak256([]byte("loadKeccak256PreimagePart(uint256,bytes)"))[:4]`
`9`	`10`	`)`
Original file line number	Diff line number	Diff line change
`@@ -108,7 +108,7 @@ func fullTest(t testing.T, vmState fast.VMState, po *testOracle, symbols fast.`
`108`	`108`	`require.NoError(t, err)`
`109`	`109`
`110`	`110`	`if runSlow {`
`111`		`- slowPostHash, err := slow.Step(wit.EncodeStepInput(), po)`
	`111`	`+ slowPostHash, err := slow.Step(wit.EncodeStepInput(fast.LocalContext{}), po)`
`112`	`112`	`require.NoErrorf(t, err, "slow VM err at step %d, PC %08x: %v", i, vmState.PC, err)`
`113`	`113`	`require.Equal(t, fastStateHash, slowPostHash, "fast post-state must match slow post-state")`
`114`	`114`	`}`