From 718ca2e1f70d337989ac9b784c03e0e67dc22485 Mon Sep 17 00:00:00 2001
From: yzang2019 <zymfrank@gmail.com>
Date: Wed, 10 Jul 2024 10:58:17 -0700
Subject: [PATCH 1/2] Add CacheMultiStoreForExport for wasm snapshot export

---
 store/cachemulti/store.go  | 16 ++++++++++++++++
 store/rootmulti/store.go   |  4 ++++
 store/types/store.go       |  5 +++++
 storev2/rootmulti/store.go | 30 ++++++++++++++++++++++++++++++
 4 files changed, 55 insertions(+)

diff --git a/store/cachemulti/store.go b/store/cachemulti/store.go
index 96ce20dfc..5cd02d33b 100644
--- a/store/cachemulti/store.go
+++ b/store/cachemulti/store.go
@@ -30,6 +30,7 @@ type Store struct {
 	traceContext types.TraceContext
 
 	listeners map[types.StoreKey][]types.WriteListener
+	closers   []io.Closer
 }
 
 var _ types.CacheMultiStore = Store{}
@@ -52,6 +53,7 @@ func NewFromKVStore(
 		traceWriter:  traceWriter,
 		traceContext: traceContext,
 		listeners:    listeners,
+		closers:      []io.Closer{},
 	}
 
 	for key, store := range stores {
@@ -225,3 +227,17 @@ func (cms Store) SetKVStores(handler func(sk types.StoreKey, s types.KVStore) ty
 	}
 	return cms
 }
+
+func (cms Store) CacheMultiStoreForExport(_ int64) (types.CacheMultiStore, error) {
+	panic("Not implemented")
+}
+
+func (cms Store) AddCloser(closer io.Closer) {
+	cms.closers = append(cms.closers, closer)
+}
+
+func (cms Store) Close() {
+	for _, closer := range cms.closers {
+		closer.Close()
+	}
+}
diff --git a/store/rootmulti/store.go b/store/rootmulti/store.go
index 77f7a728c..1a4619aca 100644
--- a/store/rootmulti/store.go
+++ b/store/rootmulti/store.go
@@ -587,6 +587,10 @@ func (rs *Store) CacheMultiStoreWithVersion(version int64) (types.CacheMultiStor
 	return cachemulti.NewStore(rs.db, cachedStores, rs.keysByName, rs.traceWriter, rs.getTracingContext(), rs.listeners), nil
 }
 
+func (rs *Store) CacheMultiStoreForExport(version int64) (types.CacheMultiStore, error) {
+	return rs.CacheMultiStoreWithVersion(version)
+}
+
 // GetStore returns a mounted Store for a given StoreKey. If the StoreKey does
 // not exist, it will panic. If the Store is wrapped in an inter-block cache, it
 // will be unwrapped prior to being returned.
diff --git a/store/types/store.go b/store/types/store.go
index b66e4a53b..701eb97c2 100644
--- a/store/types/store.go
+++ b/store/types/store.go
@@ -113,6 +113,9 @@ type MultiStore interface {
 	// each stored is loaded at a specific version (height).
 	CacheMultiStoreWithVersion(version int64) (CacheMultiStore, error)
 
+	// CacheMultiStoreForExport create a cache multistore specifically for state export
+	CacheMultiStoreForExport(version int64) (CacheMultiStore, error)
+
 	// Convenience for fetching substores.
 	// If the store does not exist, panics.
 	GetStore(StoreKey) Store
@@ -159,6 +162,8 @@ type CacheMultiStore interface {
 
 	// Writes operations to underlying KVStore
 	Write()
+
+	Close()
 }
 
 // CommitMultiStore is an interface for a MultiStore without cache capabilities.
diff --git a/storev2/rootmulti/store.go b/storev2/rootmulti/store.go
index cbb4ef89c..cb37c2c7c 100644
--- a/storev2/rootmulti/store.go
+++ b/storev2/rootmulti/store.go
@@ -269,6 +269,36 @@ func (rs *Store) CacheMultiStoreWithVersion(version int64) (types.CacheMultiStor
 	return cachemulti.NewStore(nil, stores, rs.storeKeys, nil, nil, nil), nil
 }
 
+func (rs *Store) CacheMultiStoreForExport(version int64) (types.CacheMultiStore, error) {
+	if version <= 0 || (rs.lastCommitInfo != nil && version == rs.lastCommitInfo.Version) {
+		return rs.CacheMultiStore(), nil
+	}
+	rs.mtx.RLock()
+	defer rs.mtx.RUnlock()
+	stores := make(map[types.StoreKey]types.CacheWrapper)
+	// add the transient/mem stores registered in current app.
+	for k, store := range rs.ckvStores {
+		if store.GetStoreType() != types.StoreTypeIAVL {
+			stores[k] = store
+		}
+	}
+	// add SC stores for historical queries
+	scStore, err := rs.scStore.LoadVersion(version, true)
+	if err != nil {
+		return nil, err
+	}
+	for k, store := range rs.ckvStores {
+		if store.GetStoreType() == types.StoreTypeIAVL {
+			tree := scStore.GetTreeByName(k.Name())
+			stores[k] = commitment.NewStore(tree, rs.logger)
+		}
+	}
+	cacheMs := cachemulti.NewStore(nil, stores, rs.storeKeys, nil, nil, nil)
+	// We need this because we need to make sure sc is closed after being used to release the resources
+	cacheMs.AddCloser(scStore)
+	return cacheMs, nil
+}
+
 // GetStore Implements interface MultiStore
 func (rs *Store) GetStore(key types.StoreKey) types.Store {
 	return rs.ckvStores[key]

From fcc046d8f39b035d215e62995225d916673b0f3a Mon Sep 17 00:00:00 2001
From: yzang2019 <zymfrank@gmail.com>
Date: Wed, 10 Jul 2024 11:04:46 -0700
Subject: [PATCH 2/2] Fix test

---
 server/mock/store.go | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/server/mock/store.go b/server/mock/store.go
index be3523caa..3bdf67dfd 100644
--- a/server/mock/store.go
+++ b/server/mock/store.go
@@ -29,6 +29,10 @@ func (ms multiStore) CacheMultiStoreWithVersion(_ int64) (sdk.CacheMultiStore, e
 	panic("not implemented")
 }
 
+func (ms multiStore) CacheMultiStoreForExport(version int64) (store.CacheMultiStore, error) {
+	panic("not implemented")
+}
+
 func (ms multiStore) CacheWrap(_ store.StoreKey) sdk.CacheWrap {
 	panic("not implemented")
 }