From 01ca69241ef9ec3f034acc277b24d246dfa04c93 Mon Sep 17 00:00:00 2001
From: yihuang <yi.codeplayer@gmail.com>
Date: Mon, 1 Apr 2024 16:46:09 +0800
Subject: [PATCH] add Discard method to CacheWrap

---
 store/CHANGELOG.md        |  1 +
 store/cachekv/store.go    |  4 ++++
 store/cachemulti/store.go | 12 +++++++++---
 store/types/store.go      |  3 +++
 4 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/store/CHANGELOG.md b/store/CHANGELOG.md
index 850f6a98f19d5..39cfdacfa33ac 100644
--- a/store/CHANGELOG.md
+++ b/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)
 
diff --git a/store/cachekv/store.go b/store/cachekv/store.go
index c11ddd9e2b219..5d174acaaeb8a 100644
--- a/store/cachekv/store.go
+++ b/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)
diff --git a/store/cachemulti/store.go b/store/cachemulti/store.go
index c30b02f27462c..73979be54b19e 100644
--- a/store/cachemulti/store.go
+++ b/store/cachemulti/store.go
@@ -126,6 +126,12 @@ func (cms Store) Write() {
 	}
 }
 
+func (cms Store) Discard() {
+	for _, store := range cms.stores {
+		store.Discard()
+	}
+}
+
 // Implements CacheWrapper.
 func (cms Store) CacheWrap() types.CacheWrap {
 	return cms.CacheMultiStore().(types.CacheWrap)
@@ -202,10 +208,10 @@ func (cms Store) Restore(other Store) {
 		cms.stores[k].(types.BranchStore).Restore(v.(types.BranchStore))
 	}
 
-	for k := range cms.stores {
+	for k, v := range cms.stores {
 		if _, ok := other.stores[k]; !ok {
-			// delete the store if it's not in the other
-			delete(cms.stores, k)
+			// clear the cache store if it's not in the other
+			v.Discard()
 		}
 	}
 }
diff --git a/store/types/store.go b/store/types/store.go
index e63ad6d9a79b4..67bd140f5e753 100644
--- a/store/types/store.go
+++ b/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 {