From 6e4084b88e019a4b35e2bc56db22e477f53e6acc Mon Sep 17 00:00:00 2001
From: stepit <48993133+0xstepit@users.noreply.github.com>
Date: Mon, 1 Jul 2024 19:18:22 +0200
Subject: [PATCH]  chore(cachemulti): add Copy method (#52)

Co-authored-by: tom <tomasguerraalda@hotmail.com>
Co-authored-by: Tom <54514587+GAtom22@users.noreply.github.com>
---
 CHANGELOG.md              |  4 ++++
 store/cachekv/store.go    | 32 ++++++++++++++++++++++++++++++++
 store/cachemulti/store.go | 32 ++++++++++++++++++++++++++++++++
 store/types/store.go      |  3 ++-
 4 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2381da211b9d..537ddc1623ea 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -35,6 +35,10 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 # Changelog
 
+## Unreleased
+
+* (cache-store) [#52](https://github.com/evmos/cosmos-sdk/pull/52) Add a deep copy method for the store.
+
 ## [v0.47.12-evmos](https://github.com/cosmos/evmos/releases/tag/v0.47.12-evmos) - 2024-06-20
 
 ## [v0.47.13](https://github.com/cosmos/cosmos-sdk/releases/tag/v0.47.13) - 2024-07-15
diff --git a/store/cachekv/store.go b/store/cachekv/store.go
index a8f468979cd9..ec4819ac9050 100644
--- a/store/cachekv/store.go
+++ b/store/cachekv/store.go
@@ -153,6 +153,38 @@ func (store *Store) CacheWrapWithTrace(w io.Writer, tc types.TraceContext) types
 	return NewStore(tracekv.NewStore(store, w, tc))
 }
 
+// Copy creates a deep copy of the Store object
+func (store *Store) Copy() *Store {
+	store.mtx.Lock()
+	defer store.mtx.Unlock()
+
+	// Copy cache
+	cacheCopy := make(map[string]*cValue, len(store.cache))
+	for key, val := range store.cache {
+		newVal := *val // Create a copy of the cValue
+		cacheCopy[key] = &newVal
+	}
+
+	// Copy unsortedCache
+	unsortedCacheCopy := make(map[string]struct{}, len(store.unsortedCache))
+	for key := range store.unsortedCache {
+		unsortedCacheCopy[key] = struct{}{}
+	}
+
+	// Copy sortedCache
+	sortedCacheCopy := store.sortedCache.Copy()
+
+	// Create new Store with copied values
+	newStore := &Store{
+		cache:         cacheCopy,
+		unsortedCache: unsortedCacheCopy,
+		sortedCache:   sortedCacheCopy,
+		parent:        store.parent,
+	}
+
+	return newStore
+}
+
 //----------------------------------------
 // Iteration
 
diff --git a/store/cachemulti/store.go b/store/cachemulti/store.go
index 86927466a905..c94b1c42c15a 100644
--- a/store/cachemulti/store.go
+++ b/store/cachemulti/store.go
@@ -168,3 +168,35 @@ 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 {
+	// Deep copy the db field
+	newDB := cms.db.(*cachekv.Store).Copy()
+
+	// Deep copy the cachekv stores map
+	newStores := make(map[types.StoreKey]types.CacheWrap, len(cms.stores))
+	for key, store := range cms.stores {
+		store, ok := store.(*cachekv.Store)
+		if ok {
+			newStores[key] = store.Copy()
+		}
+	}
+
+	// Deep copy the keys map
+	newKeys := make(map[string]types.StoreKey, len(cms.keys))
+	for key, value := range cms.keys {
+		newKeys[key] = value
+	}
+
+	// Create new Store with copied values
+	newStore := Store{
+		db:           newDB,
+		stores:       newStores,
+		keys:         newKeys,
+		traceWriter:  cms.traceWriter,
+		traceContext: cms.traceContext,
+	}
+
+	return newStore
+}
diff --git a/store/types/store.go b/store/types/store.go
index 992173195661..5b4d25a83118 100644
--- a/store/types/store.go
+++ b/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 // Returns a deep copy of the CacheMultiStore
 }
 
 // CommitMultiStore is an interface for a MultiStore without cache capabilities.