feat: before- and after-block hooks with additional arguments

Author: ARR4N

libevm/precompiles/parallel/parallel.go

@@ -29,6 +29,7 @@ import (
 	"github.com/ava-labs/libevm/core/types"
 	"github.com/ava-labs/libevm/core/vm"
 	"github.com/ava-labs/libevm/libevm"
+	"github.com/ava-labs/libevm/libevm/stateconf"
 	"github.com/ava-labs/libevm/params"
 )
 
@@ -43,10 +44,22 @@ import (
 //
 // Scenario (2) allows precompile access to be determined through inspection of
 // the [types.Transaction] alone, without the need for execution.
+//
+// All [libevm.StateReader] instances are opened to the state at the beginning
+// of the block. The [StateDB] is the same one used to execute the block,
+// before being committed, and MAY be written to.
 type Handler[Result any] interface {
-	BeforeBlock(*types.Header)
+	BeforeBlock(libevm.StateReader, *types.Block)
 	Gas(*types.Transaction) (gas uint64, process bool)
 	Process(sdb libevm.StateReader, index int, tx *types.Transaction) Result
+	AfterBlock(StateDB, *types.Block, types.Receipts)
+}
+
+// StateDB is the subset of [state.StateDB] methods that MAY be called by
+// [Handler.AfterBlock].
+type StateDB interface {
+	libevm.StateReader
+	SetState(_ common.Address, key, val common.Hash, _ ...stateconf.StateDBStateOption)
 }
 
 // A Processor orchestrates dispatch and collection of results from a [Handler].
@@ -162,9 +175,19 @@ func (p *Processor[R]) Close() {
 // StartBlock dispatches transactions to the [Handler] and returns immediately.
 // It MUST be paired with a call to [Processor.FinishBlock], without overlap of
 // blocks.
-func (p *Processor[R]) StartBlock(b *types.Block, rules params.Rules, sdb *state.StateDB) error {
+func (p *Processor[R]) StartBlock(sdb *state.StateDB, rules params.Rules, b *types.Block) error {
+	// The distribution mechanism copies the StateDB so we don't need to do it
+	// here, but the [Handler] is called directly so we do copy.
 	p.stateShare.distribute(sdb)
-	p.handler.BeforeBlock(types.CopyHeader(b.Header()))
+	p.handler.BeforeBlock(
+		sdb.Copy(),
+		types.NewBlockWithHeader(
+			b.Header(),
+		).WithBody(
+			*b.Body(),
+		),
+	)
+
 	txs := b.Transactions()
 	jobs := make([]*job, 0, len(txs))
 
@@ -206,14 +229,15 @@ func (p *Processor[R]) StartBlock(b *types.Block, rules params.Rules, sdb *state
 // FinishBlock returns the [Processor] to a state ready for the next block. A
 // return from FinishBlock guarantees that all dispatched work from the
 // respective call to [Processor.StartBlock] has been completed.
-func (p *Processor[R]) FinishBlock(b *types.Block) {
+func (p *Processor[R]) FinishBlock(sdb vm.StateDB, b *types.Block, rs types.Receipts) {
 	for i := range len(b.Transactions()) {
 		// Every result channel is guaranteed to have some value in its buffer
 		// because [Processor.BeforeBlock] either sends a nil *R or it
 		// dispatches a job, which will send a non-nil *R.
 		tx := (<-p.results[i]).tx
 		delete(p.txGas, tx)
 	}
+	p.handler.AfterBlock(sdb, b, rs)
 }
 
 // Result blocks until the i'th transaction passed to [Processor.StartBlock] has

libevm/precompiles/parallel/parallel_test.go

@@ -54,8 +54,8 @@ type concat struct {
 	gas         uint64
 }
 
-func (c *concat) BeforeBlock(h *types.Header) {
-	c.headerExtra = slices.Clone(h.Extra)
+func (c *concat) BeforeBlock(_ libevm.StateReader, b *types.Block) {
+	c.headerExtra = slices.Clone(b.Header().Extra)
 }
 
 func (c *concat) Gas(tx *types.Transaction) (uint64, bool) {
@@ -77,6 +77,8 @@ func (c *concat) Process(sdb libevm.StateReader, i int, tx *types.Transaction) [
 	)
 }
 
+func (*concat) AfterBlock(StateDB, *types.Block, types.Receipts) {}
+
 func TestProcessor(t *testing.T) {
 	handler := &concat{
 		addr:     common.Address{'c', 'o', 'n', 'c', 'a', 't'},
@@ -173,8 +175,8 @@ func TestProcessor(t *testing.T) {
 
 			extra := []byte("extra")
 			block := types.NewBlock(&types.Header{Extra: extra}, txs, nil, nil, trie.NewStackTrie(nil))
-			require.NoError(t, p.StartBlock(block, rules, sdb), "StartBlock()")
-			defer p.FinishBlock(block)
+			require.NoError(t, p.StartBlock(sdb, rules, block), "StartBlock()")
+			defer p.FinishBlock(sdb, block, nil)
 
 			for i, tx := range txs {
 				wantOK := wantProcessed[i]
@@ -305,8 +307,7 @@ func TestIntegration(t *testing.T) {
 	}
 
 	block := types.NewBlock(header, txs, nil, nil, trie.NewStackTrie(nil))
-	require.NoError(t, sut.StartBlock(block, rules, state), "StartBlock()")
-	defer sut.FinishBlock(block)
+	require.NoError(t, sut.StartBlock(state, rules, block), "StartBlock()")
 
 	pool := core.GasPool(math.MaxUint64)
 	var got []*types.Receipt
@@ -332,4 +333,5 @@ func TestIntegration(t *testing.T) {
 	if diff := cmp.Diff(want, got, ignore); diff != "" {
 		t.Errorf("%T diff (-want +got):\n%s", got, diff)
 	}
+	sut.FinishBlock(state, block, got)
 }