Problem: Restore don't work with snapshot revert usage

store/CHANGELOG.md

@@ -33,6 +33,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * [#241](https://github.com/crypto-org-chain/cosmos-sdk/pull/241) Refactor the cache store to be btree backed, prepare to support copy-on-write atomic branching.
 * [#242](https://github.com/crypto-org-chain/cosmos-sdk/pull/242) Init cache on cache lazily, save memory allocations.
 * [#243](https://github.com/crypto-org-chain/cosmos-sdk/pull/243) Support `RunAtomic` API to use new CoW cache store.
+* [#244](https://github.com/crypto-org-chain/cosmos-sdk/pull/244) Add `Discard` method to CacheWrap to discard the write buffer.
 
 ## v1.1.0 (March 20, 2024)
 

store/cachekv/store.go

@@ -119,6 +119,10 @@ func (store *GStore[V]) Write() {
 	store.writeSet.Clear()
 }
 
+func (store *GStore[V]) Discard() {
+	store.writeSet.Clear()
+}
+
 // CacheWrap implements CacheWrapper.
 func (store *GStore[V]) CacheWrap() types.CacheWrap {
 	return NewGStore(store, store.isZero, store.valueLen)

store/cachemulti/store.go

@@ -45,7 +45,7 @@
 	}
 
 	for key, store := range stores {
-		cms.stores[key] = cms.initStore(key, store)
+		cms.initStore(key, store)
 	}
 
 	return cms
@@ -80,7 +80,9 @@
 			store = tracekv.NewStore(kvstore, cms.traceWriter, tctx)
 		}
 	}
-	return store.CacheWrap()
+	cache := store.CacheWrap()
+	cms.stores[key] = cache
+	return cache
 }
 
 // SetTracer sets the tracer for the MultiStore that the underlying
@@ -126,6 +128,12 @@
 	}
 }
 
+func (cms Store) Discard() {
+	for _, store := range cms.stores {
+		store.Discard()
+	}
+}
+
 // Implements CacheWrapper.
 func (cms Store) CacheWrap() types.CacheWrap {
 	return cms.CacheMultiStore().(types.CacheWrap)
@@ -138,14 +146,9 @@
 
 func (cms Store) getCacheWrap(key types.StoreKey) types.CacheWrap {
 	store, ok := cms.stores[key]
-	if !ok {
+	if !ok && cms.parentStore != nil {
 		// load on demand
-		if cms.branched {
-			store = cms.parentStore(key).(types.BranchStore).Clone().(types.CacheWrap)
-		} else if cms.parentStore != nil {
-			store = cms.initStore(key, cms.parentStore(key))
-		}
-		cms.stores[key] = store
+		store = cms.initStore(key, cms.parentStore(key))
 	}
 	if key == nil || store == nil {
 		panic(fmt.Sprintf("kv store with key %v has not been registered in stores", key))
@@ -181,12 +184,15 @@
 }
 
 func (cms Store) Clone() Store {
+	stores := make(map[types.StoreKey]types.CacheWrap, len(cms.stores))
+	for k, v := range cms.stores {
+		stores[k] = v.(types.BranchStore).Clone().(types.CacheWrap)
+	}
 	return Store{
-		stores: make(map[types.StoreKey]types.CacheWrap),
-
+		stores:       stores,
 		traceWriter:  cms.traceWriter,
 		traceContext: cms.traceContext,
-		parentStore:  cms.getCacheWrap,
+		parentStore:  cms.parentStore,
 
 		branched: true,
 	}
@@ -197,9 +203,22 @@
 		panic("cannot restore from non-branched store")
 	}
 
-	// restore the stores
+	// discard the non-exists stores
+	for k, v := range cms.stores {
+		if _, ok := other.stores[k]; !ok {
+			// clear the cache store if it's not in the other
+			v.Discard()
+		}
+	}
+
+	// restore the other stores
 	for k, v := range other.stores {
-		cms.stores[k].(types.BranchStore).Restore(v.(types.BranchStore))
+		store, ok := cms.stores[k]
+		if !ok {
+			store = cms.initStore(k, cms.parentStore(k))
+		}
+
+		store.(types.BranchStore).Restore(v.(types.BranchStore))
 	}
 }
 

store/cachemulti/store_test.go

@@ -35,22 +35,22 @@ func TestRunAtomic(t *testing.T) {
 		func(v any) int { return 1 },
 	)
 	keys := map[string]types.StoreKey{
-		"abc":  types.NewKVStoreKey("abc"),
-		"obj":  types.NewObjectStoreKey("obj"),
-		"lazy": types.NewKVStoreKey("lazy"),
+		"abc": types.NewKVStoreKey("abc"),
+		"obj": types.NewObjectStoreKey("obj"),
 	}
-	s := Store{stores: map[types.StoreKey]types.CacheWrap{
-		keys["abc"]:  store.CacheWrap(),
-		keys["obj"]:  objStore.CacheWrap(),
-		keys["lazy"]: nil,
+	parent := Store{stores: map[types.StoreKey]types.CacheWrap{
+		keys["abc"]: store.CacheWrap(),
+		keys["obj"]: objStore.CacheWrap(),
 	}}
 
+	s := Store{stores: map[types.StoreKey]types.CacheWrap{}, parentStore: parent.getCacheWrap}
 	s.RunAtomic(func(ms types.CacheMultiStore) error {
 		ms.GetKVStore(keys["abc"]).Set([]byte("key"), []byte("value"))
 		ms.GetObjKVStore(keys["obj"]).Set([]byte("key"), "value")
 		return nil
 	})
 	require.Equal(t, []byte("value"), s.GetKVStore(keys["abc"]).Get([]byte("key")))
+	require.Equal(t, []byte(nil), s.GetKVStore(keys["abc"]).Get([]byte("key-non-exist")))
 	require.Equal(t, "value", s.GetObjKVStore(keys["obj"]).Get([]byte("key")).(string))
 
 	require.Error(t, s.RunAtomic(func(ms types.CacheMultiStore) error {
@@ -61,3 +61,28 @@ func TestRunAtomic(t *testing.T) {
 	require.Equal(t, []byte("value"), s.GetKVStore(keys["abc"]).Get([]byte("key")))
 	require.Equal(t, "value", s.GetObjKVStore(keys["obj"]).Get([]byte("key")).(string))
 }
+
+func TestBranchStore(t *testing.T) {
+	store := dbadapter.Store{DB: dbm.NewMemDB()}
+	objStore := internal.NewBTreeStore(btree.NewBTree[any](),
+		func(v any) bool { return v == nil },
+		func(v any) int { return 1 },
+	)
+	keys := map[string]types.StoreKey{
+		"abc": types.NewKVStoreKey("abc"),
+		"obj": types.NewObjectStoreKey("obj"),
+	}
+	parent := Store{stores: map[types.StoreKey]types.CacheWrap{
+		keys["abc"]: store.CacheWrap(),
+		keys["obj"]: objStore.CacheWrap(),
+	}}
+
+	s := Store{stores: map[types.StoreKey]types.CacheWrap{}, parentStore: parent.getCacheWrap}
+	s.GetKVStore(keys["abc"]).Set([]byte("key"), []byte("value"))
+	snapshot := s.Clone()
+	s.GetKVStore(keys["abc"]).Set([]byte("key"), []byte("value2"))
+	s.GetObjKVStore(keys["obj"]).Set([]byte("key"), "value")
+	s.Restore(snapshot)
+	require.Equal(t, []byte("value"), s.GetKVStore(keys["abc"]).Get([]byte("key")))
+	require.Equal(t, nil, s.GetObjKVStore(keys["obj"]).Get([]byte("key")))
+}

store/types/store.go

@@ -339,6 +339,9 @@ type CacheWrap interface {
 
 	// Write syncs with the underlying store.
 	Write()
+
+	// Discard the write set
+	Discard()
 }
 
 type CacheWrapper interface {