Skip to content

Commit 999f2c0

Browse files
ARR4NARR4N
committed
feat: trienode payloads for Node, NodeSet, and MergedNodeSet 
1 parent 414b1f5 commit 999f2c0

File tree

3 files changed

+247
-1
lines changed

3 files changed

+247
-1
lines changed

‎trie/trienode/node.go‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"strings"
2323

2424
"github.com/ava-labs/libevm/common"
25+
"github.com/ava-labs/libevm/libevm/pseudo"
2526
)
2627

2728
// Node is a wrapper which contains the encoded blob of the trie node and its
@@ -30,6 +31,8 @@ import (
3031
type Node struct {
3132
Hash common.Hash // Node hash, empty for deleted node
3233
Blob []byte // Encoded node blob, nil for the deleted node
34+
35+
extra *pseudo.Type // libevm
3336
}
3437

3538
// Size returns the total memory size used by this node.
@@ -64,6 +67,8 @@ type NodeSet struct {
6467
Nodes map[string]*Node
6568
updates int // the count of updated and inserted nodes
6669
deletes int // the count of deleted nodes
70+
71+
extra *pseudo.Type // libevm
6772
}
6873

6974
// NewNodeSet initializes a node set. The owner is zero for the account trie and
@@ -97,6 +102,7 @@ func (set *NodeSet) AddNode(path []byte, n *Node) {
97102
set.updates += 1
98103
}
99104
set.Nodes[string(path)] = n
105+
set.mergePayload(path, n)
100106
}
101107

102108
// Merge adds a set of nodes into the set.
@@ -164,6 +170,8 @@ func (set *NodeSet) Summary() string {
164170
// MergedNodeSet represents a merged node set for a group of tries.
165171
type MergedNodeSet struct {
166172
Sets map[common.Hash]*NodeSet
173+
174+
extra *pseudo.Type // libevm
167175
}
168176

169177
// NewMergedNodeSet initializes an empty merged set.
@@ -180,7 +188,7 @@ func NewWithNodeSet(set *NodeSet) *MergedNodeSet {
180188

181189
// Merge merges the provided dirty nodes of a trie into the set. The assumption
182190
// is held that no duplicated set belonging to the same trie will be merged twice.
183-
func (set *MergedNodeSet) Merge(other *NodeSet) error {
191+
func (set *MergedNodeSet) merge(other *NodeSet) error {
184192
subset, present := set.Sets[other.Owner]
185193
if present {
186194
return subset.Merge(other.Owner, other.Nodes)

‎trie/trienode/node.libevm.go‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
// Copyright 2025 the libevm authors.
2+
//
3+
// The libevm additions to go-ethereum are free software: you can redistribute
4+
// them and/or modify them under the terms of the GNU Lesser General Public License
5+
// as published by the Free Software Foundation, either version 3 of the License,
6+
// or (at your option) any later version.
7+
//
8+
// The libevm additions are distributed in the hope that they will be useful,
9+
// but WITHOUT ANY WARRANTY; without even the implied warranty of
10+
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser
11+
// General Public License for more details.
12+
//
13+
// You should have received a copy of the GNU Lesser General Public License
14+
// along with the go-ethereum library. If not, see
15+
// <http://www.gnu.org/licenses/>.
16+
17+
package trienode
18+
19+
import (
20+
"github.com/ava-labs/libevm/libevm/pseudo"
21+
"github.com/ava-labs/libevm/libevm/register"
22+
)
23+
24+
// MergedNodeSetHooks
25+
type MergedNodeSetHooks interface {
26+
Merge(into *MergedNodeSet, _ *NodeSet) error
27+
}
28+
29+
// NodeSetHooks
30+
type NodeSetHooks interface {
31+
Add(into *NodeSet, path []byte, _ *Node)
32+
}
33+
34+
// RegisterExtras
35+
func RegisterExtras[
36+
MNS, NS, N any,
37+
MNSPtr interface {
38+
MergedNodeSetHooks
39+
*MNS
40+
},
41+
NSPtr interface {
42+
NodeSetHooks
43+
*NS
44+
},
45+
]() ExtraPayloads[MNSPtr, NSPtr, N] {
46+
payloads := ExtraPayloads[MNSPtr, NSPtr, N]{
47+
MergedNodeSet: pseudo.NewAccessor[*MergedNodeSet, MNSPtr](
48+
(*MergedNodeSet).extraPayload,
49+
func(s *MergedNodeSet, t *pseudo.Type) { s.extra = t },
50+
),
51+
NodeSet: pseudo.NewAccessor[*NodeSet, NSPtr](
52+
(*NodeSet).extraPayload,
53+
func(s *NodeSet, t *pseudo.Type) { s.extra = t },
54+
),
55+
Node: pseudo.NewAccessor[*Node, N](
56+
(*Node).extraPayload,
57+
func(n *Node, t *pseudo.Type) { n.extra = t },
58+
),
59+
}
60+
61+
registeredExtras.MustRegister(&extraConstructors{
62+
newMergedNodeSet: pseudo.NewConstructor[MNS]().NewPointer,
63+
newNodeSet: pseudo.NewConstructor[NS]().NewPointer,
64+
newNode: pseudo.NewConstructor[N]().Zero,
65+
hooks: payloads,
66+
})
67+
68+
return payloads
69+
}
70+
71+
// TestOnlyClearRegisteredExtras
72+
func TestOnlyClearRegisteredExtras() {
73+
registeredExtras.TestOnlyClear()
74+
}
75+
76+
var registeredExtras register.AtMostOnce[*extraConstructors]
77+
78+
type extraConstructors struct {
79+
newMergedNodeSet func() *pseudo.Type
80+
newNodeSet func() *pseudo.Type
81+
newNode func() *pseudo.Type
82+
hooks interface {
83+
hooksFromMNS(*MergedNodeSet) MergedNodeSetHooks
84+
hooksFromNS(*NodeSet) NodeSetHooks
85+
}
86+
}
87+
88+
// Merge merges the provided dirty nodes of a trie into the set. The assumption
89+
// is held that no duplicated set belonging to the same trie will be merged twice.
90+
func (set *MergedNodeSet) Merge(other *NodeSet) error {
91+
if err := set.merge(other); err != nil {
92+
return err
93+
}
94+
if r := registeredExtras; r.Registered() {
95+
return r.Get().hooks.hooksFromMNS(set).Merge(set, other)
96+
}
97+
return nil
98+
}
99+
100+
func (set *NodeSet) mergePayload(path []byte, n *Node) {
101+
if r := registeredExtras; r.Registered() {
102+
r.Get().hooks.hooksFromNS(set).Add(set, path, n)
103+
}
104+
}
105+
106+
// ExtraPayloads
107+
type ExtraPayloads[
108+
MNS MergedNodeSetHooks,
109+
NS NodeSetHooks,
110+
N any,
111+
] struct {
112+
MergedNodeSet pseudo.Accessor[*MergedNodeSet, MNS]
113+
NodeSet pseudo.Accessor[*NodeSet, NS]
114+
Node pseudo.Accessor[*Node, N]
115+
}
116+
117+
func (e ExtraPayloads[MNS, NS, N]) hooksFromMNS(s *MergedNodeSet) MergedNodeSetHooks {
118+
return e.MergedNodeSet.Get(s)
119+
}
120+
121+
func (e ExtraPayloads[MNS, NS, N]) hooksFromNS(s *NodeSet) NodeSetHooks {
122+
return e.NodeSet.Get(s)
123+
}
124+
125+
func extraPayloadOrSetDefault(field **pseudo.Type, construct func(*extraConstructors) *pseudo.Type) *pseudo.Type {
126+
r := registeredExtras
127+
if !r.Registered() {
128+
// See params.ChainConfig.extraPayload() for panic rationale.
129+
panic("<T>.extraPayload() called before RegisterExtras()")
130+
}
131+
if *field == nil {
132+
*field = construct(r.Get())
133+
}
134+
return *field
135+
}
136+
137+
func (set *MergedNodeSet) extraPayload() *pseudo.Type {
138+
return extraPayloadOrSetDefault(&set.extra, func(c *extraConstructors) *pseudo.Type {
139+
return c.newMergedNodeSet()
140+
})
141+
}
142+
143+
func (set *NodeSet) extraPayload() *pseudo.Type {
144+
return extraPayloadOrSetDefault(&set.extra, func(c *extraConstructors) *pseudo.Type {
145+
return c.newNodeSet()
146+
})
147+
}
148+
149+
func (n *Node) extraPayload() *pseudo.Type {
150+
return extraPayloadOrSetDefault(&n.extra, func(c *extraConstructors) *pseudo.Type {
151+
return c.newNode()
152+
})
153+
}

‎trie/trienode/node.libevm_test.go‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
// Copyright 2025 the libevm authors.
2+
//
3+
// The libevm additions to go-ethereum are free software: you can redistribute
4+
// them and/or modify them under the terms of the GNU Lesser General Public License
5+
// as published by the Free Software Foundation, either version 3 of the License,
6+
// or (at your option) any later version.
7+
//
8+
// The libevm additions are distributed in the hope that they will be useful,
9+
// but WITHOUT ANY WARRANTY; without even the implied warranty of
10+
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser
11+
// General Public License for more details.
12+
//
13+
// You should have received a copy of the GNU Lesser General Public License
14+
// along with the go-ethereum library. If not, see
15+
// <http://www.gnu.org/licenses/>.
16+
17+
package trienode
18+
19+
import (
20+
"maps"
21+
"testing"
22+
23+
"github.com/ava-labs/libevm/common"
24+
"github.com/google/go-cmp/cmp"
25+
"github.com/stretchr/testify/require"
26+
)
27+
28+
type nodePayload struct {
29+
x uint64
30+
}
31+
32+
type setPayload struct {
33+
added map[string]uint64
34+
}
35+
36+
func (p *setPayload) Add(_ *NodeSet, path []byte, n *Node) {
37+
if p.added == nil {
38+
p.added = make(map[string]uint64)
39+
}
40+
p.added[string(path)] = extras.Node.Get(n).x
41+
}
42+
43+
type mergedSetPayload struct {
44+
merged []map[string]uint64
45+
}
46+
47+
func (p *mergedSetPayload) Merge(_ *MergedNodeSet, ns *NodeSet) error {
48+
p.merged = append(p.merged, maps.Clone(extras.NodeSet.Get(ns).added))
49+
return nil
50+
}
51+
52+
var extras ExtraPayloads[*mergedSetPayload, *setPayload, nodePayload]
53+
54+
func TestExtras(t *testing.T) {
55+
extras = RegisterExtras[mergedSetPayload, setPayload, nodePayload]()
56+
t.Cleanup(TestOnlyClearRegisteredExtras)
57+
58+
n1 := New(common.Hash{0}, nil)
59+
extras.Node.Set(n1, nodePayload{x: 1})
60+
n42 := New(common.Hash{1}, nil)
61+
extras.Node.Set(n42, nodePayload{x: 42})
62+
63+
set := NewNodeSet(common.Hash{})
64+
merge := NewMergedNodeSet()
65+
set.AddNode([]byte("n1"), n1)
66+
require.NoError(t, merge.Merge(set))
67+
68+
set.AddNode([]byte("n42"), n42)
69+
require.NoError(t, merge.Merge(set))
70+
71+
got := extras.MergedNodeSet.Get(merge).merged
72+
want := []map[string]uint64{
73+
{
74+
"n1": 1,
75+
},
76+
{
77+
"n1": 1,
78+
"n42": 42,
79+
},
80+
}
81+
82+
if diff := cmp.Diff(want, got); diff != "" {
83+
t.Errorf("%T payload diff (-want +got):\n%s", merge, diff)
84+
}
85+
}

0 commit comments

Comments
 (0)