refactor(stf/appmanager): remove the need of recurseStateChanges (#20911)

server/v2/appmanager/appmanager.go

@@ -9,15 +9,25 @@ import (
 	appmanager "cosmossdk.io/core/app"
 	corestore "cosmossdk.io/core/store"
 	"cosmossdk.io/core/transaction"
-	"cosmossdk.io/server/v2/appmanager/store"
 )
 
+// Store defines the underlying storage behavior needed by AppManager.
+type Store interface {
+	// StateLatest returns a readonly view over the latest
+	// committed state of the store. Alongside the version
+	// associated with it.
+	StateLatest() (uint64, corestore.ReaderMap, error)
+
+	// StateAt returns a readonly view over the provided
+	// state. Must error when the version does not exist.
+	StateAt(version uint64) (corestore.ReaderMap, error)
+}
+
 // AppManager is a coordinator for all things related to an application
-// TODO: add exportGenesis function
 type AppManager[T transaction.Tx] struct {
 	config Config
 
-	db store.Store
+	db Store
 
 	initGenesis   InitGenesis
 	exportGenesis ExportGenesis
@@ -40,7 +50,7 @@ func (a AppManager[T]) InitGenesis(
 	}
 
 	var genTxs []T
-	zeroState, err = a.stf.RunWithCtx(ctx, zeroState, func(ctx context.Context) error {
+	genesisState, err := a.stf.RunWithCtx(ctx, zeroState, func(ctx context.Context) error {
 		return a.initGenesis(ctx, bytes.NewBuffer(initGenesisJSON), func(jsonTx json.RawMessage) error {
 			genTx, err := txDecoder.DecodeJSON(jsonTx)
 			if err != nil {
@@ -55,16 +65,26 @@ func (a AppManager[T]) InitGenesis(
 	}
 	// run block
 	// TODO: in an ideal world, genesis state is simply an initial state being applied
-	// unaware of what that state means in relation to every other, so here we can
-	// chain genesis
+	// unaware of what that state means in relation to every other
 	blockRequest.Txs = genTxs
 
-	blockresponse, genesisState, err := a.stf.DeliverBlock(ctx, blockRequest, zeroState)
+	blockResponse, blockZeroState, err := a.stf.DeliverBlock(ctx, blockRequest, genesisState)
 	if err != nil {
-		return blockresponse, nil, fmt.Errorf("failed to deliver block %d: %w", blockRequest.Height, err)
+		return blockResponse, nil, fmt.Errorf("failed to deliver block %d: %w", blockRequest.Height, err)
 	}
 
-	return blockresponse, genesisState, err
+	// after executing block 0, we extract the changes and apply them to the genesis state.
+	blockZeroStateChanges, err := blockZeroState.GetStateChanges()
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to get block zero state changes: %w", err)
+	}
+
+	err = genesisState.ApplyStateChanges(blockZeroStateChanges)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to apply blcok zero state changes to genesis state: %w", err)
+	}
+
+	return blockResponse, genesisState, err
 	// consensus server will need to set the version of the store
 }
 

server/v2/appmanager/appmanager_builder.go

@@ -2,14 +2,13 @@ package appmanager
 
 import (
 	"cosmossdk.io/core/transaction"
-	"cosmossdk.io/server/v2/appmanager/store"
 )
 
 // Builder is a struct that represents the application builder for managing transactions.
 // It contains various fields and methods for initializing the application and handling transactions.
 type Builder[T transaction.Tx] struct {
 	STF StateTransitionFunction[T] // The state transition function for processing transactions.
-	DB  store.Store                // The database for storing application data.
+	DB  Store                      // The database for storing application data.
 
 	// Gas limits for validating, querying, and simulating transactions.
 	ValidateTxGasLimit uint64

server/v2/appmanager/store/types.go

@@ -1,17 +0,0 @@
-package store
-
-import (
-	"cosmossdk.io/core/store"
-)
-
-// Store defines the underlying storage engine of an app.
-type Store interface {
-	// StateLatest returns a readonly view over the latest
-	// committed state of the store. Alongside the version
-	// associated with it.
-	StateLatest() (uint64, store.ReaderMap, error)
-
-	// StateAt returns a readonly view over the provided
-	// state. Must error when the version does not exist.
-	StateAt(version uint64) (store.ReaderMap, error)
-}

server/v2/stf/branch/writer_map.go

@@ -53,47 +53,22 @@ func (b WriterMap) ApplyStateChanges(stateChanges []store.StateChanges) error {
 	return nil
 }
 
-// GetStateChanges returns the state changes for all actors in the WriterMap, including all direct
-// ancesotors from which this WriterMap was derived.
-// See WriterMap.recurseStateChanges for more details.
-// Subject to possible renaming to ensure a developer can retrieve only changes in *this* branch
-// context (not ancestors) if that is desired.
-// see: https://github.com/cosmos/cosmos-sdk/pull/20412#discussion_r1618771230
+// GetStateChanges returns the state changes for all actors in the WriterMap.
 func (b WriterMap) GetStateChanges() ([]store.StateChanges, error) {
-	var (
-		changes = make(map[string][]store.KVPair)
-		sc      []store.StateChanges
-	)
-	if err := b.recurseStateChanges(changes); err != nil {
-		return nil, err
-	}
-
-	for account, kvPairs := range changes {
+	sc := make([]store.StateChanges, 0, len(b.branchedWriterState))
+	for acc, w := range b.branchedWriterState {
+		accBytes := []byte(acc)
+		kvChanges, err := w.ChangeSets()
+		if err != nil {
+			return nil, fmt.Errorf("unable to get actor writer changes %x: %w", accBytes, err)
+		}
 		sc = append(sc, store.StateChanges{
-			Actor:        []byte(account),
-			StateChanges: kvPairs,
+			Actor:        accBytes,
+			StateChanges: kvChanges,
 		})
 	}
-	return sc, nil
-}
 
-// recurseStateChanges will recursively collect state changes from the tree of
-// WriterMap's and write them to the `changes` map.
-func (b WriterMap) recurseStateChanges(changes map[string][]store.KVPair) error {
-	// depth first
-	if wr, ok := b.state.(WriterMap); ok {
-		if err := wr.recurseStateChanges(changes); err != nil {
-			return err
-		}
-	}
-	for account, stateChange := range b.branchedWriterState {
-		kvChanges, err := stateChange.ChangeSets()
-		if err != nil {
-			return err
-		}
-		changes[account] = append(changes[account], kvChanges...)
-	}
-	return nil
+	return sc, nil
 }
 
 func (b WriterMap) applyStateChange(sc store.StateChanges) error {