rvgo:Add stopAtPreimage flag

rvgo/cmd/run.go

@@ -3,6 +3,9 @@ package cmd
 import (
 	"fmt"
 	"os"
+	"slices"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
@@ -70,24 +73,41 @@ func Run(ctx *cli.Context) error {
 	errLog := &LoggingWriter{Name: "program std-err", Log: l}
 
 	stopAtAnyPreimage := false
-	var stopAtPreimageTypeByte preimage.KeyType
-	switch ctx.String(cannon.RunStopAtPreimageTypeFlag.Name) {
-	case "local":
-		stopAtPreimageTypeByte = preimage.LocalKeyType
-	case "keccak":
-		stopAtPreimageTypeByte = preimage.Keccak256KeyType
-	case "sha256":
-		stopAtPreimageTypeByte = preimage.Sha256KeyType
-	case "blob":
-		stopAtPreimageTypeByte = preimage.BlobKeyType
-	case "precompile":
-		stopAtPreimageTypeByte = preimage.PrecompileKeyType
-	case "any":
-		stopAtAnyPreimage = true
-	case "":
-		// 0 preimage type is forbidden so will not stop at any preimage
-	default:
-		return fmt.Errorf("invalid preimage type %q", ctx.String(cannon.RunStopAtPreimageTypeFlag.Name))
+	var stopAtPreimageKeyPrefix []byte
+	stopAtPreimageOffset := uint64(0)
+	if ctx.IsSet(cannon.RunStopAtPreimageFlag.Name) {
+		val := ctx.String(cannon.RunStopAtPreimageFlag.Name)
+		parts := strings.Split(val, "@")
+		if len(parts) > 2 {
+			return fmt.Errorf("invalid %v: %v", cannon.RunStopAtPreimageFlag.Name, val)
+		}
+		stopAtPreimageKeyPrefix = common.FromHex(parts[0])
+		if len(parts) == 2 {
+			x, err := strconv.ParseUint(parts[1], 10, 64)
+			if err != nil {
+				return fmt.Errorf("invalid preimage offset: %w", err)
+			}
+			stopAtPreimageOffset = x
+		}
+	} else {
+		switch ctx.String(cannon.RunStopAtPreimageTypeFlag.Name) {
+		case "local":
+			stopAtPreimageKeyPrefix = []byte{byte(preimage.LocalKeyType)}
+		case "keccak":
+			stopAtPreimageKeyPrefix = []byte{byte(preimage.Keccak256KeyType)}
+		case "sha256":
+			stopAtPreimageKeyPrefix = []byte{byte(preimage.Sha256KeyType)}
+		case "blob":
+			stopAtPreimageKeyPrefix = []byte{byte(preimage.BlobKeyType)}
+		case "precompile":
+			stopAtPreimageKeyPrefix = []byte{byte(preimage.PrecompileKeyType)}
+		case "any":
+			stopAtAnyPreimage = true
+		case "":
+			// 0 preimage type is forbidden so will not stop at any preimage
+		default:
+			return fmt.Errorf("invalid preimage type %q", ctx.String(cannon.RunStopAtPreimageTypeFlag.Name))
+		}
 	}
 	stopAtPreimageLargerThan := ctx.Int(cannon.RunStopAtPreimageLargerThanFlag.Name)
 
@@ -177,8 +197,6 @@ func Run(ctx *cli.Context) error {
 			}
 		}
 
-		prevPreimageOffset := state.PreimageOffset
-
 		if proofAt(state) {
 			preStateHash, err := state.EncodeWitness().StateHash()
 			if err != nil {
@@ -214,14 +232,21 @@ func Run(ctx *cli.Context) error {
 			}
 		}
 
-		if preimageRead := state.PreimageOffset > prevPreimageOffset; preimageRead {
+		lastPreimageKey, lastPreimageValue, lastPreimageOffset := us.LastPreimage()
+		if lastPreimageOffset != ^uint64(0) {
 			if stopAtAnyPreimage {
+				l.Info("Stopping at preimage read")
 				break
 			}
-			if state.PreimageKey[0] == byte(stopAtPreimageTypeByte) {
-				break
+			if len(stopAtPreimageKeyPrefix) > 0 &&
+				slices.Equal(lastPreimageKey[:len(stopAtPreimageKeyPrefix)], stopAtPreimageKeyPrefix) {
+				if stopAtPreimageOffset == lastPreimageOffset {
+					l.Info("Stopping at preimage read", "keyPrefix", common.Bytes2Hex(stopAtPreimageKeyPrefix), "offset", lastPreimageOffset)
+					break
+				}
 			}
-			if stopAtPreimageLargerThan != 0 && len(us.LastPreimage()) > stopAtPreimageLargerThan {
+			if stopAtPreimageLargerThan != 0 && len(lastPreimageValue) > stopAtPreimageLargerThan {
+				l.Info("Stopping at preimage read", "size", len(lastPreimageValue), "min", stopAtPreimageLargerThan)
 				break
 			}
 		}
@@ -250,6 +275,7 @@ var RunCommand = &cli.Command{
 		cannon.RunSnapshotAtFlag,
 		cannon.RunSnapshotFmtFlag,
 		cannon.RunStopAtFlag,
+		cannon.RunStopAtPreimageFlag,
 		cannon.RunStopAtPreimageTypeFlag,
 		cannon.RunStopAtPreimageLargerThanFlag,
 		cannon.RunMetaFlag,

rvgo/fast/instrumented.go

@@ -116,6 +116,6 @@ func (m *InstrumentedState) verifyMemChange(effAddr uint64, proofIndex uint8) {
 	}
 }
 
-func (m *InstrumentedState) LastPreimage() []byte {
-	return m.lastPreimage
+func (m *InstrumentedState) LastPreimage() ([32]byte, []byte, uint64) {
+	return m.lastPreimageKey, m.lastPreimage, m.lastPreimageOffset
 }