From 1c80a4793160f21afb835c0986a35ee2c0ce08c5 Mon Sep 17 00:00:00 2001
From: yihuang <huang@crypto.com>
Date: Wed, 3 Apr 2024 09:43:14 +0800
Subject: [PATCH] feat: save some memory allocations on un-used cache store
 (#248)

* feat: save some memory allocations on un-used cache store

* Update CHANGELOG.md

Signed-off-by: yihuang <huang@crypto.com>

* Update store/internal/btree/btree.go

Signed-off-by: yihuang <huang@crypto.com>

---------

Signed-off-by: yihuang <huang@crypto.com>
---
 CHANGELOG.md                        |  1 +
 store/internal/btree/btree.go       | 36 +++++++++++++++++++++++------
 store/internal/btree/memiterator.go |  9 ++++++++
 3 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a69702176003..28720c3441ed 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -44,6 +44,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (baseapp) [#206](https://github.com/crypto-org-chain/cosmos-sdk/pull/206) Support mount object store in baseapp, add `ObjectStore` api in context.
 * (bank) [#237](https://github.com/crypto-org-chain/cosmos-sdk/pull/237) Support virtual accounts in sending coins.
 * [#243](https://github.com/crypto-org-chain/cosmos-sdk/pull/243) Support `RunAtomic` API in `Context` to use new CoW branched cache store.
+* [#248](https://github.com/crypto-org-chain/cosmos-sdk/pull/248) Init btree store lazily to save allocations when no content insert into it.
 
 ## [Unreleased-Upstream]
 
diff --git a/store/internal/btree/btree.go b/store/internal/btree/btree.go
index bef4309412ff..920b373ff491 100644
--- a/store/internal/btree/btree.go
+++ b/store/internal/btree/btree.go
@@ -28,25 +28,38 @@ type BTree[V any] struct {
 
 // NewBTree creates a wrapper around `btree.BTreeG`.
 func NewBTree[V any]() BTree[V] {
-	return BTree[V]{
-		tree: btree.NewBTreeGOptions(byKeys[V], btree.Options{
-			Degree:  bTreeDegree,
-			NoLocks: false,
-		}),
+	return BTree[V]{}
+}
+
+func (bt *BTree[V]) init() {
+	if bt.tree != nil {
+		return
 	}
+	bt.tree = btree.NewBTreeGOptions(byKeys[V], btree.Options{
+		Degree:  bTreeDegree,
+		NoLocks: false,
+	})
 }
 
 // Set supports nil as value when used as overlay
-func (bt BTree[V]) Set(key []byte, value V) {
+func (bt *BTree[V]) Set(key []byte, value V) {
+	bt.init()
 	bt.tree.Set(newItem(key, value))
 }
 
 func (bt BTree[V]) Get(key []byte) (V, bool) {
+	if bt.tree == nil {
+		var zero V
+		return zero, false
+	}
 	i, found := bt.tree.Get(newItemWithKey[V](key))
 	return i.value, found
 }
 
-func (bt BTree[V]) Delete(key []byte) {
+func (bt *BTree[V]) Delete(key []byte) {
+	if bt.tree == nil {
+		return
+	}
 	bt.tree.Delete(newItemWithKey[V](key))
 }
 
@@ -67,16 +80,25 @@ func (bt BTree[V]) ReverseIterator(start, end []byte) (types.GIterator[V], error
 // Copy the tree. This is a copy-on-write operation and is very fast because
 // it only performs a shadowed copy.
 func (bt BTree[V]) Copy() BTree[V] {
+	if bt.tree == nil {
+		return BTree[V]{}
+	}
 	return BTree[V]{
 		tree: bt.tree.Copy(),
 	}
 }
 
 func (bt BTree[V]) Clear() {
+	if bt.tree == nil {
+		return
+	}
 	bt.tree.Clear()
 }
 
 func (bt BTree[V]) Scan(cb func(key []byte, value V) bool) {
+	if bt.tree == nil {
+		return
+	}
 	bt.tree.Scan(func(i item[V]) bool {
 		return cb(i.key, i.value)
 	})
diff --git a/store/internal/btree/memiterator.go b/store/internal/btree/memiterator.go
index e98ae9cd834e..b3caabd4584b 100644
--- a/store/internal/btree/memiterator.go
+++ b/store/internal/btree/memiterator.go
@@ -24,6 +24,15 @@ type memIterator[V any] struct {
 }
 
 func newMemIterator[V any](start, end []byte, items BTree[V], ascending bool) *memIterator[V] {
+	if items.tree == nil {
+		return &memIterator[V]{
+			start:     start,
+			end:       end,
+			ascending: ascending,
+			valid:     false,
+		}
+	}
+
 	var (
 		valid bool
 		empty V