feat: copy cached store

store/cachekv/store.go

@@ -388,3 +388,29 @@ func (store *Store) setCacheValue(key, value []byte, dirty bool) {
 		store.unsortedCache[keyStr] = struct{}{}
 	}
 }
+
+// Copy branches a new cache kv store with same parent and caches
+func (store *Store) Copy() *Store {
+	store.mtx.Lock()
+	defer store.mtx.Unlock()
+
+	// new cache
+	newCache := make(map[string]*cValue, len(store.cache))
+	for k, v := range store.cache {
+		tmp := *v // copy to a new struct
+		newCache[k] = &tmp
+	}
+
+	// new unsorted cache
+	newUnsortedCache := make(map[string]struct{}, len(store.unsortedCache))
+	for key := range store.unsortedCache {
+		newUnsortedCache[key] = struct{}{}
+	}
+
+	return &Store{
+		cache:         newCache,
+		unsortedCache: newUnsortedCache,
+		sortedCache:   store.sortedCache.Copy(),
+		parent:        store.parent,
+	}
+}

store/cachemulti/store.go

@@ -3,6 +3,7 @@ package cachemulti
 import (
 	"fmt"
 	"io"
+	"maps"
 
 	dbm "github.com/cometbft/cometbft-db"
 
@@ -168,3 +169,22 @@ func (cms Store) GetKVStore(key types.StoreKey) types.KVStore {
 	}
 	return store.(types.KVStore)
 }
+
+// Copy creates a deep copy of the Store object
+func (cms Store) Copy() types.CacheMultiStore {
+	// new stores
+	newStores := make(map[types.StoreKey]types.CacheWrap, len(cms.stores))
+	for key, store := range cms.stores {
+		if cacheStore, ok := store.(*cachekv.Store); ok {
+			newStores[key] = cacheStore.Copy()
+		}
+	}
+
+	return Store{
+		db:           cms.db.(*cachekv.Store).Copy(),
+		stores:       newStores,
+		keys:         maps.Clone(cms.keys),
+		traceWriter:  cms.traceWriter,
+		traceContext: cms.traceContext,
+	}
+}

store/types/store.go

@@ -135,7 +135,8 @@ type MultiStore interface {
 // From MultiStore.CacheMultiStore()....
 type CacheMultiStore interface {
 	MultiStore
-	Write() // Writes operations to underlying KVStore
+	Write()                // Writes operations to underlying KVStore
+	Copy() CacheMultiStore // Deep copy a cached store
 }
 
 // CommitMultiStore is an interface for a MultiStore without cache capabilities.