From 2486622de9bdb83814d6fdaf00f955d448f777a9 Mon Sep 17 00:00:00 2001
From: forcodedancing <just.haha.it@gmail.com>
Date: Tue, 4 Jul 2023 10:59:16 +0800
Subject: [PATCH] feat: enable diff on iavl store (#231)

---
 store/iavl/store.go | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/store/iavl/store.go b/store/iavl/store.go
index 77f6113805..e59c834a7d 100644
--- a/store/iavl/store.go
+++ b/store/iavl/store.go
@@ -38,6 +38,7 @@ var (
 type Store struct {
 	tree   Tree
 	logger log.Logger
+	diff   map[string]struct{}
 }
 
 // LoadStore returns an IAVL Store as a CommitKVStore. Internally, it will load the
@@ -190,6 +191,18 @@ func (st *Store) CacheWrapWithTrace(w io.Writer, tc types.TraceContext) types.Ca
 	return cachekv.NewStore(tracekv.NewStore(st, w, tc))
 }
 
+func (st *Store) EnableDiff() {
+	st.diff = map[string]struct{}{}
+}
+
+func (st *Store) GetDiff() map[string]struct{} {
+	return st.diff
+}
+
+func (st *Store) ResetDiff() {
+	st.diff = map[string]struct{}{}
+}
+
 // Implements types.KVStore.
 func (st *Store) Set(key, value []byte) {
 	types.AssertValidKey(key)
@@ -198,6 +211,10 @@ func (st *Store) Set(key, value []byte) {
 	if err != nil && st.logger != nil {
 		st.logger.Error("iavl set error", "error", err.Error())
 	}
+
+	if st.diff != nil {
+		st.diff[string(key)] = struct{}{}
+	}
 }
 
 // Implements types.KVStore.
@@ -224,6 +241,10 @@ func (st *Store) Has(key []byte) (exists bool) {
 func (st *Store) Delete(key []byte) {
 	defer telemetry.MeasureSince(time.Now(), "store", "iavl", "delete")
 	st.tree.Remove(key)
+
+	if st.diff != nil {
+		st.diff[string(key)] = struct{}{}
+	}
 }
 
 // DeleteVersions deletes a series of versions from the MutableTree. An error