From de817e5b419d1a040dceac3dab6433a60db2dd83 Mon Sep 17 00:00:00 2001
From: Aleksandr Bezobchuk <aleks.bezobchuk@gmail.com>
Date: Wed, 4 Sep 2019 15:52:38 -0400
Subject: [PATCH 1/3] Fix CacheMultiStoreWithVersion

---
 store/cache/cache.go     | 9 +++++++++
 store/rootmulti/store.go | 6 ++++++
 store/types/store.go     | 3 +++
 3 files changed, 18 insertions(+)

diff --git a/store/cache/cache.go b/store/cache/cache.go
index 2293c3bbde3e..9eea90ad942e 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/rootmulti/store.go b/store/rootmulti/store.go
index af3b03039ef7..c3c7d7efd8a5 100644
--- a/store/rootmulti/store.go
+++ b/store/rootmulti/store.go
@@ -321,6 +321,12 @@ 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 {
+				store = rs.interBlockCache.Unwrap(key)
+			}
+
 			// 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 8dbc0e31074b..5b5e64742552 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()
 }

From 864a8e6ccd79f193a057dd303506e1f47f1f1a9d Mon Sep 17 00:00:00 2001
From: Aleksandr Bezobchuk <aleks.bezobchuk@gmail.com>
Date: Wed, 4 Sep 2019 15:54:49 -0400
Subject: [PATCH 2/3] Add nil check

---
 store/rootmulti/store.go | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/store/rootmulti/store.go b/store/rootmulti/store.go
index c3c7d7efd8a5..92736c0b8c2f 100644
--- a/store/rootmulti/store.go
+++ b/store/rootmulti/store.go
@@ -324,7 +324,9 @@ func (rs *Store) CacheMultiStoreWithVersion(version int64) (types.CacheMultiStor
 			// 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 {
-				store = rs.interBlockCache.Unwrap(key)
+				if ckvs := rs.interBlockCache.Unwrap(key); ckvs != nil {
+					store = ckvs
+				}
 			}
 
 			// Attempt to lazy-load an already saved IAVL store version. If the

From 5b6b9ccc21932ad1cd70e5f4fac3c46d6ad34bd3 Mon Sep 17 00:00:00 2001
From: Aleksandr Bezobchuk <aleks.bezobchuk@gmail.com>
Date: Wed, 4 Sep 2019 15:57:01 -0400
Subject: [PATCH 3/3] Implement unit test

---
 store/cache/cache_test.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/store/cache/cache_test.go b/store/cache/cache_test.go
index e2898a2c83a2..c430e33678fe 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)