Merge pull request #5446 from ethereum-optimism/op-program-trie

op-program: trie list pre-image read/write util funcs

op-program/client/mpt/db.go

+package mpt
+
+import "github.com/ethereum/go-ethereum/ethdb"
+
+type Hooks struct {
+	Get    func(key []byte) []byte
+	Put    func(key []byte, value []byte)
+	Delete func(key []byte)
+}
+
+// DB implements the ethdb.Database to back the StateDB of Geth.
+type DB struct {
+	db Hooks
+}
+
+func (p *DB) Has(key []byte) (bool, error) {
+	panic("not supported")
+}
+
+func (p *DB) Get(key []byte) ([]byte, error) {
+	return p.db.Get(key), nil
+}
+
+func (p *DB) Put(key []byte, value []byte) error {
+	p.db.Put(key, value)
+	return nil
+}
+
+func (p DB) Delete(key []byte) error {
+	p.db.Delete(key)
+	return nil
+}
+
+func (p DB) Stat(property string) (string, error) {
+	panic("not supported")
+}
+
+func (p DB) NewBatch() ethdb.Batch {
+	panic("not supported")
+}
+
+func (p DB) NewBatchWithSize(size int) ethdb.Batch {
+	panic("not supported")
+}
+
+func (p DB) NewIterator(prefix []byte, start []byte) ethdb.Iterator {
+	panic("not supported")
+}
+
+func (p DB) Compact(start []byte, limit []byte) error {
+	return nil // no-op
+}
+
+func (p DB) NewSnapshot() (ethdb.Snapshot, error) {
+	panic("not supported")
+}
+
+func (p DB) Close() error {
+	return nil
+}
+
+// We implement the full ethdb.Database bloat because the StateDB takes this full interface,
+// even though it only uses the KeyValue subset.
+
+func (p *DB) HasAncient(kind string, number uint64) (bool, error) {
+	panic("not supported")
+}
+
+func (p *DB) Ancient(kind string, number uint64) ([]byte, error) {
+	panic("not supported")
+}
+
+func (p *DB) AncientRange(kind string, start, count, maxBytes uint64) ([][]byte, error) {
+	panic("not supported")
+}
+
+func (p *DB) Ancients() (uint64, error) {
+	panic("not supported")
+}
+
+func (p *DB) Tail() (uint64, error) {
+	panic("not supported")
+}
+
+func (p *DB) AncientSize(kind string) (uint64, error) {
+	panic("not supported")
+}
+
+func (p *DB) ReadAncients(fn func(ethdb.AncientReaderOp) error) (err error) {
+	panic("not supported")
+}
+
+func (p *DB) ModifyAncients(f func(ethdb.AncientWriteOp) error) (int64, error) {
+	panic("not supported")
+}
+
+func (p *DB) TruncateHead(n uint64) error {
+	panic("not supported")
+}
+
+func (p *DB) TruncateTail(n uint64) error {
+	panic("not supported")
+}
+
+func (p *DB) Sync() error {
+	panic("not supported")
+}
+
+func (p *DB) MigrateTable(s string, f func([]byte) ([]byte, error)) error {
+	panic("not supported")
+}
+
+func (p *DB) AncientDatadir() (string, error) {
+	panic("not supported")
+}
+
+var _ ethdb.KeyValueStore = (*DB)(nil)

op-program/client/mpt/trie.go

+package mpt
+
+import (
+	"bytes"
+	"fmt"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/common/hexutil"
+	"github.com/ethereum/go-ethereum/core/types"
+	"github.com/ethereum/go-ethereum/rlp"
+	"github.com/ethereum/go-ethereum/trie"
+)
+
+// ReadTrie takes a Merkle Patricia Trie (MPT) root of a "DerivableList", and a pre-image oracle getter,
+// and traverses the implied MPT to collect all raw leaf nodes in order, which are then returned.
+func ReadTrie(root common.Hash, getPreimage func(key common.Hash) []byte) []hexutil.Bytes {
+	odb := &DB{db: Hooks{
+		Get: func(key []byte) []byte {
+			if len(key) != 32 {
+				panic(fmt.Errorf("expected 32 byte key query, but got %d bytes: %x", len(key), key))
+			}
+			return getPreimage(*(*[32]byte)(key))
+		},
+		Put: func(key []byte, value []byte) {
+			panic("put not supported")
+		},
+		Delete: func(key []byte) {
+			panic("delete not supported")
+		},
+	}}
+
+	// trie.New backed with a trie.NodeReader and trie.Reader seems really promising
+	// for a simple node-fetching backend, but the interface is half-private,
+	// while we already have the full database code for doing the same thing.
+	// Maybe it's still worth a small diff in geth to expose it?
+	// Diff would be:
+	//
+	//      type Node = node
+	//
+	//      func DecodeNode(hash, buf []byte) (node, error) {
+	//      	return decodeNode(hash, buf)
+	//      }
+	//
+	// And then still some code here to implement the trie.NodeReader and trie.Reader
+	// interfaces to map to the getPreimageFunction.
+	//
+	// For now we just use the state DB trie approach.
+
+	tdb := trie.NewDatabase(odb)
+	tr, err := trie.New(trie.TrieID(root), tdb)
+	if err != nil {
+		panic(err)
+	}
+	iter := tr.NodeIterator(nil)
+
+	// With small lists the iterator seems to use 0x80 (RLP empty string, unlike the others)
+	// as key for item 0, causing it to come last.
+	// Let's just remember the keys, and reorder them in the canonical order, to ensure it is correct.
+	var values [][]byte
+	var keys []uint64
+	for iter.Next(true) {
+		if iter.Leaf() {
+			k := iter.LeafKey()
+			var x uint64
+			err := rlp.DecodeBytes(k, &x)
+			if err != nil {
+				panic(fmt.Errorf("invalid key: %w", err))
+			}
+			keys = append(keys, x)
+			values = append(values, iter.LeafBlob())
+		}
+	}
+	out := make([]hexutil.Bytes, len(values))
+	for i, x := range keys {
+		if x >= uint64(len(values)) {
+			panic(fmt.Errorf("bad key: %d", x))
+		}
+		if out[x] != nil {
+			panic(fmt.Errorf("duplicate key %d", x))
+		}
+		out[x] = values[i]
+	}
+	return out
+}
+
+type rawList []hexutil.Bytes
+
+func (r rawList) Len() int {
+	return len(r)
+}
+
+func (r rawList) EncodeIndex(i int, buf *bytes.Buffer) {
+	buf.Write(r[i])
+}
+
+var _ types.DerivableList = rawList(nil)
+
+type noResetHasher struct {
+	*trie.StackTrie
+}
+
+// Reset is intercepted and is no-op, because we want to retain the writing function when calling types.DeriveSha
+func (n noResetHasher) Reset() {}
+
+// WriteTrie takes a list of values, and merkleizes them as a "DerivableList":
+// a Merkle Patricia Trie (MPT) with values keyed by their RLP encoded index.
+// This merkleization matches that of transactions, receipts, and withdrawals lists in the block header
+// (at least up to the Shanghai L1 update).
+// This then returns the MPT root and a list of pre-images of the trie.
+// Note: empty values are illegal, and there may be less pre-images returned than values,
+// if any values are less than 32 bytes and fit into branch-node slots that way.
+func WriteTrie(values []hexutil.Bytes) (common.Hash, []hexutil.Bytes) {
+	var out []hexutil.Bytes
+	st := noResetHasher{trie.NewStackTrie(
+		func(owner common.Hash, path []byte, hash common.Hash, blob []byte) {
+			out = append(out, common.CopyBytes(blob)) // the stack hasher may mutate the blob bytes, so copy them.
+		})}
+	root := types.DeriveSha(rawList(values), st)
+	return root, out
+}

op-program/client/mpt/trie_test.go

+package mpt
+
+import (
+	"fmt"
+	"math/rand"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ethereum/go-ethereum/common"
+	"github.com/ethereum/go-ethereum/common/hexutil"
+	"github.com/ethereum/go-ethereum/crypto"
+)
+
+type trieCase struct {
+	name     string
+	elements []hexutil.Bytes
+}
+
+func (tc *trieCase) run(t *testing.T) {
+	root, preimages := WriteTrie(tc.elements)
+	byHash := make(map[common.Hash][]byte)
+	for _, v := range preimages {
+		k := crypto.Keccak256Hash(v)
+		byHash[k] = v
+	}
+	results := ReadTrie(root, func(key common.Hash) []byte {
+		v, ok := byHash[key]
+		if !ok {
+			panic(fmt.Errorf("missing key %s", key))
+		}
+		return v
+	})
+	require.Equal(t, len(tc.elements), len(results), "expected equal amount of values")
+	for i, result := range results {
+		// hex encoded for debugging readability
+		require.Equal(t, tc.elements[i].String(), result.String(),
+			"value %d does not match, expected equal value data", i)
+	}
+}
+
+func TestListTrieRoundtrip(t *testing.T) {
+	testCases := []trieCase{
+		{name: "empty list", elements: []hexutil.Bytes{}},
+		{name: "nil list", elements: nil},
+		{name: "simple", elements: []hexutil.Bytes{[]byte("hello"), []byte("world")}},
+	}
+	rng := rand.New(rand.NewSource(1234))
+	// add some randomized cases
+	for i := 0; i < 30; i++ {
+		n := rng.Intn(300)
+		elems := make([]hexutil.Bytes, n)
+		for i := range elems {
+			length := 1 + rng.Intn(300) // empty items not allowed
+			data := make([]byte, length)
+			rng.Read(data[:])
+			elems[i] = data
+		}
+		testCases = append(testCases, trieCase{name: fmt.Sprintf("rand_%d", i), elements: elems})
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, tc.run)
+	}
+}