refactor(store/v2): simplify genesis flow

core/store/changeset.go

@@ -6,6 +6,7 @@ import (
 
 // Changeset is a list of changes to be written to disk
 type Changeset struct {
+	Version uint64
 	Changes []StateChanges
 }
 
@@ -29,11 +30,11 @@ type KVPair = struct {
 	Remove bool
 }
 
-func NewChangeset() *Changeset {
-	return &Changeset{}
+func NewChangeset(version uint64) *Changeset {
+	return &Changeset{Version: version}
 }
 
-func NewChangesetWithPairs(pairs map[string]KVPairs) *Changeset {
+func NewChangesetWithPairs(version uint64, pairs map[string]KVPairs) *Changeset {
 	changes := make([]StateChanges, len(pairs))
 	i := 0
 	for storeKey, kvPairs := range pairs {
@@ -44,6 +45,7 @@ func NewChangesetWithPairs(pairs map[string]KVPairs) *Changeset {
 		i++
 	}
 	return &Changeset{
+		Version: version,
 		Changes: changes,
 	}
 }

runtime/v2/go.mod

@@ -5,6 +5,7 @@ go 1.23
 // server v2 integration
 replace (
 	cosmossdk.io/api => ../../api
+	cosmossdk.io/core => ../../core
 	cosmossdk.io/core/testing => ../../core/testing
 	cosmossdk.io/server/v2/appmanager => ../../server/v2/appmanager
 	cosmossdk.io/server/v2/stf => ../../server/v2/stf

runtime/v2/go.sum

@@ -2,8 +2,6 @@ buf.build/gen/go/cometbft/cometbft/protocolbuffers/go v1.35.1-20240701160653-fed
 buf.build/gen/go/cometbft/cometbft/protocolbuffers/go v1.35.1-20240701160653-fedbb9acfd2f.1/go.mod h1:JTBMfyi+qAXUHumX+rcD2WIq9FNWmdcNh5MjBnSw0L0=
 buf.build/gen/go/cosmos/gogo-proto/protocolbuffers/go v1.35.1-20240130113600-88ef6483f90f.1 h1:F78ecjvMtgd1aZ1Aj9cvBjURxVGCYvRM+OOy5eR+pjw=
 buf.build/gen/go/cosmos/gogo-proto/protocolbuffers/go v1.35.1-20240130113600-88ef6483f90f.1/go.mod h1:zqi/LZjZhyvjCMTEVIwAf5VRlkLduuCfqmZxgoormq0=
-cosmossdk.io/core v1.0.0-alpha.5 h1:McjYXAQ6XcT20v2uHyH7PhoWH8V+mebzfVFqT3GinsI=
-cosmossdk.io/core v1.0.0-alpha.5/go.mod h1:3u9cWq1FAVtiiCrDPpo4LhR+9V6k/ycSG4/Y/tREWCY=
 cosmossdk.io/depinject v1.1.0 h1:wLan7LG35VM7Yo6ov0jId3RHWCGRhe8E8bsuARorl5E=
 cosmossdk.io/depinject v1.1.0/go.mod h1:kkI5H9jCGHeKeYWXTqYdruogYrEeWvBQCw1Pj4/eCFI=
 cosmossdk.io/errors/v2 v2.0.0-20240731132947-df72853b3ca5 h1:IQNdY2kB+k+1OM2DvqFG1+UgeU1JzZrWtwuWzI3ZfwA=

server/v2/cometbft/abci.go

@@ -329,7 +329,7 @@ func (c *Consensus[T]) InitChain(ctx context.Context, req *abciproto.InitChainRe
 		IsGenesis: true,
 	}
 
-	blockresponse, genesisState, err := c.app.InitGenesis(
+	blockResponse, genesisState, err := c.app.InitGenesis(
 		ctx,
 		br,
 		req.AppStateBytes,
@@ -338,17 +338,16 @@ func (c *Consensus[T]) InitChain(ctx context.Context, req *abciproto.InitChainRe
 		return nil, fmt.Errorf("genesis state init failure: %w", err)
 	}
 
-	for _, txRes := range blockresponse.TxResults {
+	for _, txRes := range blockResponse.TxResults {
 		if err := txRes.Error; err != nil {
-			space, code, log := errorsmod.ABCIInfo(err, c.cfg.AppTomlConfig.Trace)
-			c.logger.Warn("genesis tx failed", "codespace", space, "code", code, "log", log)
+			space, code, txLog := errorsmod.ABCIInfo(err, c.cfg.AppTomlConfig.Trace)
+			c.logger.Warn("genesis tx failed", "codespace", space, "code", code, "log", txLog)
 		}
 	}
 
-	validatorUpdates := intoABCIValidatorUpdates(blockresponse.ValidatorUpdates)
+	validatorUpdates := intoABCIValidatorUpdates(blockResponse.ValidatorUpdates)
 
-	// set the initial version of the store
-	if err := c.store.SetInitialVersion(uint64(req.InitialHeight)); err != nil {
+	if err := c.store.SetInitialVersion(uint64(req.InitialHeight - 1)); err != nil {
 		return nil, fmt.Errorf("failed to set initial version: %w", err)
 	}
 
@@ -357,9 +356,10 @@ func (c *Consensus[T]) InitChain(ctx context.Context, req *abciproto.InitChainRe
 		return nil, err
 	}
 	cs := &store.Changeset{
+		Version: uint64(req.InitialHeight - 1),
 		Changes: stateChanges,
 	}
-	stateRoot, err := c.store.WorkingHash(cs)
+	stateRoot, err := c.store.Commit(cs)
 	if err != nil {
 		return nil, fmt.Errorf("unable to write the changeset: %w", err)
 	}
@@ -455,18 +455,6 @@ func (c *Consensus[T]) FinalizeBlock(
 		return nil, err
 	}
 
-	// we don't need to deliver the block in the genesis block
-	if req.Height == int64(c.initialHeight) {
-		appHash, err := c.store.Commit(store.NewChangeset())
-		if err != nil {
-			return nil, fmt.Errorf("unable to commit the changeset: %w", err)
-		}
-		c.lastCommittedHeight.Store(req.Height)
-		return &abciproto.FinalizeBlockResponse{
-			AppHash: appHash,
-		}, nil
-	}
-
 	// TODO(tip): can we expect some txs to not decode? if so, what we do in this case? this does not seem to be the case,
 	// considering that prepare and process always decode txs, assuming they're the ones providing txs we should never
 	// have a tx that fails decoding.
@@ -507,7 +495,7 @@ func (c *Consensus[T]) FinalizeBlock(
 	if err != nil {
 		return nil, err
 	}
-	appHash, err := c.store.Commit(&store.Changeset{Changes: stateChanges})
+	appHash, err := c.store.Commit(&store.Changeset{Version: uint64(req.Height), Changes: stateChanges})
 	if err != nil {
 		return nil, fmt.Errorf("unable to commit the changeset: %w", err)
 	}

server/v2/cometbft/abci_test.go

@@ -570,7 +570,8 @@ func TestConsensus_Query(t *testing.T) {
 	c := setUpConsensus(t, 100_000, cometmock.MockMempool[mock.Tx]{})
 
 	// Write data to state storage
-	err := c.store.GetStateStorage().ApplyChangeset(1, &store.Changeset{
+	err := c.store.GetStateStorage().ApplyChangeset(&store.Changeset{
+		Version: 1,
 		Changes: []store.StateChanges{
 			{
 				Actor: actorName,
@@ -687,7 +688,9 @@ func setUpConsensus(t *testing.T, gasLimit uint64, mempool mempool.Mempool[mock.
 		nil,
 	)
 
-	return NewConsensus[mock.Tx](log.NewNopLogger(), "testing-app", am, func() error { return nil }, mempool, map[string]struct{}{}, nil, mockStore, Config{AppTomlConfig: DefaultAppTomlConfig()}, mock.TxCodec{}, "test")
+	return NewConsensus[mock.Tx](log.NewNopLogger(), "testing-app", am, func() error { return nil },
+		mempool, map[string]struct{}{}, nil, mockStore,
+		Config{AppTomlConfig: DefaultAppTomlConfig()}, mock.TxCodec{}, "test")
 }
 
 // Check target version same with store's latest version

server/v2/cometbft/go.mod

@@ -4,6 +4,7 @@ go 1.23.1
 
 replace (
 	cosmossdk.io/api => ../../../api
+	cosmossdk.io/core => ../../../core
 	cosmossdk.io/core/testing => ../../../core/testing
 	cosmossdk.io/server/v2 => ../
 	cosmossdk.io/server/v2/appmanager => ../appmanager

server/v2/cometbft/go.sum

@@ -6,8 +6,6 @@ cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMT
 cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 cosmossdk.io/collections v0.4.0 h1:PFmwj2W8szgpD5nOd8GWH6AbYNi1f2J6akWXJ7P5t9s=
 cosmossdk.io/collections v0.4.0/go.mod h1:oa5lUING2dP+gdDquow+QjlF45eL1t4TJDypgGd+tv0=
-cosmossdk.io/core v1.0.0-alpha.5 h1:McjYXAQ6XcT20v2uHyH7PhoWH8V+mebzfVFqT3GinsI=
-cosmossdk.io/core v1.0.0-alpha.5/go.mod h1:3u9cWq1FAVtiiCrDPpo4LhR+9V6k/ycSG4/Y/tREWCY=
 cosmossdk.io/depinject v1.1.0 h1:wLan7LG35VM7Yo6ov0jId3RHWCGRhe8E8bsuARorl5E=
 cosmossdk.io/depinject v1.1.0/go.mod h1:kkI5H9jCGHeKeYWXTqYdruogYrEeWvBQCw1Pj4/eCFI=
 cosmossdk.io/errors v1.0.1 h1:bzu+Kcr0kS/1DuPBtUFdWjzLqyUuCiyHjyJB6srBV/0=

server/v2/cometbft/internal/mock/mock_store.go

@@ -60,7 +60,7 @@ func (s *MockStore) StateLatest() (uint64, corestore.ReaderMap, error) {
 
 func (s *MockStore) Commit(changeset *corestore.Changeset) (corestore.Hash, error) {
 	v, _, _ := s.StateLatest()
-	err := s.Storage.ApplyChangeset(v, changeset)
+	err := s.Storage.ApplyChangeset(changeset)
 	if err != nil {
 		return []byte{}, err
 	}
@@ -127,17 +127,3 @@ func (s *MockStore) LastCommitID() (proof.CommitID, error) {
 func (s *MockStore) SetInitialVersion(v uint64) error {
 	return s.Committer.SetInitialVersion(v)
 }
-
-func (s *MockStore) WorkingHash(changeset *corestore.Changeset) (corestore.Hash, error) {
-	v, _, _ := s.StateLatest()
-	err := s.Storage.ApplyChangeset(v, changeset)
-	if err != nil {
-		return []byte{}, err
-	}
-
-	err = s.Committer.WriteChangeset(changeset)
-	if err != nil {
-		return []byte{}, err
-	}
-	return []byte{}, nil
-}

server/v2/cometbft/types/store.go

@@ -19,10 +19,6 @@ type Store interface {
 	// SetInitialVersion sets the initial version of the store.
 	SetInitialVersion(uint64) error
 
-	// WorkingHash writes the provided changeset to the state and returns
-	// the working hash of the state.
-	WorkingHash(*store.Changeset) (store.Hash, error)
-
 	// Commit commits the provided changeset and returns
 	// the new state root of the state.
 	Commit(*store.Changeset) (store.Hash, error)

server/v2/store/server.go

@@ -47,7 +47,7 @@ func (s *Server[T]) Start(context.Context) error {
 }
 
 func (s *Server[T]) Stop(context.Context) error {
-	return s.store.Close()
+	return nil
 }
 
 func (s *Server[T]) CLICommands() serverv2.CLIConfig {

simapp/v2/go.mod

@@ -291,6 +291,7 @@ replace (
 // server v2 integration
 replace (
 	cosmossdk.io/api => ../../api
+	cosmossdk.io/core => ../../core
 	cosmossdk.io/core/testing => ../../core/testing
 	cosmossdk.io/runtime/v2 => ../../runtime/v2
 	cosmossdk.io/server/v2 => ../../server/v2

simapp/v2/go.sum

@@ -192,8 +192,6 @@ cloud.google.com/go/webrisk v1.4.0/go.mod h1:Hn8X6Zr+ziE2aNd8SliSDWpEnSS1u4R9+xX
 cloud.google.com/go/webrisk v1.5.0/go.mod h1:iPG6fr52Tv7sGk0H6qUFzmL3HHZev1htXuWDEEsqMTg=
 cloud.google.com/go/workflows v1.6.0/go.mod h1:6t9F5h/unJz41YqfBmqSASJSXccBLtD1Vwf+KmJENM0=
 cloud.google.com/go/workflows v1.7.0/go.mod h1:JhSrZuVZWuiDfKEFxU0/F1PQjmpnpcoISEXH2bcHC3M=
-cosmossdk.io/core v1.0.0-alpha.5 h1:McjYXAQ6XcT20v2uHyH7PhoWH8V+mebzfVFqT3GinsI=
-cosmossdk.io/core v1.0.0-alpha.5/go.mod h1:3u9cWq1FAVtiiCrDPpo4LhR+9V6k/ycSG4/Y/tREWCY=
 cosmossdk.io/depinject v1.1.0 h1:wLan7LG35VM7Yo6ov0jId3RHWCGRhe8E8bsuARorl5E=
 cosmossdk.io/depinject v1.1.0/go.mod h1:kkI5H9jCGHeKeYWXTqYdruogYrEeWvBQCw1Pj4/eCFI=
 cosmossdk.io/errors v1.0.1 h1:bzu+Kcr0kS/1DuPBtUFdWjzLqyUuCiyHjyJB6srBV/0=

store/v2/commitment/iavl/tree.go

@@ -51,7 +51,14 @@ func (t *IavlTree) Hash() []byte {
 	return t.tree.Hash()
 }
 
+// Version returns the current version of the tree.
+func (t *IavlTree) Version() uint64 {
+	return uint64(t.tree.Version())
+}
+
 // WorkingHash returns the working hash of the tree.
+// Danger! iavl.MutableTree.WorkingHash() is a mutating operation!
+// It advances the tree version by 1.
 func (t *IavlTree) WorkingHash() []byte {
 	return t.tree.WorkingHash()
 }

store/v2/commitment/mem/tree.go

@@ -26,8 +26,8 @@ func (t *Tree) Hash() []byte {
 	return nil
 }
 
-func (t *Tree) WorkingHash() []byte {
-	return nil
+func (t *Tree) Version() uint64 {
+	return 0
 }
 
 func (t *Tree) LoadVersion(version uint64) error {

store/v2/commitment/store.go

@@ -82,28 +82,6 @@ func (c *CommitStore) WriteChangeset(cs *corestore.Changeset) error {
 	return nil
 }
 
-func (c *CommitStore) WorkingCommitInfo(version uint64) *proof.CommitInfo {
-	storeInfos := make([]proof.StoreInfo, 0, len(c.multiTrees))
-	for storeKey, tree := range c.multiTrees {
-		if internal.IsMemoryStoreKey(storeKey) {
-			continue
-		}
-		bz := []byte(storeKey)
-		storeInfos = append(storeInfos, proof.StoreInfo{
-			Name: bz,
-			CommitID: proof.CommitID{
-				Version: version,
-				Hash:    tree.WorkingHash(),
-			},
-		})
-	}
-
-	return &proof.CommitInfo{
-		Version:    version,
-		StoreInfos: storeInfos,
-	}
-}
-
 func (c *CommitStore) LoadVersion(targetVersion uint64) error {
 	storeKeys := make([]string, 0, len(c.multiTrees))
 	for storeKey := range c.multiTrees {
@@ -184,7 +162,10 @@ func (c *CommitStore) loadVersion(targetVersion uint64, storeKeys []string) erro
 	// If the target version is greater than the latest version, it is the snapshot
 	// restore case, we should create a new commit info for the target version.
 	if targetVersion > latestVersion {
-		cInfo := c.WorkingCommitInfo(targetVersion)
+		cInfo, err := c.GetCommitInfo(targetVersion)
+		if err != nil {
+			return err
+		}
 		return c.metadata.flushCommitInfo(targetVersion, cInfo)
 	}
 
@@ -198,29 +179,16 @@ func (c *CommitStore) Commit(version uint64) (*proof.CommitInfo, error) {
 		if internal.IsMemoryStoreKey(storeKey) {
 			continue
 		}
-		// If a commit event execution is interrupted, a new iavl store's version
-		// will be larger than the RMS's metadata, when the block is replayed, we
-		// should avoid committing that iavl store again.
-		var commitID proof.CommitID
-		v, err := tree.GetLatestVersion()
+		hash, cversion, err := tree.Commit()
 		if err != nil {
 			return nil, err
 		}
-		if v >= version {
-			commitID.Version = version
-			commitID.Hash = tree.Hash()
-		} else {
-			hash, cversion, err := tree.Commit()
-			if err != nil {
-				return nil, err
-			}
-			if cversion != version {
-				return nil, fmt.Errorf("commit version %d does not match the target version %d", cversion, version)
-			}
-			commitID = proof.CommitID{
-				Version: version,
-				Hash:    hash,
-			}
+		if cversion != version {
+			return nil, fmt.Errorf("commit version %d does not match the target version %d", cversion, version)
+		}
+		commitID := proof.CommitID{
+			Version: version,
+			Hash:    hash,
 		}
 		storeInfos = append(storeInfos, proof.StoreInfo{
 			Name:     []byte(storeKey),
@@ -541,7 +509,40 @@ loop:
 }
 
 func (c *CommitStore) GetCommitInfo(version uint64) (*proof.CommitInfo, error) {
-	return c.metadata.GetCommitInfo(version)
+	// if the commit info is already stored, return it
+	ci, err := c.metadata.GetCommitInfo(version)
+	if err != nil {
+		return nil, err
+	}
+	if ci != nil {
+		return ci, nil
+	}
+	// otherwise built the commit info from the trees
+	storeInfos := make([]proof.StoreInfo, 0, len(c.multiTrees))
+	for storeKey, tree := range c.multiTrees {
+		if internal.IsMemoryStoreKey(storeKey) {
+			continue
+		}
+		v := tree.Version()
+		if v != version {
+			return nil, fmt.Errorf("tree version %d does not match the target version %d", v, version)
+		}
+		bz := []byte(storeKey)
+		storeInfos = append(storeInfos, proof.StoreInfo{
+			Name: bz,
+			CommitID: proof.CommitID{
+				Version: v,
+				Hash:    tree.Hash(),
+			},
+		})
+	}
+
+	ci = &proof.CommitInfo{
+		Version:    version,
+		StoreInfos: storeInfos,
+	}
+	ci.Hash()
+	return ci, nil
 }
 
 func (c *CommitStore) GetLatestVersion() (uint64, error) {
@@ -554,6 +555,5 @@ func (c *CommitStore) Close() error {
 			return err
 		}
 	}
-
 	return nil
 }