Repository: maypok86/otter (go)

Issue:
Add InitialCapacity function to builder
# Clarification and motivation

Since otter supports cost-based eviction, we cannot pre-initialize internal structures, but in reality this feature is rarely used, and other users may not be happy because otter allocated more memory than it could. So we would like to add an `InitialCapacity` function with which users could preallocate otter's internal data structures.

The usage is envisioned like this
```go
c, err := otter.MustBuilder[int, int](size).
    InitialCapacity(size).
    WithTTL(time.Second).
    Build()
```

# Acceptance criteria

1. The `InitialCapacity` function is added to the builder.
2. Implementation tests are written


Pull request:
[#47] Add InitialCapacity function
## Description

Added `InitialCapacity` function and function for creating a hash table of the required size.

## Related issue(s)

- Resolves #47 

## :white_check_mark: Checklist for your Pull Request

<!--
Ideally a PR has all of the checkmarks set.

If something in this list is irrelevant to your PR, you should still set this
checkmark indicating that you are sure it is dealt with (be that by irrelevance).

If you don't set a checkmark (e. g. don't add a test for new functionality),
you must be able to justify that.
-->

#### Related changes (conditional)

- Tests
    - [x] If I added new functionality, I added tests covering it.
    - [x] If I fixed a bug, I added a regression test to prevent the bug from
      silently reappearing again.

- Documentation
    - [x] I checked whether I should update the docs and did so if necessary:
        - [README](../README.md)

- Public contracts
    - [x] My changes doesn't break project license.

#### Stylistic guide (mandatory)

- [x] My code complies with the [styles guide](https://github.com/uber-go/guide/blob/master/style.md).
- [x] My commit history is clean (only contains changes relating to my
  issue/pull request and no reverted-my-earlier-commit changes) and commit
  messages start with identifiers of related issues in square brackets.

  **Example:** `[#42] Short commit description`

  If necessary both of these can be achieved even after the commits have been
  made/pushed using [rebase and squash](https://git-scm.com/docs/git-rebase).

#### Before merging (mandatory)
- [x] Check __target__  branch of PR is set correctly


Patch:
diff --git a/builder.go b/builder.go
index d8542bd..aba7b6a 100644
--- a/builder.go
+++ b/builder.go
@@ -21,36 +21,60 @@ import (
 	"github.com/maypok86/otter/internal/core"
 )
 
+const (
+	unsetCapacity = -1
+)
+
 var (
 	// ErrIllegalCapacity means that a non-positive capacity has been passed to the NewBuilder.
 	ErrIllegalCapacity = errors.New("capacity should be positive")
+	// ErrIllegalInitialCapacity means that a non-positive capacity has been passed to the Builder.InitialCapacity.
+	ErrIllegalInitialCapacity = errors.New("initial capacity should be positive")
+	// ErrNilCostFunc means that a nil cost func has been passed to the Builder.Cost.
+	ErrNilCostFunc = errors.New("setCostFunc func should not be nil")
 	// ErrIllegalTTL means that a non-positive ttl has been passed to the Builder.WithTTL.
 	ErrIllegalTTL = errors.New("ttl should be positive")
 )
 
 type baseOptions[K comparable, V any] struct {
-	capacity     int
-	statsEnabled bool
-	costFunc     func(key K, value V) uint32
+	capacity        int
+	initialCapacity int
+	statsEnabled    bool
+	costFunc        func(key K, value V) uint32
 }
 
 func (o *baseOptions[K, V]) collectStats() {
 	o.statsEnabled = true
 }
 
-func (o *baseOptions[K, V]) cost(costFunc func(key K, value V) uint32) {
+func (o *baseOptions[K, V]) setCostFunc(costFunc func(key K, value V) uint32) {
 	o.costFunc = costFunc
 }
 
+func (o *baseOptions[K, V]) setInitialCapacity(initialCapacity int) {
+	o.initialCapacity = initialCapacity
+}
+
 func (o *baseOptions[K, V]) validate() error {
+	if o.initialCapacity <= 0 && o.initialCapacity != unsetCapacity {
+		return ErrIllegalInitialCapacity
+	}
+	if o.costFunc == nil {
+		return ErrNilCostFunc
+	}
 	return nil
 }
 
 func (o *baseOptions[K, V]) toConfig() core.Config[K, V] {
+	var initialCapacity *int
+	if o.initialCapacity != unsetCapacity {
+		initialCapacity = &o.initialCapacity
+	}
 	return core.Config[K, V]{
-		Capacity:     o.capacity,
-		StatsEnabled: o.statsEnabled,
-		CostFunc:     o.costFunc,
+		Capacity:        o.capacity,
+		InitialCapacity: initialCapacity,
+		StatsEnabled:    o.statsEnabled,
+		CostFunc:        o.costFunc,
 	}
 }
 
@@ -108,8 +132,9 @@ func NewBuilder[K comparable, V any](capacity int) (*Builder[K, V], error) {
 
 	return &Builder[K, V]{
 		baseOptions: baseOptions[K, V]{
-			capacity:     capacity,
-			statsEnabled: false,
+			capacity:        capacity,
+			initialCapacity: unsetCapacity,
+			statsEnabled:    false,
 			costFunc: func(key K, value V) uint32 {
 				return 1
 			},
@@ -125,11 +150,19 @@ func (b *Builder[K, V]) CollectStats() *Builder[K, V] {
 	return b
 }
 
+// InitialCapacity sets the minimum total size for the internal data structures. Providing a large enough estimate
+// at construction time avoids the need for expensive resizing operations later, but setting this
+// value unnecessarily high wastes memory.
+func (b *Builder[K, V]) InitialCapacity(initialCapacity int) *Builder[K, V] {
+	b.setInitialCapacity(initialCapacity)
+	return b
+}
+
 // Cost sets a function to dynamically calculate the cost of an item.
 //
 // By default, this function always returns 1.
 func (b *Builder[K, V]) Cost(costFunc func(key K, value V) uint32) *Builder[K, V] {
-	b.cost(costFunc)
+	b.setCostFunc(costFunc)
 	return b
 }
 
@@ -179,11 +212,19 @@ func (b *ConstTTLBuilder[K, V]) CollectStats() *ConstTTLBuilder[K, V] {
 	return b
 }
 
+// InitialCapacity sets the minimum total size for the internal data structures. Providing a large enough estimate
+// at construction time avoids the need for expensive resizing operations later, but setting this
+// value unnecessarily high wastes memory.
+func (b *ConstTTLBuilder[K, V]) InitialCapacity(initialCapacity int) *ConstTTLBuilder[K, V] {
+	b.setInitialCapacity(initialCapacity)
+	return b
+}
+
 // Cost sets a function to dynamically calculate the cost of an item.
 //
 // By default, this function always returns 1.
 func (b *ConstTTLBuilder[K, V]) Cost(costFunc func(key K, value V) uint32) *ConstTTLBuilder[K, V] {
-	b.cost(costFunc)
+	b.setCostFunc(costFunc)
 	return b
 }
 
@@ -210,11 +251,19 @@ func (b *VariableTTLBuilder[K, V]) CollectStats() *VariableTTLBuilder[K, V] {
 	return b
 }
 
+// InitialCapacity sets the minimum total size for the internal data structures. Providing a large enough estimate
+// at construction time avoids the need for expensive resizing operations later, but setting this
+// value unnecessarily high wastes memory.
+func (b *VariableTTLBuilder[K, V]) InitialCapacity(initialCapacity int) *VariableTTLBuilder[K, V] {
+	b.setInitialCapacity(initialCapacity)
+	return b
+}
+
 // Cost sets a function to dynamically calculate the cost of an item.
 //
 // By default, this function always returns 1.
 func (b *VariableTTLBuilder[K, V]) Cost(costFunc func(key K, value V) uint32) *VariableTTLBuilder[K, V] {
-	b.cost(costFunc)
+	b.setCostFunc(costFunc)
 	return b
 }
 
diff --git a/cache.go b/cache.go
index 23d66af..f2e017f 100644
--- a/cache.go
+++ b/cache.go
@@ -125,7 +125,7 @@ func newCache[K comparable, V any](c core.Config[K, V]) Cache[K, V] {
 
 // Set associates the value with the key in this cache.
 //
-// If it returns false, then the key-value item had too much cost and the Set was dropped.
+// If it returns false, then the key-value item had too much setCostFunc and the Set was dropped.
 func (c Cache[K, V]) Set(key K, value V) bool {
 	return c.cache.Set(key, value)
 }
@@ -134,7 +134,7 @@ func (c Cache[K, V]) Set(key K, value V) bool {
 //
 // If the specified key is not already associated with a value, then it returns false.
 //
-// Also, it returns false if the key-value item had too much cost and the SetIfAbsent was dropped.
+// Also, it returns false if the key-value item had too much setCostFunc and the SetIfAbsent was dropped.
 func (c Cache[K, V]) SetIfAbsent(key K, value V) bool {
 	return c.cache.SetIfAbsent(key, value)
 }
@@ -153,7 +153,7 @@ func newCacheWithVariableTTL[K comparable, V any](c core.Config[K, V]) CacheWith
 
 // Set associates the value with the key in this cache and sets the custom ttl for this key-value item.
 //
-// If it returns false, then the key-value item had too much cost and the Set was dropped.
+// If it returns false, then the key-value item had too much setCostFunc and the Set was dropped.
 func (c CacheWithVariableTTL[K, V]) Set(key K, value V, ttl time.Duration) bool {
 	return c.cache.SetWithTTL(key, value, ttl)
 }
@@ -163,7 +163,7 @@ func (c CacheWithVariableTTL[K, V]) Set(key K, value V, ttl time.Duration) bool
 //
 // If the specified key is not already associated with a value, then it returns false.
 //
-// Also, it returns false if the key-value item had too much cost and the SetIfAbsent was dropped.
+// Also, it returns false if the key-value item had too much setCostFunc and the SetIfAbsent was dropped.
 func (c CacheWithVariableTTL[K, V]) SetIfAbsent(key K, value V, ttl time.Duration) bool {
 	return c.cache.SetIfAbsentWithTTL(key, value, ttl)
 }
diff --git a/internal/core/cache.go b/internal/core/cache.go
index ece04a2..cf5cff6 100644
--- a/internal/core/cache.go
+++ b/internal/core/cache.go
@@ -43,6 +43,7 @@ func getExpiration(ttl time.Duration) uint32 {
 // Config is a set of cache settings.
 type Config[K comparable, V any] struct {
 	Capacity        int
+	InitialCapacity *int
 	StatsEnabled    bool
 	TTL             *time.Duration
 	WithVariableTTL bool
@@ -81,8 +82,15 @@ func NewCache[K comparable, V any](c Config[K, V]) *Cache[K, V] {
 		readBuffers = append(readBuffers, lossy.New[node.Node[K, V]]())
 	}
 
+	var hashmap *hashtable.Map[K, V]
+	if c.InitialCapacity == nil {
+		hashmap = hashtable.New[K, V]()
+	} else {
+		hashmap = hashtable.NewWithSize[K, V](*c.InitialCapacity)
+	}
+
 	cache := &Cache[K, V]{
-		hashmap:     hashtable.New[K, V](),
+		hashmap:     hashmap,
 		policy:      s3fifo.NewPolicy[K, V](uint32(c.Capacity)),
 		readBuffers: readBuffers,
 		writeBuffer: queue.NewMPSC[node.WriteTask[K, V]](writeBufferCapacity),
diff --git a/internal/hashtable/map.go b/internal/hashtable/map.go
index 5e9df24..839301f 100644
--- a/internal/hashtable/map.go
+++ b/internal/hashtable/map.go
@@ -18,6 +18,7 @@ import (
 	"github.com/dolthub/maphash"
 
 	"github.com/maypok86/otter/internal/node"
+	"github.com/maypok86/otter/internal/xmath"
 	"github.com/maypok86/otter/internal/xruntime"
 )
 
@@ -119,11 +120,28 @@ type paddedCounter struct {
 	counter
 }
 
+// NewWithSize creates a new Map instance with capacity enough
+// to hold size nodes. If size is zero or negative, the value
+// is ignored.
+func NewWithSize[K comparable, V any](size int) *Map[K, V] {
+	return newMap[K, V](size)
+}
+
 // New creates a new Map instance.
 func New[K comparable, V any]() *Map[K, V] {
+	return newMap[K, V](minNodeCount)
+}
+
+func newMap[K comparable, V any](size int) *Map[K, V] {
 	m := &Map[K, V]{}
 	m.resizeCond = *sync.NewCond(&m.resizeMutex)
-	t := newTable(minBucketCount, maphash.NewHasher[K]())
+	var t *table[K]
+	if size <= minNodeCount {
+		t = newTable(minBucketCount, maphash.NewHasher[K]())
+	} else {
+		bucketCount := xmath.RoundUpPowerOf2(uint32(size / bucketSize))
+		t = newTable(int(bucketCount), maphash.NewHasher[K]())
+	}
 	atomic.StorePointer(&m.table, unsafe.Pointer(t))
 	return m
 }


Test patch:
diff --git a/builder_test.go b/builder_test.go
index 063928e..32b34dd 100644
--- a/builder_test.go
+++ b/builder_test.go
@@ -38,17 +38,47 @@ func TestBuilder_NewFailed(t *testing.T) {
 		t.Fatalf("should fail with an error %v, but got %v", ErrIllegalCapacity, err)
 	}
 
-	_, err = MustBuilder[int, int](100).WithTTL(-1).Build()
+	capacity := 100
+	// negative const ttl
+	_, err = MustBuilder[int, int](capacity).WithTTL(-1).Build()
 	if err == nil || !errors.Is(err, ErrIllegalTTL) {
 		t.Fatalf("should fail with an error %v, but got %v", ErrIllegalTTL, err)
 	}
+
+	// negative initial capacity
+	_, err = MustBuilder[int, int](capacity).InitialCapacity(-2).Build()
+	if err == nil || !errors.Is(err, ErrIllegalInitialCapacity) {
+		t.Fatalf("should fail with an error %v, but got %v", ErrIllegalInitialCapacity, err)
+	}
+
+	_, err = MustBuilder[int, int](capacity).WithTTL(time.Hour).InitialCapacity(0).Build()
+	if err == nil || !errors.Is(err, ErrIllegalInitialCapacity) {
+		t.Fatalf("should fail with an error %v, but got %v", ErrIllegalInitialCapacity, err)
+	}
+
+	_, err = MustBuilder[int, int](capacity).WithVariableTTL().InitialCapacity(-5).Build()
+	if err == nil || !errors.Is(err, ErrIllegalInitialCapacity) {
+		t.Fatalf("should fail with an error %v, but got %v", ErrIllegalInitialCapacity, err)
+	}
+
+	// nil cost func
+	_, err = MustBuilder[int, int](capacity).Cost(nil).Build()
+	if err == nil || !errors.Is(err, ErrNilCostFunc) {
+		t.Fatalf("should fail with an error %v, but got %v", ErrNilCostFunc, err)
+	}
 }
 
 func TestBuilder_BuildSuccess(t *testing.T) {
 	b := MustBuilder[int, int](10)
 
+	_, err := b.InitialCapacity(unsetCapacity).Build()
+	if err != nil {
+		t.Fatalf("builded cache with error: %v", err)
+	}
+
 	c, err := b.
 		CollectStats().
+		InitialCapacity(10).
 		Cost(func(key int, value int) uint32 {
 			return 2
 		}).Build()
diff --git a/cache_test.go b/cache_test.go
index 7a280bf..db57795 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -133,7 +133,10 @@ func TestCache_SetIfAbsent(t *testing.T) {
 
 func TestCache_SetWithTTL(t *testing.T) {
 	size := 256
-	c, err := MustBuilder[int, int](size).WithTTL(time.Second).Build()
+	c, err := MustBuilder[int, int](size).
+		InitialCapacity(size).
+		WithTTL(time.Second).
+		Build()
 	if err != nil {
 		t.Fatalf("can not create builder: %v", err)
 	}
@@ -184,7 +187,10 @@ func TestCache_SetWithTTL(t *testing.T) {
 
 func TestBaseCache_DeleteByFunc(t *testing.T) {
 	size := 256
-	c, err := MustBuilder[int, int](size).WithTTL(time.Hour).Build()
+	c, err := MustBuilder[int, int](size).
+		InitialCapacity(size).
+		WithTTL(time.Hour).
+		Build()
 	if err != nil {
 		t.Fatalf("can not create builder: %v", err)
 	}
diff --git a/internal/hashtable/map_test.go b/internal/hashtable/map_test.go
index 87acafe..beefefd 100644
--- a/internal/hashtable/map_test.go
+++ b/internal/hashtable/map_test.go
@@ -121,8 +121,8 @@ type hasher struct {
 }
 
 func TestMap_SetWithCollisions(t *testing.T) {
-	const numEntries = 1000
-	m := New[int, int]()
+	const numNodes = 1000
+	m := NewWithSize[int, int](numNodes)
 	table := (*table[int])(atomic.LoadPointer(&m.table))
 	hasher := (*hasher)((unsafe.Pointer)(&table.hasher))
 	hasher.hash = func(ptr unsafe.Pointer, seed uintptr) uintptr {
@@ -130,10 +130,10 @@ func TestMap_SetWithCollisions(t *testing.T) {
 		// that the map copes with key collisions.
 		return 42
 	}
-	for i := 0; i < numEntries; i++ {
+	for i := 0; i < numNodes; i++ {
 		m.Set(newNode(i, i))
 	}
-	for i := 0; i < numEntries; i++ {
+	for i := 0; i < numNodes; i++ {
 		v, ok := m.Get(i)
 		if !ok {
 			t.Fatalf("value not found for %d", i)