From bc3f63e815ab8e7177a388c89f1107444e6182c7 Mon Sep 17 00:00:00 2001
From: pdrobnjak <drobnjakpavle95@gmail.com>
Date: Wed, 18 Feb 2026 14:43:40 +0100
Subject: [PATCH] perf: hollow CMS for OCC prepareTask

CacheMultiStoreForOCC creates a hollow CMS where stores are populated
directly from VersionIndexedStores, skipping creation of intermediate
cachekv instances that would be immediately discarded by SetKVStores.

Cherry-picked from de0dd9cc4, adapted for this branch (no fast store),
GC tuning changes excluded.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 giga/deps/tasks/scheduler.go         | 36 ++++++++++++++++++++--------
 sei-cosmos/store/cachemulti/store.go | 33 +++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/giga/deps/tasks/scheduler.go b/giga/deps/tasks/scheduler.go
index 16c97e3fdb..4170526e99 100644
--- a/giga/deps/tasks/scheduler.go
+++ b/giga/deps/tasks/scheduler.go
@@ -490,9 +490,6 @@ func (s *scheduler) prepareTask(task *deliverTxTask) {
 
 	// if there are no stores, don't try to wrap, because there's nothing to wrap
 	if len(s.multiVersionStores) > 0 {
-		// non-blocking
-		cms := ctx.MultiStore().CacheMultiStore()
-
 		// init version stores by store key
 		vs := make(map[store.StoreKey]*multiversion.VersionIndexedStore)
 		for storeKey, mvs := range s.multiVersionStores {
@@ -501,14 +498,33 @@ func (s *scheduler) prepareTask(task *deliverTxTask) {
 
 		// save off version store so we can ask it things later
 		task.VersionStores = vs
-		ms := cms.SetKVStores(func(k store.StoreKey, kvs sdk.KVStore) store.CacheWrap {
-			return vs[k]
-		})
-		ms = ms.(store.GigaMultiStore).SetGigaKVStores(func(k store.StoreKey, kvs sdk.KVStore) store.KVStore {
-			return vs[k]
-		})
 
-		ctx = ctx.WithMultiStore(ms)
+		// Try the fast OCC path: create a hollow CMS with stores populated
+		// directly from VersionIndexedStores, skipping intermediate store creation.
+		type occPreparer interface {
+			CacheMultiStoreForOCC(
+				func(store.StoreKey) store.CacheWrap,
+				func(store.StoreKey) store.KVStore,
+			) store.CacheMultiStore
+		}
+		parentMS := ctx.MultiStore()
+		if occ, ok := parentMS.(occPreparer); ok {
+			ms := occ.CacheMultiStoreForOCC(
+				func(k store.StoreKey) store.CacheWrap { return vs[k] },
+				func(k store.StoreKey) store.KVStore { return vs[k] },
+			)
+			ctx = ctx.WithMultiStore(ms)
+		} else {
+			// Fallback: create full CMS and replace stores
+			cms := parentMS.CacheMultiStore()
+			ms := cms.SetKVStores(func(k store.StoreKey, kvs sdk.KVStore) store.CacheWrap {
+				return vs[k]
+			})
+			ms = ms.(store.GigaMultiStore).SetGigaKVStores(func(k store.StoreKey, kvs sdk.KVStore) store.KVStore {
+				return vs[k]
+			})
+			ctx = ctx.WithMultiStore(ms)
+		}
 	}
 
 	if task.TxTracer != nil {
diff --git a/sei-cosmos/store/cachemulti/store.go b/sei-cosmos/store/cachemulti/store.go
index eb4f8bce05..14deb7f1e3 100644
--- a/sei-cosmos/store/cachemulti/store.go
+++ b/sei-cosmos/store/cachemulti/store.go
@@ -228,6 +228,39 @@ func (cms Store) CacheMultiStore() types.CacheMultiStore {
 	return newCacheMultiStoreFromCMS(cms)
 }
 
+// CacheMultiStoreForOCC creates a hollow CMS where all stores are directly
+// provided by the caller via handler functions. This skips creating
+// intermediate cachekv instances that would be immediately discarded when
+// the OCC scheduler replaces all stores with VersionIndexedStores.
+func (cms Store) CacheMultiStoreForOCC(
+	kvHandler func(sk types.StoreKey) types.CacheWrap,
+	gigaHandler func(sk types.StoreKey) types.KVStore,
+) types.CacheMultiStore {
+	// Use cms.keys as the canonical set of store keys — no need to
+	// materialize parent stores just to discover which keys exist.
+	result := Store{
+		db:              cachekv.NewStore(cms.db, nil, types.DefaultCacheSizeLimit),
+		stores:          make(map[types.StoreKey]types.CacheWrap, len(cms.keys)),
+		parents:         make(map[types.StoreKey]types.CacheWrapper),
+		keys:            cms.keys,
+		gigaKeys:        cms.gigaKeys,
+		gigaStores:      make(map[types.StoreKey]types.KVStore, len(cms.gigaKeys)),
+		traceWriter:     cms.traceWriter,
+		traceContext:    cms.traceContext,
+		mu:              &sync.RWMutex{},
+		materializeOnce: &sync.Once{},
+	}
+
+	for _, k := range cms.keys {
+		result.stores[k] = kvHandler(k)
+	}
+	for _, k := range cms.gigaKeys {
+		result.gigaStores[k] = gigaHandler(k)
+	}
+
+	return result
+}
+
 // CacheMultiStoreWithVersion implements the MultiStore interface. It will panic
 // as an already cached multi-store cannot load previous versions.
 //