trie: re-add stacktrie binary marshalling (OBS! format change)

trie/stacktrie.go

@@ -420,88 +420,3 @@ func (stack *StackTrie) Commit() (h common.Hash, err error) {
 	stack.writeFn(stack.owner, nil, h, st.val)
 	return h, nil
 }
-
-//// NewFromBinary initialises a serialized stacktrie with the given db.
-//func NewFromBinary(data []byte, writeFn NodeWriteFunc) (*StackTrie, error) {
-//	var st StackTrie
-//	if err := st.UnmarshalBinary(data); err != nil {
-//		return nil, err
-//	}
-//	// If a database is used, we need to recursively add it to every child
-//	if writeFn != nil {
-//		st.setWriter(writeFn)
-//	}
-//	return &st, nil
-//}
-//
-//// MarshalBinary implements encoding.BinaryMarshaler
-//func (st *StackTrie) MarshalBinary() (data []byte, err error) {
-//	var (
-//		b bytes.Buffer
-//		w = bufio.NewWriter(&b)
-//	)
-//	if err := gob.NewEncoder(w).Encode(struct {
-//		Owner    common.Hash
-//		NodeType uint8
-//		Val      []byte
-//		Key      []byte
-//	}{
-//		st.owner,
-//		st.nodeType,
-//		st.val,
-//		st.key,
-//	}); err != nil {
-//		return nil, err
-//	}
-//	for _, child := range st.children {
-//		if child == nil {
-//			w.WriteByte(0)
-//			continue
-//		}
-//		w.WriteByte(1)
-//		if childData, err := child.MarshalBinary(); err != nil {
-//			return nil, err
-//		} else {
-//			w.Write(childData)
-//		}
-//	}
-//	w.Flush()
-//	return b.Bytes(), nil
-//}
-//
-//// UnmarshalBinary implements encoding.BinaryUnmarshaler
-//func (st *StackTrie) UnmarshalBinary(data []byte) error {
-//	r := bytes.NewReader(data)
-//	return st.unmarshalBinary(r)
-//}
-//
-//func (st *StackTrie) unmarshalBinary(r io.Reader) error {
-//	var dec struct {
-//		Owner    common.Hash
-//		NodeType uint8
-//		Val      []byte
-//		Key      []byte
-//	}
-//	if err := gob.NewDecoder(r).Decode(&dec); err != nil {
-//		return err
-//	}
-//	st.owner = dec.Owner
-//	st.nodeType = dec.NodeType
-//	st.val = dec.Val
-//	st.key = dec.Key
-//
-//	var hasChild = make([]byte, 1)
-//	for i := range st.children {
-//		if _, err := r.Read(hasChild); err != nil {
-//			return err
-//		} else if hasChild[0] == 0 {
-//			continue
-//		}
-//		var child StackTrie
-//		if err := child.unmarshalBinary(r); err != nil {
-//			return err
-//		}
-//		st.children[i] = &child
-//	}
-//	return nil
-//}

trie/stacktrie_marshalling.go

@@ -0,0 +1,114 @@
+// Copyright 2023 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package trie
+
+import (
+	"bufio"
+	"bytes"
+	"encoding"
+	"encoding/gob"
+)
+
+var ( //Compile-time interface checks
+	_ = encoding.BinaryMarshaler((*StackTrie)(nil))
+	_ = encoding.BinaryUnmarshaler((*StackTrie)(nil))
+)
+
+// NewFromBinaryV2 initialises a serialized stacktrie with the given db.
+// OBS! Format was changed along with the name of this constructor.
+func NewFromBinaryV2(data []byte, writeFn NodeWriteFunc) (*StackTrie, error) {
+	stack := NewStackTrie(writeFn)
+	if err := stack.UnmarshalBinary(data); err != nil {
+		return nil, err
+	}
+	return stack, nil
+}
+
+// UnmarshalBinary implements encoding.BinaryMarshaler
+func (st *StackTrie) MarshalBinary() (data []byte, err error) {
+	var (
+		b bytes.Buffer
+		w = bufio.NewWriter(&b)
+	)
+	if err := gob.NewEncoder(w).Encode(st.owner); err != nil {
+		return nil, err
+	}
+	if err := st.root.marshalInto(w); err != nil {
+		return nil, err
+	}
+	w.Flush()
+	return b.Bytes(), nil
+}
+
+// UnmarshalBinary implements encoding.BinaryUnmarshaler
+func (stack *StackTrie) UnmarshalBinary(data []byte) error {
+	r := bytes.NewReader(data)
+	if err := gob.NewDecoder(r).Decode(&stack.owner); err != nil {
+		return err
+	}
+	if err := stack.root.unmarshalFrom(r); err != nil {
+		return err
+	}
+	return nil
+}
+
+type encodedNode struct {
+	NodeType uint8
+	Val      []byte
+	Key      []byte
+}
+
+func (st *stNode) marshalInto(w *bufio.Writer) (err error) {
+	if err := gob.NewEncoder(w).Encode(encodedNode{st.nodeType, st.val, st.key}); err != nil {
+		return err
+	}
+	for _, child := range st.children {
+		if child == nil {
+			w.WriteByte(0)
+			continue
+		}
+		w.WriteByte(1)
+		if err := child.marshalInto(w); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (st *stNode) unmarshalFrom(r *bytes.Reader) error {
+	var dec encodedNode
+	if err := gob.NewDecoder(r).Decode(&dec); err != nil {
+		return err
+	}
+	st.nodeType = dec.NodeType
+	st.val = dec.Val
+	st.key = dec.Key
+
+	for i := range st.children {
+		if b, err := r.ReadByte(); err != nil {
+			return err
+		} else if b == 0 {
+			continue
+		}
+		var child stNode
+		if err := child.unmarshalFrom(r); err != nil {
+			return err
+		}
+		st.children[i] = &child
+	}
+	return nil
+}

trie/stacktrie_test.go

@@ -379,45 +379,48 @@ func TestStacktrieNotModifyValues(t *testing.T) {
 
 // TestStacktrieSerialization tests that the stacktrie works well if we
 // serialize/unserialize it a lot
-//func TestStacktrieSerialization(t *testing.T) {
+func TestStacktrieSerialization(t *testing.T) {
-//	var (
+	var (
-//		st       = NewStackTrie(nil)
+		st       = NewStackTrieWithOwner(nil, common.Hash{0x12})
-//		nt       = NewEmpty(NewDatabase(rawdb.NewMemoryDatabase(), nil))
+		nt       = NewEmpty(NewDatabase(rawdb.NewMemoryDatabase(), nil))
-//		keyB     = big.NewInt(1)
+		keyB     = big.NewInt(1)
-//		keyDelta = big.NewInt(1)
+		keyDelta = big.NewInt(1)
-//		vals     [][]byte
+		vals     [][]byte
-//		keys     [][]byte
+		keys     [][]byte
-//	)
+	)
-//	getValue := func(i int) []byte {
+	getValue := func(i int) []byte {
-//		if i%2 == 0 { // large
+		if i%2 == 0 { // large
-//			return crypto.Keccak256(big.NewInt(int64(i)).Bytes())
+			return crypto.Keccak256(big.NewInt(int64(i)).Bytes())
-//		} else { //small
+		} else { //small
-//			return big.NewInt(int64(i)).Bytes()
+			return big.NewInt(int64(i)).Bytes()
-//		}
+		}
-//	}
+	}
-//	for i := 0; i < 10; i++ {
+	for i := 0; i < 10; i++ {
-//		vals = append(vals, getValue(i))
+		vals = append(vals, getValue(i))
-//		keys = append(keys, common.BigToHash(keyB).Bytes())
+		keys = append(keys, common.BigToHash(keyB).Bytes())
-//		keyB = keyB.Add(keyB, keyDelta)
+		keyB = keyB.Add(keyB, keyDelta)
-//		keyDelta.Add(keyDelta, common.Big1)
+		keyDelta.Add(keyDelta, common.Big1)
-//	}
+	}
-//	for i, k := range keys {
+	for i, k := range keys {
-//		nt.Update(k, common.CopyBytes(vals[i]))
+		nt.Update(k, common.CopyBytes(vals[i]))
-//	}
+	}
-//
+
-//	for i, k := range keys {
+	for i, k := range keys {
-//		blob, err := st.MarshalBinary()
+		blob, err := st.MarshalBinary()
-//		if err != nil {
+		if err != nil {
-//			t.Fatal(err)
+			t.Fatal(err)
-//		}
+		}
-//		newSt, err := NewFromBinary(blob, nil)
+		newSt, err := NewFromBinaryV2(blob, nil)
-//		if err != nil {
+		if err != nil {
-//			t.Fatal(err)
+			t.Fatal(err)
-//		}
+		}
-//		st = newSt
+		st = newSt
-//		st.Update(k, common.CopyBytes(vals[i]))
+		st.Update(k, common.CopyBytes(vals[i]))
-//	}
+	}
-//	if have, want := st.Hash(), nt.Hash(); have != want {
+	if have, want := st.Hash(), nt.Hash(); have != want {
-//		t.Fatalf("have %#x want %#x", have, want)
+		t.Fatalf("have %#x want %#x", have, want)
-//	}
+	}
-//}
+	if have, want := st.owner, (common.Hash{0x12}); have != want {
+		t.Fatalf("have %#x want %#x", have, want)
+	}
+}