From 45aa49859e2963681a73272073117d65298a2f16 Mon Sep 17 00:00:00 2001
From: mmsqe <mavis@crypto.com>
Date: Mon, 22 Jul 2024 11:44:46 +0800
Subject: [PATCH] add GetObjKVStore

---
 store/rootmulti/store.go | 9 +++++++++
 versiondb/multistore.go  | 5 +++++
 2 files changed, 14 insertions(+)

diff --git a/store/rootmulti/store.go b/store/rootmulti/store.go
index 205f2f4b17..f4fa26a27c 100644
--- a/store/rootmulti/store.go
+++ b/store/rootmulti/store.go
@@ -246,6 +246,15 @@ func (rs *Store) GetKVStore(key types.StoreKey) types.KVStore {
 	return s
 }
 
+// GetObjKVStore returns a mounted ObjKVStore for a given StoreKey.
+func (rs *Store) GetObjKVStore(key types.StoreKey) types.ObjKVStore {
+	s, ok := rs.GetStore(key).(types.ObjKVStore)
+	if !ok {
+		panic(fmt.Sprintf("store with key %v is not KVStore", key))
+	}
+	return s
+}
+
 // Implements interface MultiStore
 func (rs *Store) TracingEnabled() bool {
 	return false
diff --git a/versiondb/multistore.go b/versiondb/multistore.go
index a88850d0d3..5749a72b9c 100644
--- a/versiondb/multistore.go
+++ b/versiondb/multistore.go
@@ -99,6 +99,11 @@ func (s *MultiStore) GetKVStore(storeKey types.StoreKey) types.KVStore {
 	return s.GetStore(storeKey).(types.KVStore)
 }
 
+// GetObjKVStore implements `MultiStore` interface
+func (s *MultiStore) GetObjKVStore(storeKey types.StoreKey) types.ObjKVStore {
+	return s.GetStore(storeKey).(types.ObjKVStore)
+}
+
 // SetTracer sets the tracer for the MultiStore that the underlying
 // stores will utilize to trace operations. A MultiStore is returned.
 func (s *MultiStore) SetTracer(w io.Writer) types.MultiStore {