refactor(store/v2)!: simplify storage

runtime/v2/builder.go

@@ -134,6 +134,7 @@ func (a *AppBuilder[T]) initGenesis(ctx context.Context, src io.Reader, txHandle
 	if err != nil {
 		return nil, fmt.Errorf("failed to read import state: %w", err)
 	}
+
 	var genesisJSON map[string]json.RawMessage
 	if err = json.Unmarshal(bz, &genesisJSON); err != nil {
 		return nil, err

server/v2/cometbft/abci_test.go

@@ -591,7 +591,7 @@ 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(&store.Changeset{
+	err := c.store.GetStateCommitment().WriteChangeset(&store.Changeset{
 		Version: 1,
 		Changes: []store.StateChanges{
 			{
@@ -691,9 +691,8 @@ func setUpConsensus(t *testing.T, gasLimit uint64, mempool mempool.Mempool[mock.
 	)
 	require.NoError(t, err)
 
-	ss := cometmock.NewMockStorage(log.NewNopLogger(), t.TempDir())
 	sc := cometmock.NewMockCommiter(log.NewNopLogger(), string(actorName), "stf")
-	mockStore := cometmock.NewMockStore(ss, sc)
+	mockStore := cometmock.NewMockStore(sc)
 
 	am := appmanager.New(appmanager.Config{
 		ValidateTxGasLimit: gasLimit,
@@ -786,6 +785,7 @@ func TestOptimisticExecution(t *testing.T) {
 		Txs:    ppReq.Txs,
 	}
 	fbResp, err := c.FinalizeBlock(context.Background(), fbReq)
+	require.Nil(t, fbResp)
 	require.Error(t, err)
 	require.ErrorContains(t, err, "test error") // from optimisticMockFunc
 	require.Equal(t, 1, calledTimes)

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

@@ -39,7 +39,7 @@ func NewMockReader(v uint64, rs *MockStore, actor []byte) *MockReader {
 }
 
 func (roa *MockReader) Has(key []byte) (bool, error) {
-	val, err := roa.store.GetStateStorage().Has(roa.actor, roa.version, key)
+	val, err := roa.store.GetStateCommitment().Has(roa.actor, roa.version, key)
 	if err != nil {
 		return false, err
 	}
@@ -48,7 +48,7 @@ func (roa *MockReader) Has(key []byte) (bool, error) {
 }
 
 func (roa *MockReader) Get(key []byte) ([]byte, error) {
-	result, err := roa.store.GetStateStorage().Get(roa.actor, roa.version, key)
+	result, err := roa.store.GetStateCommitment().Get(roa.actor, roa.version, key)
 	if err != nil {
 		return nil, err
 	}
@@ -57,9 +57,9 @@ func (roa *MockReader) Get(key []byte) ([]byte, error) {
 }
 
 func (roa *MockReader) Iterator(start, end []byte) (corestore.Iterator, error) {
-	return roa.store.GetStateStorage().Iterator(roa.actor, roa.version, start, end)
+	return roa.store.GetStateCommitment().Iterator(roa.actor, roa.version, start, end)
 }
 
 func (roa *MockReader) ReverseIterator(start, end []byte) (corestore.Iterator, error) {
-	return roa.store.GetStateStorage().ReverseIterator(roa.actor, roa.version, start, end)
+	return roa.store.GetStateCommitment().ReverseIterator(roa.actor, roa.version, start, end)
 }

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

@@ -11,21 +11,12 @@ import (
 	"cosmossdk.io/store/v2/commitment/iavl"
 	dbm "cosmossdk.io/store/v2/db"
 	"cosmossdk.io/store/v2/proof"
-	"cosmossdk.io/store/v2/storage"
-	"cosmossdk.io/store/v2/storage/pebbledb"
 )
 
 type MockStore struct {
-	Storage   storev2.VersionedWriter
 	Committer storev2.Committer
 }
 
-func NewMockStorage(logger log.Logger, dir string) storev2.VersionedWriter {
-	storageDB, _ := pebbledb.New(dir)
-	ss := storage.NewStorageStore(storageDB, logger)
-	return ss
-}
-
 func NewMockCommiter(logger log.Logger, actors ...string) storev2.Committer {
 	treeMap := make(map[string]commitment.Tree)
 	for _, actor := range actors {
@@ -36,8 +27,8 @@ func NewMockCommiter(logger log.Logger, actors ...string) storev2.Committer {
 	return sc
 }
 
-func NewMockStore(ss storev2.VersionedWriter, sc storev2.Committer) *MockStore {
-	return &MockStore{Storage: ss, Committer: sc}
+func NewMockStore(sc storev2.Committer) *MockStore {
+	return &MockStore{Committer: sc}
 }
 
 func (s *MockStore) GetLatestVersion() (uint64, error) {
@@ -59,12 +50,7 @@ func (s *MockStore) StateLatest() (uint64, corestore.ReaderMap, error) {
 }
 
 func (s *MockStore) Commit(changeset *corestore.Changeset) (corestore.Hash, error) {
-	err := s.Storage.ApplyChangeset(changeset)
-	if err != nil {
-		return []byte{}, err
-	}
-
-	err = s.Committer.WriteChangeset(changeset)
+	err := s.Committer.WriteChangeset(changeset)
 	if err != nil {
 		return []byte{}, err
 	}
@@ -81,10 +67,6 @@ func (s *MockStore) StateAt(version uint64) (corestore.ReaderMap, error) {
 	return NewMockReaderMap(version, s), nil
 }
 
-func (s *MockStore) GetStateStorage() storev2.VersionedWriter {
-	return s.Storage
-}
-
 func (s *MockStore) GetStateCommitment() storev2.Committer {
 	return s.Committer
 }

server/v2/cometbft/server.go

@@ -127,7 +127,6 @@ func New[T transaction.Tx](
 		indexEvents[e] = struct{}{}
 	}
 
-	ss := store.GetStateStorage().(snapshots.StorageSnapshotter)
 	sc := store.GetStateCommitment().(snapshots.CommitSnapshotter)
 
 	snapshotStore, err := GetSnapshotStore(srv.config.ConfigTomlConfig.RootDir)
@@ -155,7 +154,6 @@ func New[T transaction.Tx](
 		snapshotStore,
 		srv.serverOptions.SnapshotOptions(cfg),
 		sc,
-		ss,
 		nil, // extensions snapshotter registered below
 		logger,
 	)

server/v2/stf/branch/bench_test.go

@@ -105,6 +105,7 @@ func Benchmark_Iterate(b *testing.B) {
 
 // makeBranchStack creates a branch stack of the given size and initializes it with unique key-value pairs.
 func makeBranchStack(b *testing.B, stackSize int) Store[store.KVStore] {
+	b.Helper()
 	parent := coretesting.NewMemKV()
 	branch := NewStore[store.KVStore](parent)
 	for i := 1; i < stackSize; i++ {

server/v2/store/snapshot.go

@@ -375,10 +375,11 @@ func createSnapshotsManager(
 	}
 
 	sm := snapshots.NewManager(
-		snapshotStore, snapshots.NewSnapshotOptions(interval, uint32(keepRecent)),
+		snapshotStore,
+		snapshots.NewSnapshotOptions(interval, uint32(keepRecent)),
 		store.GetStateCommitment().(snapshots.CommitSnapshotter),
-		store.GetStateStorage().(snapshots.StorageSnapshotter),
-		nil, logger)
+		nil,
+		logger)
 	return sm, nil
 }
 

store/iavl/store_test.go

@@ -79,8 +79,6 @@ func TestLoadStore(t *testing.T) {
 	cIDHp := types.CommitID{Version: verHp, Hash: hash}
 	require.Nil(t, err)
 
-	// TODO: Prune this height
-
 	// Create current height Hc
 	updated, err = tree.Set([]byte("hello"), []byte("ciao"))
 	require.NoError(t, err)

store/v2/commitment/iavl/tree.go

@@ -83,6 +83,16 @@ func (t *IavlTree) Commit() ([]byte, uint64, error) {
 
 // GetProof returns a proof for the given key and version.
 func (t *IavlTree) GetProof(version uint64, key []byte) (*ics23.CommitmentProof, error) {
+	// the mutable tree is empty at genesis & when the storekey is removed, but the immutable tree is not but the immutable tree is not empty when the storekey is removed
+	// by checking the latest version we can determine if we are in genesis or have a key that has been removed
+	lv, err := t.tree.GetLatestVersion()
+	if err != nil {
+		return nil, err
+	}
+	if lv == 0 {
+		return t.tree.GetProof(key)
+	}
+
 	immutableTree, err := t.tree.GetImmutable(int64(version))
 	if err != nil {
 		return nil, fmt.Errorf("failed to get immutable tree at version %d: %w", version, err)
@@ -93,6 +103,16 @@ func (t *IavlTree) GetProof(version uint64, key []byte) (*ics23.CommitmentProof,
 
 // Get implements the Reader interface.
 func (t *IavlTree) Get(version uint64, key []byte) ([]byte, error) {
+	// the mutable tree is empty at genesis & when the storekey is removed, but the immutable tree is not but the immutable tree is not empty when the storekey is removed
+	// by checking the latest version we can determine if we are in genesis or have a key that has been removed
+	lv, err := t.tree.GetLatestVersion()
+	if err != nil {
+		return nil, err
+	}
+	if lv == 0 {
+		return t.tree.Get(key)
+	}
+
 	immutableTree, err := t.tree.GetImmutable(int64(version))
 	if err != nil {
 		return nil, fmt.Errorf("failed to get immutable tree at version %d: %w", version, err)
@@ -103,6 +123,16 @@ func (t *IavlTree) Get(version uint64, key []byte) ([]byte, error) {
 
 // Iterator implements the Reader interface.
 func (t *IavlTree) Iterator(version uint64, start, end []byte, ascending bool) (corestore.Iterator, error) {
+	// the mutable tree is empty at genesis & when the storekey is removed, but the immutable tree is not empty when the storekey is removed
+	// by checking the latest version we can determine if we are in genesis or have a key that has been removed
+	lv, err := t.tree.GetLatestVersion()
+	if err != nil {
+		return nil, err
+	}
+	if lv == 0 {
+		return t.tree.Iterator(start, end, ascending)
+	}
+
 	immutableTree, err := t.tree.GetImmutable(int64(version))
 	if err != nil {
 		return nil, fmt.Errorf("failed to get immutable tree at version %d: %w", version, err)

store/v2/commitment/store.go

@@ -233,6 +233,7 @@ func (c *CommitStore) SetInitialVersion(version uint64) error {
 	return nil
 }
 
+// GetProof returns a proof for the given key and version.
 func (c *CommitStore) GetProof(storeKey []byte, version uint64, key []byte) ([]proof.CommitmentOp, error) {
 	rawStoreKey := conv.UnsafeBytesToStr(storeKey)
 	tree, ok := c.multiTrees[rawStoreKey]
@@ -268,8 +269,12 @@ func (c *CommitStore) GetProof(storeKey []byte, version uint64, key []byte) ([]p
 // WARNING: This function is only used during the migration process. The SC layer
 // generally does not provide a reader for the CommitStore.
 func (c *CommitStore) getReader(storeKey string) (Reader, error) {
-	tree, ok := c.multiTrees[storeKey]
-	if !ok {
+	var tree Tree
+	if storeTree, ok := c.oldTrees[storeKey]; ok {
+		tree = storeTree
+	} else if storeTree, ok := c.multiTrees[storeKey]; ok {
+		tree = storeTree
+	} else {
 		return nil, fmt.Errorf("store %s not found", storeKey)
 	}
 
@@ -283,6 +288,14 @@ func (c *CommitStore) getReader(storeKey string) (Reader, error) {
 
 // VersionExists implements store.VersionedReader.
 func (c *CommitStore) VersionExists(version uint64) (bool, error) {
+	latestVersion, err := c.metadata.GetLatestVersion()
+	if err != nil {
+		return false, err
+	}
+	if latestVersion == 0 {
+		return version == 0, nil
+	}
+
 	ci, err := c.metadata.GetCommitInfo(version)
 	return ci != nil, err
 }
@@ -435,12 +448,10 @@ func (c *CommitStore) Restore(
 	version uint64,
 	format uint32,
 	protoReader protoio.Reader,
-	chStorage chan<- *corestore.StateChanges,
 ) (snapshotstypes.SnapshotItem, error) {
 	var (
 		importer     Importer
 		snapshotItem snapshotstypes.SnapshotItem
-		storeKey     []byte
 	)
 
 loop:
@@ -463,8 +474,6 @@ loop:
 					return snapshotstypes.SnapshotItem{}, fmt.Errorf("failed to close importer: %w", err)
 				}
 			}
-
-			storeKey = []byte(item.Store.Name)
 			tree := c.multiTrees[item.Store.Name]
 			if tree == nil {
 				return snapshotstypes.SnapshotItem{}, fmt.Errorf("store %s not found", item.Store.Name)
@@ -493,17 +502,6 @@ loop:
 				if node.Value == nil {
 					node.Value = []byte{}
 				}
-
-				// If the node is a leaf node, it will be written to the storage.
-				chStorage <- &corestore.StateChanges{
-					Actor: storeKey,
-					StateChanges: []corestore.KVPair{
-						{
-							Key:   node.Key,
-							Value: node.Value,
-						},
-					},
-				}
 			}
 			err := importer.Add(node)
 			if err != nil {