From bcd9697c61983d8193bc430ae1bd424ab63d101c Mon Sep 17 00:00:00 2001
From: XuShuo <xushuo@xiaohongshu.net>
Date: Mon, 2 Dec 2024 23:31:16 +0800
Subject: [PATCH] feat: Add NewOrderedMapWithCapacity function for initializing
 an orderedmap with enough pre-allocatDed space

---
 orderedmap.go         |  8 ++++++++
 orderedmap_test.go    | 15 +++++++++++++++
 v2/orderedmap.go      |  8 ++++++++
 v2/orderedmap_test.go | 15 +++++++++++++++
 4 files changed, 46 insertions(+)

diff --git a/orderedmap.go b/orderedmap.go
index 9907ddf..d14f361 100644
--- a/orderedmap.go
+++ b/orderedmap.go
@@ -11,6 +11,14 @@ func NewOrderedMap() *OrderedMap {
 	}
 }
 
+// NewOrderedMapWithCapacity creates a map with enough pre-allocated space to
+// hold the specified number of elements.
+func NewOrderedMapWithCapacity(capacity int) *OrderedMap {
+	return &OrderedMap{
+		kv: make(map[interface{}]*Element, capacity),
+	}
+}
+
 // Get returns the value for a key. If the key does not exist, the second return
 // parameter will be false and the value will be nil.
 func (m *OrderedMap) Get(key interface{}) (interface{}, bool) {
diff --git a/orderedmap_test.go b/orderedmap_test.go
index 519512d..713aec6 100644
--- a/orderedmap_test.go
+++ b/orderedmap_test.go
@@ -761,6 +761,21 @@ func BenchmarkBigOrderedMap_Set(b *testing.B) {
 	benchmarkBigOrderedMap_Set()(b)
 }
 
+func benchmarkBigMapWithCapacity_Set() func(b *testing.B) {
+	return func(b *testing.B) {
+		for j := 0; j < b.N; j++ {
+			m := orderedmap.NewOrderedMapWithCapacity(10000000)
+			for i := 0; i < 10000000; i++ {
+				m.Set(i, true)
+			}
+		}
+	}
+}
+
+func BenchmarkBigMapWithCapacity_Set(b *testing.B) {
+	benchmarkBigMapWithCapacity_Set()(b)
+}
+
 func benchmarkBigMap_Get() func(b *testing.B) {
 	m := make(map[int]bool)
 	for i := 0; i < 10000000; i++ {
diff --git a/v2/orderedmap.go b/v2/orderedmap.go
index a9f665d..f6ab5f3 100644
--- a/v2/orderedmap.go
+++ b/v2/orderedmap.go
@@ -11,6 +11,14 @@ func NewOrderedMap[K comparable, V any]() *OrderedMap[K, V] {
 	}
 }
 
+// NewOrderedMapWithCapacity creates a map with enough pre-allocated space to
+// hold the specified number of elements.
+func NewOrderedMapWithCapacity[K comparable, V any](capacity int) *OrderedMap[K, V] {
+	return &OrderedMap[K, V]{
+		kv: make(map[K]*Element[K, V], capacity),
+	}
+}
+
 func NewOrderedMapWithElements[K comparable, V any](els ...*Element[K, V]) *OrderedMap[K, V] {
 	om := &OrderedMap[K, V]{
 		kv: make(map[K]*Element[K, V], len(els)),
diff --git a/v2/orderedmap_test.go b/v2/orderedmap_test.go
index bd0982d..d39c264 100644
--- a/v2/orderedmap_test.go
+++ b/v2/orderedmap_test.go
@@ -730,6 +730,21 @@ func BenchmarkBigOrderedMap_Set(b *testing.B) {
 	benchmarkBigOrderedMap_Set()(b)
 }
 
+func benchmarkBigMapWithCapacity_Set() func(b *testing.B) {
+	return func(b *testing.B) {
+		for j := 0; j < b.N; j++ {
+			m := orderedmap.NewOrderedMapWithCapacity[int, bool](10000000)
+			for i := 0; i < 10000000; i++ {
+				m.Set(i, true)
+			}
+		}
+	}
+}
+
+func BenchmarkBigMapWithCapacity_Set(b *testing.B) {
+	benchmarkBigMapWithCapacity_Set()(b)
+}
+
 func benchmarkBigMap_Get() func(b *testing.B) {
 	m := make(map[int]bool)
 	for i := 0; i < 10000000; i++ {