From 8d2a022fd5fbe71ee011004295d6991d894e811b Mon Sep 17 00:00:00 2001
From: Nanthen Hale <abusizhishen@gmail.com>
Date: Thu, 15 Feb 2024 04:15:30 +0800
Subject: [PATCH] Add support for specifying bitcount unit as byte or bit, byte
 default (#2887)

* Add support for specifying bitcount unit as byte or bit, byte default

* Add bitcount test

* Test bitcount without unit specified

---------

Co-authored-by: wanghongwei5 <wanghongwei5@360.cn>
Co-authored-by: ofekshenawa <104765379+ofekshenawa@users.noreply.github.com>
---
 bitmap_commands.go      | 14 ++++++
 bitmap_commands_test.go | 98 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 112 insertions(+)
 create mode 100644 bitmap_commands_test.go

diff --git a/bitmap_commands.go b/bitmap_commands.go
index 550d7f52..d9fc50dc 100644
--- a/bitmap_commands.go
+++ b/bitmap_commands.go
@@ -2,6 +2,7 @@ package redis
 
 import (
 	"context"
+	"errors"
 )
 
 type BitMapCmdable interface {
@@ -37,15 +38,28 @@ func (c cmdable) SetBit(ctx context.Context, key string, offset int64, value int
 
 type BitCount struct {
 	Start, End int64
+	Unit       string // BYTE(default) | BIT
 }
 
+const BitCountIndexByte string = "BYTE"
+const BitCountIndexBit string = "BIT"
+
 func (c cmdable) BitCount(ctx context.Context, key string, bitCount *BitCount) *IntCmd {
 	args := []interface{}{"bitcount", key}
 	if bitCount != nil {
+		if bitCount.Unit == "" {
+			bitCount.Unit = "BYTE"
+		}
+		if bitCount.Unit != BitCountIndexByte && bitCount.Unit != BitCountIndexBit {
+			cmd := NewIntCmd(ctx)
+			cmd.SetErr(errors.New("redis: invalid bitcount index"))
+			return cmd
+		}
 		args = append(
 			args,
 			bitCount.Start,
 			bitCount.End,
+			string(bitCount.Unit),
 		)
 	}
 	cmd := NewIntCmd(ctx, args...)
diff --git a/bitmap_commands_test.go b/bitmap_commands_test.go
new file mode 100644
index 00000000..f3cc3205
--- /dev/null
+++ b/bitmap_commands_test.go
@@ -0,0 +1,98 @@
+package redis_test
+
+import (
+	. "github.com/bsm/ginkgo/v2"
+	. "github.com/bsm/gomega"
+	"github.com/redis/go-redis/v9"
+)
+
+type bitCountExpected struct {
+	Start    int64
+	End      int64
+	Expected int64
+}
+
+var _ = Describe("BitCountBite", func() {
+	var client *redis.Client
+	key := "bit_count_test"
+
+	BeforeEach(func() {
+		client = redis.NewClient(redisOptions())
+		Expect(client.FlushDB(ctx).Err()).NotTo(HaveOccurred())
+		values := []int{0, 1, 0, 0, 1, 0, 1, 0, 1, 1}
+		for i, v := range values {
+			cmd := client.SetBit(ctx, key, int64(i), v)
+			Expect(cmd.Err()).NotTo(HaveOccurred())
+		}
+	})
+
+	AfterEach(func() {
+		Expect(client.Close()).NotTo(HaveOccurred())
+	})
+
+	It("bit count bite", func() {
+		var expected = []bitCountExpected{
+			{0, 0, 0},
+			{0, 1, 1},
+			{0, 2, 1},
+			{0, 3, 1},
+			{0, 4, 2},
+			{0, 5, 2},
+			{0, 6, 3},
+			{0, 7, 3},
+			{0, 8, 4},
+			{0, 9, 5},
+		}
+
+		for _, e := range expected {
+			cmd := client.BitCount(ctx, key, &redis.BitCount{Start: e.Start, End: e.End, Unit: redis.BitCountIndexBit})
+			Expect(cmd.Err()).NotTo(HaveOccurred())
+			Expect(cmd.Val()).To(Equal(e.Expected))
+		}
+	})
+})
+
+var _ = Describe("BitCountByte", func() {
+	var client *redis.Client
+	key := "bit_count_test"
+
+	BeforeEach(func() {
+		client = redis.NewClient(redisOptions())
+		Expect(client.FlushDB(ctx).Err()).NotTo(HaveOccurred())
+		values := []int{0, 0, 0, 0, 0, 0, 0, 1, 1, 1}
+		for i, v := range values {
+			cmd := client.SetBit(ctx, key, int64(i), v)
+			Expect(cmd.Err()).NotTo(HaveOccurred())
+		}
+	})
+
+	AfterEach(func() {
+		Expect(client.Close()).NotTo(HaveOccurred())
+	})
+
+	It("bit count byte", func() {
+		var expected = []bitCountExpected{
+			{0, 0, 1},
+			{0, 1, 3},
+		}
+
+		for _, e := range expected {
+			cmd := client.BitCount(ctx, key, &redis.BitCount{Start: e.Start, End: e.End, Unit: redis.BitCountIndexByte})
+			Expect(cmd.Err()).NotTo(HaveOccurred())
+			Expect(cmd.Val()).To(Equal(e.Expected))
+		}
+	})
+
+	It("bit count byte with no unit specified", func() {
+		var expected = []bitCountExpected{
+			{0, 0, 1},
+			{0, 1, 3},
+		}
+
+		for _, e := range expected {
+			cmd := client.BitCount(ctx, key, &redis.BitCount{Start: e.Start, End: e.End})
+			Expect(cmd.Err()).NotTo(HaveOccurred())
+			Expect(cmd.Val()).To(Equal(e.Expected))
+		}
+	})
+})