Add AVM state.Chain interface (ava-labs#2411)

Co-authored-by: Chloe <[email protected]>

vms/avm/blocks/parser.go

@@ -5,8 +5,11 @@ package blocks
 
 import (
 	"fmt"
+	"reflect"
 
 	"github.com/ava-labs/avalanchego/codec"
+	"github.com/ava-labs/avalanchego/utils/logging"
+	"github.com/ava-labs/avalanchego/utils/timer/mockable"
 	"github.com/ava-labs/avalanchego/utils/wrappers"
 	"github.com/ava-labs/avalanchego/vms/avm/fxs"
 	"github.com/ava-labs/avalanchego/vms/avm/txs"
@@ -49,6 +52,29 @@ func NewParser(fxs []fxs.Fx) (Parser, error) {
 	}, errs.Err
 }
 
+func NewCustomParser(
+	typeToFxIndex map[reflect.Type]int,
+	clock *mockable.Clock,
+	log logging.Logger,
+	fxs []fxs.Fx,
+) (Parser, error) {
+	p, err := txs.NewCustomParser(typeToFxIndex, clock, log, fxs)
+	if err != nil {
+		return nil, err
+	}
+	c := p.CodecRegistry()
+	gc := p.GenesisCodecRegistry()
+
+	errs := wrappers.Errs{}
+	errs.Add(
+		c.RegisterType(&StandardBlock{}),
+		gc.RegisterType(&StandardBlock{}),
+	)
+	return &parser{
+		Parser: p,
+	}, errs.Err
+}
+
 func (p *parser) ParseBlock(bytes []byte) (Block, error) {
 	return parse(p.Codec(), bytes)
 }

vms/avm/index_test.go

@@ -76,9 +76,7 @@ func TestIndexTransaction_Ordered(t *testing.T) {
 		utxo := buildPlatformUTXO(utxoID, txAssetID, addr)
 
 		// save utxo to state
-		if err := vm.state.PutUTXO(utxo); err != nil {
-			t.Fatal("Error saving utxo", err)
-		}
+		vm.state.AddUTXO(utxo)
 
 		// issue transaction
 		if _, err := vm.IssueTx(tx.Bytes()); err != nil {
@@ -168,9 +166,7 @@ func TestIndexTransaction_MultipleTransactions(t *testing.T) {
 		utxo := buildPlatformUTXO(utxoID, txAssetID, addr)
 
 		// save utxo to state
-		if err := vm.state.PutUTXO(utxo); err != nil {
-			t.Fatal("Error saving utxo", err)
-		}
+		vm.state.AddUTXO(utxo)
 
 		// issue transaction
 		if _, err := vm.IssueTx(tx.Bytes()); err != nil {
@@ -264,9 +260,7 @@ func TestIndexTransaction_MultipleAddresses(t *testing.T) {
 	utxo := buildPlatformUTXO(utxoID, txAssetID, addr)
 
 	// save utxo to state
-	if err := vm.state.PutUTXO(utxo); err != nil {
-		t.Fatal("Error saving utxo", err)
-	}
+	vm.state.AddUTXO(utxo)
 
 	var inputUTXOs []*avax.UTXO //nolint:prealloc
 	for _, utxoID := range tx.Unsigned.InputUTXOs() {
@@ -325,9 +319,7 @@ func TestIndexTransaction_UnorderedWrites(t *testing.T) {
 		utxo := buildPlatformUTXO(utxoID, txAssetID, addr)
 
 		// save utxo to state
-		if err := vm.state.PutUTXO(utxo); err != nil {
-			t.Fatal("Error saving utxo", err)
-		}
+		vm.state.AddUTXO(utxo)
 
 		// issue transaction
 		if _, err := vm.IssueTx(tx.Bytes()); err != nil {

vms/avm/service_test.go

@@ -304,8 +304,8 @@ func TestServiceGetBalanceStrict(t *testing.T) {
 		},
 	}
 	// Insert the UTXO
-	err = vm.state.PutUTXO(twoOfTwoUTXO)
-	require.NoError(t, err)
+	vm.state.AddUTXO(twoOfTwoUTXO)
+	require.NoError(t, vm.state.Commit())
 
 	// Check the balance with IncludePartial set to true
 	balanceArgs := &GetBalanceArgs{
@@ -349,8 +349,8 @@ func TestServiceGetBalanceStrict(t *testing.T) {
 		},
 	}
 	// Insert the UTXO
-	err = vm.state.PutUTXO(oneOfTwoUTXO)
-	require.NoError(t, err)
+	vm.state.AddUTXO(oneOfTwoUTXO)
+	require.NoError(t, vm.state.Commit())
 
 	// Check the balance with IncludePartial set to true
 	balanceArgs = &GetBalanceArgs{
@@ -396,8 +396,8 @@ func TestServiceGetBalanceStrict(t *testing.T) {
 		},
 	}
 	// Insert the UTXO
-	err = vm.state.PutUTXO(futureUTXO)
-	require.NoError(t, err)
+	vm.state.AddUTXO(futureUTXO)
+	require.NoError(t, vm.state.Commit())
 
 	// Check the balance with IncludePartial set to true
 	balanceArgs = &GetBalanceArgs{
@@ -500,8 +500,8 @@ func TestServiceGetAllBalances(t *testing.T) {
 		},
 	}
 	// Insert the UTXO
-	err = vm.state.PutUTXO(twoOfTwoUTXO)
-	require.NoError(t, err)
+	vm.state.AddUTXO(twoOfTwoUTXO)
+	require.NoError(t, vm.state.Commit())
 
 	// Check the balance with IncludePartial set to true
 	balanceArgs := &GetAllBalancesArgs{
@@ -542,8 +542,8 @@ func TestServiceGetAllBalances(t *testing.T) {
 		},
 	}
 	// Insert the UTXO
-	err = vm.state.PutUTXO(oneOfTwoUTXO)
-	require.NoError(t, err)
+	vm.state.AddUTXO(oneOfTwoUTXO)
+	require.NoError(t, vm.state.Commit())
 
 	// Check the balance with IncludePartial set to true
 	balanceArgs = &GetAllBalancesArgs{
@@ -587,8 +587,8 @@ func TestServiceGetAllBalances(t *testing.T) {
 		},
 	}
 	// Insert the UTXO
-	err = vm.state.PutUTXO(futureUTXO)
-	require.NoError(t, err)
+	vm.state.AddUTXO(futureUTXO)
+	require.NoError(t, vm.state.Commit())
 
 	// Check the balance with IncludePartial set to true
 	balanceArgs = &GetAllBalancesArgs{
@@ -630,8 +630,8 @@ func TestServiceGetAllBalances(t *testing.T) {
 		},
 	}
 	// Insert the UTXO
-	err = vm.state.PutUTXO(otherAssetUTXO)
-	require.NoError(t, err)
+	vm.state.AddUTXO(otherAssetUTXO)
+	require.NoError(t, vm.state.Commit())
 
 	// Check the balance with IncludePartial set to true
 	balanceArgs = &GetAllBalancesArgs{
@@ -1829,10 +1829,9 @@ func TestServiceGetUTXOs(t *testing.T) {
 				},
 			},
 		}
-		if err := vm.state.PutUTXO(utxo); err != nil {
-			t.Fatal(err)
-		}
+		vm.state.AddUTXO(utxo)
 	}
+	require.NoError(t, vm.state.Commit())
 
 	sm := m.NewSharedMemory(constants.PlatformChainID)