88package blockstest
99
1010import (
11+ "slices"
12+ "sync"
1113 "testing"
1214
15+ "github.com/ava-labs/libevm/common"
1316 "github.com/ava-labs/libevm/core/types"
17+ "github.com/ava-labs/libevm/libevm/options"
1418
1519 "github.com/ava-labs/strevm/blocks"
1620)
1721
1822// A ChainBuilder builds a chain of blocks, maintaining necessary invariants.
1923type ChainBuilder struct {
20- chain []* blocks.Block
24+ chain []* blocks.Block
25+ blocksByHash sync.Map
26+
27+ defaultOpts []ChainOption
2128}
2229
2330// NewChainBuilder returns a new ChainBuilder starting from the provided block,
@@ -28,17 +35,78 @@ func NewChainBuilder(genesis *blocks.Block) *ChainBuilder {
2835 }
2936}
3037
38+ // A ChainOption configures [ChainBuilder.NewBlock].
39+ type ChainOption = options.Option [chainOptions ]
40+
41+ // SetDefaultOptions sets the default options upon which all
42+ // additional options passed to [ChainBuilder.NewBlock] are appended.
43+ func (cb * ChainBuilder ) SetDefaultOptions (opts ... ChainOption ) {
44+ cb .defaultOpts = opts
45+ }
46+
47+ type chainOptions struct {
48+ eth []EthBlockOption
49+ sae []BlockOption
50+ }
51+
52+ // WithEthBlockOptions wraps the options that [ChainBuilder.NewBlock] propagates
53+ // to [NewEthBlock].
54+ func WithEthBlockOptions (opts ... EthBlockOption ) ChainOption {
55+ return options.Func [chainOptions ](func (co * chainOptions ) {
56+ co .eth = append (co .eth , opts ... )
57+ })
58+ }
59+
60+ // WithBlockOptions wraps the options that [ChainBuilder.NewBlock] propagates to
61+ // [NewBlock].
62+ func WithBlockOptions (opts ... BlockOption ) ChainOption {
63+ return options.Func [chainOptions ](func (co * chainOptions ) {
64+ co .sae = append (co .sae , opts ... )
65+ })
66+ }
67+
3168// NewBlock constructs and returns a new block in the chain.
32- func (cb * ChainBuilder ) NewBlock (tb testing.TB , txs []* types.Transaction , opts ... EthBlockOption ) * blocks.Block {
69+ func (cb * ChainBuilder ) NewBlock (tb testing.TB , txs []* types.Transaction , opts ... ChainOption ) * blocks.Block {
3370 tb .Helper ()
71+
72+ allOpts := new (chainOptions )
73+ options .ApplyTo (allOpts , cb .defaultOpts ... )
74+ options .ApplyTo (allOpts , opts ... )
75+
3476 last := cb .Last ()
35- eth := NewEthBlock (last .EthBlock (), txs , opts ... )
36- cb .chain = append (cb .chain , NewBlock (tb , eth , last , nil )) // TODO(arr4n) support last-settled blocks
37- return cb .Last ()
77+ eth := NewEthBlock (last .EthBlock (), txs , allOpts .eth ... )
78+ b := NewBlock (tb , eth , last , nil , allOpts .sae ... ) // TODO(arr4n) support last-settled blocks
79+ cb .chain = append (cb .chain , b )
80+ cb .blocksByHash .Store (b .Hash (), b )
81+
82+ return b
3883}
3984
4085// Last returns the last block to be built by the builder, which MAY be the
4186// genesis block passed to the constructor.
4287func (cb * ChainBuilder ) Last () * blocks.Block {
4388 return cb .chain [len (cb .chain )- 1 ]
4489}
90+
91+ // AllBlocks returns all blocks, including the genesis passed to
92+ // [NewChainBuilder].
93+ func (cb * ChainBuilder ) AllBlocks () []* blocks.Block {
94+ return slices .Clone (cb .chain )
95+ }
96+
97+ // AllExceptGenesis returns all blocks created with [ChainBuilder.NewBlock].
98+ func (cb * ChainBuilder ) AllExceptGenesis () []* blocks.Block {
99+ return slices .Clone (cb .chain [1 :])
100+ }
101+
102+ // GetBlock returns the block with specified hash and height, and a flag
103+ // indicating if it was found. If either argument does not match, it returns
104+ // `nil, false`.
105+ func (cb * ChainBuilder ) GetBlock (h common.Hash , num uint64 ) (* blocks.Block , bool ) {
106+ ifc , _ := cb .blocksByHash .Load (h )
107+ b , ok := ifc .(* blocks.Block )
108+ if ! ok || b .NumberU64 () != num {
109+ return nil , false
110+ }
111+ return b , true
112+ }
0 commit comments