[#47] Add InitialCapacity function

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
 }
 

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()

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)
 }

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)
 	}

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),

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
 }

internal/hashtable/map_test.go

@@ -121,19 +121,19 @@ 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 {
 		// We intentionally use an awful hash function here to make sure
 		// 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)