feat: add unit field to bitcount (redis#538)

* feat: modify contract and commit * feat: add test * feat: address comments and convert from ptr to string * feat: remove unused util --------- Co-authored-by: Rueian <[email protected]> Co-authored-by: Anuragkillswitch <[email protected]>
SoulPancake · May 18, 2024 · 819cc5e · 819cc5e
1 parent d6aeb96
commit 819cc5e
Show file tree

Hide file tree

Showing 3 changed files with 67 additions and 3 deletions.
diff --git a/rueidiscompat/adapter.go b/rueidiscompat/adapter.go
@@ -43,7 +43,11 @@ import (
 	"github.com/redis/rueidis/internal/util"
 )
 
-const KeepTTL = -1
+const (
+	KeepTTL           = -1
+	BitCountIndexByte = "BYTE"
+	BitCountIndexBit  = "BIT"
+)
 
 type Cmdable interface {
 	Cache(ttl time.Duration) CacheCompat
@@ -1093,11 +1097,23 @@ func (c *Compat) SetBit(ctx context.Context, key string, offset int64, value int
 }
 
 func (c *Compat) BitCount(ctx context.Context, key string, bitCount *BitCount) *IntCmd {
+
 	var resp rueidis.RedisResult
 	if bitCount == nil {
 		resp = c.client.Do(ctx, c.client.B().Bitcount().Key(key).Build())
-	} else {
+		return newIntCmd(resp)
+	}
+
+	if bitCount.Unit == "" {
 		resp = c.client.Do(ctx, c.client.B().Bitcount().Key(key).Start(bitCount.Start).End(bitCount.End).Build())
+		return newIntCmd(resp)
+	}
+
+	switch bitCount.Unit {
+	case BitCountIndexByte:
+		resp = c.client.Do(ctx, c.client.B().Bitcount().Key(key).Start(bitCount.Start).End(bitCount.End).Byte().Build())
+	case BitCountIndexBit:
+		resp = c.client.Do(ctx, c.client.B().Bitcount().Key(key).Start(bitCount.Start).End(bitCount.End).Bit().Build())
 	}
 	return newIntCmd(resp)
 }
@@ -4584,8 +4600,19 @@ func (c CacheCompat) BitCount(ctx context.Context, key string, bitCount *BitCoun
 	var resp rueidis.RedisResult
 	if bitCount == nil {
 		resp = c.client.DoCache(ctx, c.client.B().Bitcount().Key(key).Cache(), c.ttl)
-	} else {
+		return newIntCmd(resp)
+	}
+
+	if bitCount.Unit == "" {
 		resp = c.client.DoCache(ctx, c.client.B().Bitcount().Key(key).Start(bitCount.Start).End(bitCount.End).Cache(), c.ttl)
+		return newIntCmd(resp)
+	}
+
+	switch bitCount.Unit {
+	case BitCountIndexByte:
+		resp = c.client.DoCache(ctx, c.client.B().Bitcount().Key(key).Start(bitCount.Start).End(bitCount.End).Byte().Cache(), c.ttl)
+	case BitCountIndexBit:
+		resp = c.client.DoCache(ctx, c.client.B().Bitcount().Key(key).Start(bitCount.Start).End(bitCount.End).Bit().Cache(), c.ttl)
 	}
 	return newIntCmd(resp)
 }

diff --git a/rueidiscompat/adapter_test.go b/rueidiscompat/adapter_test.go
@@ -1150,6 +1150,24 @@ func testAdapter(resp3 bool) {
 			})
 			Expect(bitCount.Err()).NotTo(HaveOccurred())
 			Expect(bitCount.Val()).To(Equal(int64(6)))
+
+			if resp3 {
+				bitCount = adapter.BitCount(ctx, "key", &BitCount{
+					Start: 1,
+					End:   1,
+					Unit:  "BYTE",
+				})
+				Expect(bitCount.Err()).NotTo(HaveOccurred())
+				Expect(bitCount.Val()).To(Equal(int64(6)))
+
+				bitCount = adapter.BitCount(ctx, "key", &BitCount{
+					Start: 1,
+					End:   1,
+					Unit:  "BIT",
+				})
+				Expect(bitCount.Err()).NotTo(HaveOccurred())
+				Expect(bitCount.Val()).To(Equal(int64(1)))
+			}
 		})
 
 		It("should BitOpAnd", func() {
@@ -7064,6 +7082,24 @@ func testAdapterCache(resp3 bool) {
 			})
 			Expect(bitCount.Err()).NotTo(HaveOccurred())
 			Expect(bitCount.Val()).To(Equal(int64(6)))
+
+			if resp3 {
+				bitCount = adapter.Cache(time.Hour).BitCount(ctx, "key", &BitCount{
+					Start: 1,
+					End:   1,
+					Unit:  "BIT",
+				})
+				Expect(bitCount.Err()).NotTo(HaveOccurred())
+				Expect(bitCount.Val()).To(Equal(int64(1)))
+
+				bitCount = adapter.Cache(time.Hour).BitCount(ctx, "key", &BitCount{
+					Start: 1,
+					End:   1,
+					Unit:  "BYTE",
+				})
+				Expect(bitCount.Err()).NotTo(HaveOccurred())
+				Expect(bitCount.Val()).To(Equal(int64(6)))
+			}
 		})
 
 		It("should BitPos", func() {

diff --git a/rueidiscompat/command.go b/rueidiscompat/command.go
@@ -1959,6 +1959,7 @@ type SetArgs struct {
 
 type BitCount struct {
 	Start, End int64
+	Unit       string // Stores BIT or BYTE
 }
 
 //type BitPos struct {