From 2d18f1e7e94b1b3abd94e15bf5f0e99d3c284fa0 Mon Sep 17 00:00:00 2001
From: Alexander Bezobchuk <alexanderbez@users.noreply.github.com>
Date: Thu, 5 Sep 2019 12:29:00 -0400
Subject: [PATCH] Merge PR #4997: Fix CacheMultiStoreWithVersion (inter-block
 cache)

---
 store/cache/cache.go      |  9 +++++++++
 store/cache/cache_test.go | 12 ++++++++++++
 store/rootmulti/store.go  |  8 ++++++++
 store/types/store.go      |  3 +++
 4 files changed, 32 insertions(+)

diff --git a/store/cache/cache.go b/store/cache/cache.go
index 2293c3bbde..9eea90ad94 100644
--- a/store/cache/cache.go
+++ b/store/cache/cache.go
@@ -70,6 +70,15 @@ func (cmgr *CommitKVStoreCacheManager) GetStoreCache(key types.StoreKey, store t
 	return cmgr.caches[key.Name()]
 }
 
+// Unwrap returns the underlying CommitKVStore for a given StoreKey.
+func (cmgr *CommitKVStoreCacheManager) Unwrap(key types.StoreKey) types.CommitKVStore {
+	if ckv, ok := cmgr.caches[key.Name()]; ok {
+		return ckv.(*CommitKVStoreCache).CommitKVStore
+	}
+
+	return nil
+}
+
 // Reset resets in the internal caches.
 func (cmgr *CommitKVStoreCacheManager) Reset() {
 	cmgr.caches = make(map[string]types.CommitKVStore)
diff --git a/store/cache/cache_test.go b/store/cache/cache_test.go
index e2898a2c83..c430e33678 100644
--- a/store/cache/cache_test.go
+++ b/store/cache/cache_test.go
@@ -25,6 +25,18 @@ func TestGetOrSetStoreCache(t *testing.T) {
 	require.Equal(t, store2, mngr.GetStoreCache(sKey, store))
 }
 
+func TestUnwrap(t *testing.T) {
+	db := dbm.NewMemDB()
+	mngr := cache.NewCommitKVStoreCacheManager(cache.DefaultCommitKVStoreCacheSize)
+
+	sKey := types.NewKVStoreKey("test")
+	store := iavlstore.UnsafeNewStore(iavl.NewMutableTree(db, 100), 10, 10)
+	_ = mngr.GetStoreCache(sKey, store)
+
+	require.Equal(t, store, mngr.Unwrap(sKey))
+	require.Nil(t, mngr.Unwrap(types.NewKVStoreKey("test2")))
+}
+
 func TestStoreCache(t *testing.T) {
 	db := dbm.NewMemDB()
 	mngr := cache.NewCommitKVStoreCacheManager(cache.DefaultCommitKVStoreCacheSize)
diff --git a/store/rootmulti/store.go b/store/rootmulti/store.go
index af3b03039e..92736c0b8c 100644
--- a/store/rootmulti/store.go
+++ b/store/rootmulti/store.go
@@ -321,6 +321,14 @@ func (rs *Store) CacheMultiStoreWithVersion(version int64) (types.CacheMultiStor
 	for key, store := range rs.stores {
 		switch store.GetStoreType() {
 		case types.StoreTypeIAVL:
+			// If the store is wrapped with an inter-block cache, we must first unwrap
+			// it to get the underlying IAVL store.
+			if rs.interBlockCache != nil {
+				if ckvs := rs.interBlockCache.Unwrap(key); ckvs != nil {
+					store = ckvs
+				}
+			}
+
 			// Attempt to lazy-load an already saved IAVL store version. If the
 			// version does not exist or is pruned, an error should be returned.
 			iavlStore, err := store.(*iavl.Store).GetImmutable(version)
diff --git a/store/types/store.go b/store/types/store.go
index 8dbc0e3107..5b5e647425 100644
--- a/store/types/store.go
+++ b/store/types/store.go
@@ -340,6 +340,9 @@ type MultiStorePersistentCache interface {
 	// cache.
 	GetStoreCache(key StoreKey, store CommitKVStore) CommitKVStore
 
+	// Return the underlying CommitKVStore for a StoreKey.
+	Unwrap(key StoreKey) CommitKVStore
+
 	// Reset the entire set of internal caches.
 	Reset()
 }