feat(stf): port simappv2 changes (#20587)

Co-authored-by: unknown unknown <unknown@unknown>

server/v2/stf/branch/writer_map.go

@@ -53,21 +53,47 @@ 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
 func (b WriterMap) GetStateChanges() ([]store.StateChanges, error) {
-	sc := make([]store.StateChanges, len(b.branchedWriterState))
-	for account, stateChange := range b.branchedWriterState {
-		kvChanges, err := stateChange.ChangeSets()
-		if err != nil {
-			return nil, err
-		}
+	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 = append(sc, store.StateChanges{
 			Actor:        []byte(account),
-			StateChanges: kvChanges,
+			StateChanges: kvPairs,
 		})
 	}
 	return sc, nil
 }
 
+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
+}
+
 func (b WriterMap) applyStateChange(sc store.StateChanges) error {
 	writableState, err := b.GetWriter(sc.Actor)
 	if err != nil {

server/v2/stf/core_router_service.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"reflect"
 	"strings"
 
 	"google.golang.org/protobuf/runtime/protoiface"
@@ -61,14 +62,8 @@ func (m *msgRouterService) InvokeUntyped(ctx context.Context, msg protoiface.Mes
 
 // NewQueryRouterService implements router.Service.
 func NewQueryRouterService(queryRouterBuilder *MsgRouterBuilder) router.Service {
-	queryRouter, err := queryRouterBuilder.Build()
-	if err != nil {
-		panic(fmt.Errorf("cannot create queryRouter: %w", err))
-	}
-
 	return &queryRouterService{
 		builder: queryRouterBuilder,
-		handler: queryRouter,
 	}
 }
 
@@ -100,8 +95,21 @@ func (m *queryRouterService) InvokeTyped(
 	ctx context.Context,
 	req, resp protoiface.MessageV1,
 ) error {
-	// see https://github.com/cosmos/cosmos-sdk/pull/20349
-	panic("not implemented")
+	// TODO lazy initialization is ugly and not thread safe. we don't want to check a mutex on every InvokeTyped either.
+	if m.handler == nil {
+		var err error
+		m.handler, err = m.builder.Build()
+		if err != nil {
+			return fmt.Errorf("cannot create queryRouter: %w", err)
+		}
+	}
+	// reflection is required, see https://github.com/cosmos/cosmos-sdk/pull/20349
+	res, err := m.handler(ctx, req)
+	if err != nil {
+		return err
+	}
+	reflect.Indirect(reflect.ValueOf(resp)).Set(reflect.Indirect(reflect.ValueOf(res)))
+	return nil
 }
 
 // InvokeUntyped execute a message and returns a response.

server/v2/stf/stf.go

@@ -7,6 +7,7 @@ import (
 
 	appmanager "cosmossdk.io/core/app"
 	appmodulev2 "cosmossdk.io/core/appmodule/v2"
+	corecontext "cosmossdk.io/core/context"
 	"cosmossdk.io/core/event"
 	"cosmossdk.io/core/gas"
 	"cosmossdk.io/core/header"
@@ -17,6 +18,9 @@ import (
 	"cosmossdk.io/server/v2/stf/internal"
 )
 
+// Identity defines STF's bytes identity and it's used by STF to store things in its own state.
+var Identity = []byte("stf")
+
 // STF is a struct that manages the state transition component of the app.
 type STF[T transaction.Tx] struct {
 	logger      log.Logger
@@ -108,10 +112,15 @@ func (s STF[T]) DeliverBlock(
 
 	// reset events
 	exCtx.events = make([]event.Event, 0)
+
 	// begin block
-	beginBlockEvents, err := s.beginBlock(exCtx)
-	if err != nil {
-		return nil, nil, err
+	var beginBlockEvents []event.Event
+	if !block.IsGenesis {
+		// begin block
+		beginBlockEvents, err = s.beginBlock(exCtx)
+		if err != nil {
+			return nil, nil, err
+		}
 	}
 
 	// check if we need to return early
@@ -401,11 +410,13 @@ func (s STF[T]) validatorUpdates(
 	return ctx.events, valSetUpdates, nil
 }
 
-const headerInfoPrefix = 0x0
+const headerInfoPrefix = 0x37
 
 // setHeaderInfo sets the header info in the state to be used by queries in the future.
 func (s STF[T]) setHeaderInfo(state store.WriterMap, headerInfo header.Info) error {
-	runtimeStore, err := state.GetWriter(appmanager.RuntimeIdentity)
+	// TODO storing header info is too low level here, stf should be stateless.
+	// We should have a keeper that does this.
+	runtimeStore, err := state.GetWriter(Identity)
 	if err != nil {
 		return err
 	}
@@ -422,7 +433,7 @@ func (s STF[T]) setHeaderInfo(state store.WriterMap, headerInfo header.Info) err
 
 // getHeaderInfo gets the header info from the state. It should only be used for queries
 func (s STF[T]) getHeaderInfo(state store.WriterMap) (i header.Info, err error) {
-	runtimeStore, err := state.GetWriter(appmanager.RuntimeIdentity)
+	runtimeStore, err := state.GetWriter(Identity)
 	if err != nil {
 		return header.Info{}, err
 	}
@@ -579,11 +590,12 @@ func (s STF[T]) makeContext(
 	store store.WriterMap,
 	execMode transaction.ExecMode,
 ) *executionContext {
+	valuedCtx := context.WithValue(ctx, corecontext.ExecModeKey, execMode)
 	return newExecutionContext(
 		s.makeGasMeter,
 		s.makeGasMeteredState,
 		s.branchFn,
-		ctx,
+		valuedCtx,
 		sender,
 		store,
 		execMode,