diff --git a/rueidiscompat/adapter.go b/rueidiscompat/adapter.go index 26b4a143..5a8df142 100644 --- 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 index 65700f07..5c22c5f9 100644 --- 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 index f614597a..460f0526 100644 --- 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 {