From 2d766b8f201fcdad4170874938205e7daa0a46f0 Mon Sep 17 00:00:00 2001
From: phuslu <phus.lu@gmail.com>
Date: Fri, 22 Mar 2024 22:10:56 +0800
Subject: [PATCH] simplify shard option

---
 lru_cache.go |  3 ---
 options.go   | 24 ++++++++++--------------
 2 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/lru_cache.go b/lru_cache.go
index d985e1f..4ddc5c0 100644
--- a/lru_cache.go
+++ b/lru_cache.go
@@ -4,7 +4,6 @@
 package lru
 
 import (
-	"errors"
 	"sync/atomic"
 	"unsafe"
 )
@@ -74,8 +73,6 @@ func (c *LRUCache[K, V]) Get(key K) (value V, ok bool) {
 	return (*lrushard[K, V])(unsafe.Add(unsafe.Pointer(&c.shards[0]), uintptr(hash&c.mask)*unsafe.Sizeof(c.shards[0]))).Get(hash, key)
 }
 
-var ErrLoaderIsNil = errors.New("loader is nil")
-
 // GetOrLoad returns value for key, call loader function by singleflight if value was not in cache.
 func (c *LRUCache[K, V]) GetOrLoad(key K) (value V, err error, ok bool) {
 	hash := uint32(c.hasher(noescape(unsafe.Pointer(&key)), c.seed))
diff --git a/options.go b/options.go
index dd7fac2..bcbd199 100644
--- a/options.go
+++ b/options.go
@@ -1,6 +1,7 @@
 package lru
 
 import (
+	"errors"
 	"runtime"
 	"time"
 	"unsafe"
@@ -21,32 +22,25 @@ type shardsOption[K comparable, V any] struct {
 	count uint32
 }
 
-func (o *shardsOption[K, V]) ApplyToLRUCache(c *LRUCache[K, V]) {
+func (o *shardsOption[K, V]) getcount(maxcount uint32) uint32 {
 	var shardcount uint32
 	if o.count == 0 {
 		shardcount = nextPowOf2(uint32(runtime.GOMAXPROCS(0) * 16))
 	} else {
 		shardcount = nextPowOf2(o.count)
 	}
-	if maxcount := uint32(len(c.shards)); shardcount > maxcount {
+	if shardcount > maxcount {
 		shardcount = maxcount
 	}
+	return shardcount
+}
 
-	c.mask = uint32(shardcount) - 1
+func (o *shardsOption[K, V]) ApplyToLRUCache(c *LRUCache[K, V]) {
+	c.mask = o.getcount(uint32(len(c.shards))) - 1
 }
 
 func (o *shardsOption[K, V]) ApplyToTTLCache(c *TTLCache[K, V]) {
-	var shardcount uint32
-	if o.count == 0 {
-		shardcount = nextPowOf2(uint32(runtime.GOMAXPROCS(0) * 16))
-	} else {
-		shardcount = nextPowOf2(o.count)
-	}
-	if maxcount := uint32(len(c.shards)); shardcount > maxcount {
-		shardcount = maxcount
-	}
-
-	c.mask = uint32(shardcount) - 1
+	c.mask = o.getcount(uint32(len(c.shards))) - 1
 }
 
 // WithHasher specifies the hasher function of cache.
@@ -85,6 +79,8 @@ func (o *slidingOption[K, V]) ApplyToTTLCache(c *TTLCache[K, V]) {
 	}
 }
 
+var ErrLoaderIsNil = errors.New("loader is nil")
+
 // WithLoader specifies that loader function of LoadingCache.
 func WithLoader[K comparable, V any, Loader func(key K) (value V, err error) | func(key K) (value V, ttl time.Duration, err error)](loader Loader) Option[K, V] {
 	return &loaderOption[K, V]{loader: loader}