mirror of
https://github.com/HDT3213/godis.git
synced 2025-10-06 01:07:06 +08:00
support bitmap
This commit is contained in:
@@ -2,11 +2,13 @@ package database
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"github.com/hdt3213/godis/aof"
|
"github.com/hdt3213/godis/aof"
|
||||||
|
"github.com/hdt3213/godis/datastruct/bitmap"
|
||||||
"github.com/hdt3213/godis/interface/database"
|
"github.com/hdt3213/godis/interface/database"
|
||||||
"github.com/hdt3213/godis/interface/redis"
|
"github.com/hdt3213/godis/interface/redis"
|
||||||
"github.com/hdt3213/godis/lib/utils"
|
"github.com/hdt3213/godis/lib/utils"
|
||||||
"github.com/hdt3213/godis/redis/protocol"
|
"github.com/hdt3213/godis/redis/protocol"
|
||||||
"github.com/shopspring/decimal"
|
"github.com/shopspring/decimal"
|
||||||
|
"math/bits"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
@@ -536,46 +538,200 @@ func execSetRange(db *DB, args [][]byte) redis.Reply {
|
|||||||
|
|
||||||
func execGetRange(db *DB, args [][]byte) redis.Reply {
|
func execGetRange(db *DB, args [][]byte) redis.Reply {
|
||||||
key := string(args[0])
|
key := string(args[0])
|
||||||
startIdx, errNative := strconv.ParseInt(string(args[1]), 10, 64)
|
startIdx, err2 := strconv.ParseInt(string(args[1]), 10, 64)
|
||||||
if errNative != nil {
|
if err2 != nil {
|
||||||
return protocol.MakeErrReply(errNative.Error())
|
return protocol.MakeErrReply("ERR value is not an integer or out of range")
|
||||||
}
|
}
|
||||||
endIdx, errNative := strconv.ParseInt(string(args[2]), 10, 64)
|
endIdx, err2 := strconv.ParseInt(string(args[2]), 10, 64)
|
||||||
if errNative != nil {
|
if err2 != nil {
|
||||||
return protocol.MakeErrReply(errNative.Error())
|
return protocol.MakeErrReply("ERR value is not an integer or out of range")
|
||||||
}
|
}
|
||||||
|
|
||||||
bytes, err := db.getAsString(key)
|
bs, err := db.getAsString(key)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
if bs == nil {
|
||||||
if bytes == nil {
|
|
||||||
return protocol.MakeNullBulkReply()
|
return protocol.MakeNullBulkReply()
|
||||||
}
|
}
|
||||||
|
bytesLen := int64(len(bs))
|
||||||
bytesLen := int64(len(bytes))
|
beg, end := utils.ConvertRange(startIdx, endIdx, bytesLen)
|
||||||
if startIdx < -1*bytesLen {
|
if beg < 0 {
|
||||||
return &protocol.NullBulkReply{}
|
return protocol.MakeNullBulkReply()
|
||||||
} else if startIdx < 0 {
|
|
||||||
startIdx = bytesLen + startIdx
|
|
||||||
} else if startIdx >= bytesLen {
|
|
||||||
return &protocol.NullBulkReply{}
|
|
||||||
}
|
}
|
||||||
if endIdx < -1*bytesLen {
|
return protocol.MakeBulkReply(bs[beg:end])
|
||||||
return &protocol.NullBulkReply{}
|
}
|
||||||
} else if endIdx < 0 {
|
|
||||||
endIdx = bytesLen + endIdx + 1
|
func execSetBit(db *DB, args [][]byte) redis.Reply {
|
||||||
} else if endIdx < bytesLen {
|
key := string(args[0])
|
||||||
endIdx = endIdx + 1
|
offset, err := strconv.ParseInt(string(args[1]), 10, 64)
|
||||||
|
if err != nil {
|
||||||
|
return protocol.MakeErrReply("ERR bit offset is not an integer or out of range")
|
||||||
|
}
|
||||||
|
valStr := string(args[2])
|
||||||
|
var v byte
|
||||||
|
if valStr == "1" {
|
||||||
|
v = 1
|
||||||
|
} else if valStr == "0" {
|
||||||
|
v = 0
|
||||||
} else {
|
} else {
|
||||||
endIdx = bytesLen
|
return protocol.MakeErrReply("ERR bit is not an integer or out of range\n")
|
||||||
}
|
}
|
||||||
if startIdx > endIdx {
|
bs, errReply := db.getAsString(key)
|
||||||
return protocol.MakeNullBulkReply()
|
if errReply != nil {
|
||||||
|
return errReply
|
||||||
}
|
}
|
||||||
|
bm := bitmap.FromBytes(bs)
|
||||||
|
former := bm.GetBit(offset)
|
||||||
|
bm.SetBit(offset, v)
|
||||||
|
db.PutEntity(key, &database.DataEntity{Data: bm.ToBytes()})
|
||||||
|
return protocol.MakeIntReply(int64(former))
|
||||||
|
}
|
||||||
|
|
||||||
return protocol.MakeBulkReply(bytes[startIdx:endIdx])
|
func execGetBit(db *DB, args [][]byte) redis.Reply {
|
||||||
|
key := string(args[0])
|
||||||
|
offset, err := strconv.ParseInt(string(args[1]), 10, 64)
|
||||||
|
if err != nil {
|
||||||
|
return protocol.MakeErrReply("ERR bit offset is not an integer or out of range")
|
||||||
|
}
|
||||||
|
bs, errReply := db.getAsString(key)
|
||||||
|
if errReply != nil {
|
||||||
|
return errReply
|
||||||
|
}
|
||||||
|
if bs == nil {
|
||||||
|
return protocol.MakeIntReply(0)
|
||||||
|
}
|
||||||
|
bm := bitmap.FromBytes(bs)
|
||||||
|
return protocol.MakeIntReply(int64(bm.GetBit(offset)))
|
||||||
|
}
|
||||||
|
|
||||||
|
func execBitCount(db *DB, args [][]byte) redis.Reply {
|
||||||
|
key := string(args[0])
|
||||||
|
bs, err := db.getAsString(key)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if bs == nil {
|
||||||
|
return protocol.MakeIntReply(0)
|
||||||
|
}
|
||||||
|
byteMode := true
|
||||||
|
if len(args) > 3 {
|
||||||
|
mode := strings.ToLower(string(args[3]))
|
||||||
|
if mode == "bit" {
|
||||||
|
byteMode = false
|
||||||
|
} else if mode == "byte" {
|
||||||
|
byteMode = true
|
||||||
|
} else {
|
||||||
|
return protocol.MakeErrReply("ERR syntax error")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
var size int64
|
||||||
|
bm := bitmap.FromBytes(bs)
|
||||||
|
if byteMode {
|
||||||
|
size = int64(len(*bm))
|
||||||
|
} else {
|
||||||
|
size = int64(bm.BitSize())
|
||||||
|
}
|
||||||
|
var beg, end int
|
||||||
|
if len(args) > 1 {
|
||||||
|
var err2 error
|
||||||
|
var startIdx, endIdx int64
|
||||||
|
startIdx, err2 = strconv.ParseInt(string(args[1]), 10, 64)
|
||||||
|
if err2 != nil {
|
||||||
|
return protocol.MakeErrReply("ERR value is not an integer or out of range")
|
||||||
|
}
|
||||||
|
endIdx, err2 = strconv.ParseInt(string(args[2]), 10, 64)
|
||||||
|
if err2 != nil {
|
||||||
|
return protocol.MakeErrReply("ERR value is not an integer or out of range")
|
||||||
|
}
|
||||||
|
beg, end = utils.ConvertRange(startIdx, endIdx, size)
|
||||||
|
if beg < 0 {
|
||||||
|
return protocol.MakeIntReply(0)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
var count int64
|
||||||
|
if byteMode {
|
||||||
|
bm.ForEachByte(beg, end, func(offset int64, val byte) bool {
|
||||||
|
count += int64(bits.OnesCount8(val))
|
||||||
|
return true
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
bm.ForEachBit(int64(beg), int64(end), func(offset int64, val byte) bool {
|
||||||
|
if val > 0 {
|
||||||
|
count++
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
})
|
||||||
|
}
|
||||||
|
return protocol.MakeIntReply(count)
|
||||||
|
}
|
||||||
|
|
||||||
|
func execBitPos(db *DB, args [][]byte) redis.Reply {
|
||||||
|
key := string(args[0])
|
||||||
|
bs, err := db.getAsString(key)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if bs == nil {
|
||||||
|
return protocol.MakeIntReply(-1)
|
||||||
|
}
|
||||||
|
valStr := string(args[1])
|
||||||
|
var v byte
|
||||||
|
if valStr == "1" {
|
||||||
|
v = 1
|
||||||
|
} else if valStr == "0" {
|
||||||
|
v = 0
|
||||||
|
} else {
|
||||||
|
return protocol.MakeErrReply("ERR bit is not an integer or out of range")
|
||||||
|
}
|
||||||
|
byteMode := true
|
||||||
|
if len(args) > 4 {
|
||||||
|
mode := strings.ToLower(string(args[4]))
|
||||||
|
if mode == "bit" {
|
||||||
|
byteMode = false
|
||||||
|
} else if mode == "byte" {
|
||||||
|
byteMode = true
|
||||||
|
} else {
|
||||||
|
return protocol.MakeErrReply("ERR syntax error")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
var size int64
|
||||||
|
bm := bitmap.FromBytes(bs)
|
||||||
|
if byteMode {
|
||||||
|
size = int64(len(*bm))
|
||||||
|
} else {
|
||||||
|
size = int64(bm.BitSize())
|
||||||
|
}
|
||||||
|
var beg, end int
|
||||||
|
if len(args) > 2 {
|
||||||
|
var err2 error
|
||||||
|
var startIdx, endIdx int64
|
||||||
|
startIdx, err2 = strconv.ParseInt(string(args[2]), 10, 64)
|
||||||
|
if err2 != nil {
|
||||||
|
return protocol.MakeErrReply("ERR value is not an integer or out of range")
|
||||||
|
}
|
||||||
|
endIdx, err2 = strconv.ParseInt(string(args[3]), 10, 64)
|
||||||
|
if err2 != nil {
|
||||||
|
return protocol.MakeErrReply("ERR value is not an integer or out of range")
|
||||||
|
}
|
||||||
|
beg, end = utils.ConvertRange(startIdx, endIdx, size)
|
||||||
|
if beg < 0 {
|
||||||
|
return protocol.MakeIntReply(0)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if byteMode {
|
||||||
|
beg *= 8
|
||||||
|
end *= 8
|
||||||
|
}
|
||||||
|
var offset = int64(-1)
|
||||||
|
bm.ForEachBit(int64(beg), int64(end), func(o int64, val byte) bool {
|
||||||
|
if val == v {
|
||||||
|
offset = o
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
})
|
||||||
|
return protocol.MakeIntReply(offset)
|
||||||
}
|
}
|
||||||
|
|
||||||
func init() {
|
func init() {
|
||||||
@@ -597,4 +753,9 @@ func init() {
|
|||||||
RegisterCommand("Append", execAppend, writeFirstKey, rollbackFirstKey, 3)
|
RegisterCommand("Append", execAppend, writeFirstKey, rollbackFirstKey, 3)
|
||||||
RegisterCommand("SetRange", execSetRange, writeFirstKey, rollbackFirstKey, 4)
|
RegisterCommand("SetRange", execSetRange, writeFirstKey, rollbackFirstKey, 4)
|
||||||
RegisterCommand("GetRange", execGetRange, readFirstKey, nil, 4)
|
RegisterCommand("GetRange", execGetRange, readFirstKey, nil, 4)
|
||||||
|
RegisterCommand("SetBit", execSetBit, writeFirstKey, rollbackFirstKey, 4)
|
||||||
|
RegisterCommand("GetBit", execGetBit, readFirstKey, nil, 3)
|
||||||
|
RegisterCommand("BitCount", execBitCount, readFirstKey, nil, -2)
|
||||||
|
RegisterCommand("BitPos", execBitPos, readFirstKey, nil, -3)
|
||||||
|
|
||||||
}
|
}
|
||||||
|
@@ -323,12 +323,12 @@ func TestStrLen(t *testing.T) {
|
|||||||
testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
|
testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
|
||||||
|
|
||||||
actual := testDB.Exec(nil, utils.ToCmdLine("StrLen", key))
|
actual := testDB.Exec(nil, utils.ToCmdLine("StrLen", key))
|
||||||
len, ok := actual.(*protocol.IntReply)
|
size, ok := actual.(*protocol.IntReply)
|
||||||
if !ok {
|
if !ok {
|
||||||
t.Errorf("expect int bulk protocol, get: %s", string(actual.ToBytes()))
|
t.Errorf("expect int bulk protocol, get: %s", string(actual.ToBytes()))
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
asserts.AssertIntReply(t, len, 10)
|
asserts.AssertIntReply(t, size, 10)
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestStrLen_KeyNotExist(t *testing.T) {
|
func TestStrLen_KeyNotExist(t *testing.T) {
|
||||||
@@ -648,3 +648,90 @@ func TestGetRange_StringExist_EndIdxIncorrectFormat(t *testing.T) {
|
|||||||
errorMsg := fmt.Sprintf("strconv.ParseInt: parsing \"%s\": invalid syntax", incorrectValue)
|
errorMsg := fmt.Sprintf("strconv.ParseInt: parsing \"%s\": invalid syntax", incorrectValue)
|
||||||
asserts.AssertErrReply(t, val, errorMsg)
|
asserts.AssertErrReply(t, val, errorMsg)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestSetBit(t *testing.T) {
|
||||||
|
testDB.Flush()
|
||||||
|
key := utils.RandString(10)
|
||||||
|
actual := testDB.Exec(nil, utils.ToCmdLine("SetBit", key, "15", "1"))
|
||||||
|
asserts.AssertIntReply(t, actual, 0)
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("SetBit", key, "15", "0"))
|
||||||
|
asserts.AssertIntReply(t, actual, 1)
|
||||||
|
testDB.Exec(nil, utils.ToCmdLine("SetBit", key, "13", "1"))
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("GetBit", key, "13"))
|
||||||
|
asserts.AssertIntReply(t, actual, 1)
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("GetBit", key+"1", "13")) // test not exist key
|
||||||
|
asserts.AssertIntReply(t, actual, 0)
|
||||||
|
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("SetBit", key, "13", "a"))
|
||||||
|
asserts.AssertErrReply(t, actual, "ERR bit is not an integer or out of range")
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("SetBit", key, "a", "1"))
|
||||||
|
asserts.AssertErrReply(t, actual, "ERR bit offset is not an integer or out of range")
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("GetBit", key, "a"))
|
||||||
|
asserts.AssertErrReply(t, actual, "ERR bit offset is not an integer or out of range")
|
||||||
|
|
||||||
|
key2 := utils.RandString(11)
|
||||||
|
testDB.Exec(nil, utils.ToCmdLine("rpush", key2, "1"))
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("SetBit", key2, "15", "0"))
|
||||||
|
asserts.AssertErrReply(t, actual, "WRONGTYPE Operation against a key holding the wrong kind of value")
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("GetBit", key2, "15"))
|
||||||
|
asserts.AssertErrReply(t, actual, "WRONGTYPE Operation against a key holding the wrong kind of value")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBitCount(t *testing.T) {
|
||||||
|
testDB.Flush()
|
||||||
|
key := utils.RandString(10)
|
||||||
|
testDB.Exec(nil, utils.ToCmdLine("SetBit", key, "15", "1"))
|
||||||
|
testDB.Exec(nil, utils.ToCmdLine("SetBit", key, "13", "1"))
|
||||||
|
actual := testDB.Exec(nil, utils.ToCmdLine("BitCount", key))
|
||||||
|
asserts.AssertIntReply(t, actual, 2)
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitCount", key, "14", "15", "BIT"))
|
||||||
|
asserts.AssertIntReply(t, actual, 1)
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitCount", key, "16", "20", "BIT"))
|
||||||
|
asserts.AssertIntReply(t, actual, 0)
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitCount", key, "1", "1", "BYTE"))
|
||||||
|
asserts.AssertIntReply(t, actual, 2)
|
||||||
|
|
||||||
|
key2 := utils.RandString(11)
|
||||||
|
testDB.Exec(nil, utils.ToCmdLine("rpush", key2, "1"))
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitCount", key2))
|
||||||
|
asserts.AssertErrReply(t, actual, "WRONGTYPE Operation against a key holding the wrong kind of value")
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitCount", key+"a"))
|
||||||
|
asserts.AssertIntReply(t, actual, 0)
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitCount", key, "14", "15", "B"))
|
||||||
|
asserts.AssertErrReply(t, actual, "ERR syntax error")
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitCount", key, "14", "A"))
|
||||||
|
asserts.AssertErrReply(t, actual, "ERR value is not an integer or out of range")
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitCount", key, "A", "-1"))
|
||||||
|
asserts.AssertErrReply(t, actual, "ERR value is not an integer or out of range")
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBitPos(t *testing.T) {
|
||||||
|
testDB.Flush()
|
||||||
|
key := utils.RandString(10)
|
||||||
|
testDB.Exec(nil, utils.ToCmdLine("SetBit", key, "15", "1"))
|
||||||
|
actual := testDB.Exec(nil, utils.ToCmdLine("BitPos", key, "0"))
|
||||||
|
asserts.AssertIntReply(t, actual, 0)
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitPos", key, "1"))
|
||||||
|
asserts.AssertIntReply(t, actual, 15)
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitPos", key, "1", "0", "-1", "BIT"))
|
||||||
|
asserts.AssertIntReply(t, actual, 15)
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitPos", key, "1", "1", "1", "BYTE"))
|
||||||
|
asserts.AssertIntReply(t, actual, 15)
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitPos", key, "0", "1", "1", "BYTE"))
|
||||||
|
asserts.AssertIntReply(t, actual, 8)
|
||||||
|
|
||||||
|
key2 := utils.RandString(12)
|
||||||
|
testDB.Exec(nil, utils.ToCmdLine("rpush", key2, "1"))
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitPos", key2, "1"))
|
||||||
|
asserts.AssertErrReply(t, actual, "WRONGTYPE Operation against a key holding the wrong kind of value")
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitPos", key+"a", "1"))
|
||||||
|
asserts.AssertIntReply(t, actual, -1)
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitPos", key, "1", "1", "15", "B"))
|
||||||
|
asserts.AssertErrReply(t, actual, "ERR syntax error")
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitPos", key, "1", "14", "A"))
|
||||||
|
asserts.AssertErrReply(t, actual, "ERR value is not an integer or out of range")
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitPos", key, "1", "a", "14"))
|
||||||
|
asserts.AssertErrReply(t, actual, "ERR value is not an integer or out of range")
|
||||||
|
actual = testDB.Exec(nil, utils.ToCmdLine("BitPos", key, "-1"))
|
||||||
|
asserts.AssertErrReply(t, actual, "ERR bit is not an integer or out of range")
|
||||||
|
}
|
||||||
|
97
datastruct/bitmap/bitmap.go
Normal file
97
datastruct/bitmap/bitmap.go
Normal file
@@ -0,0 +1,97 @@
|
|||||||
|
package bitmap
|
||||||
|
|
||||||
|
type BitMap []byte
|
||||||
|
|
||||||
|
func New() *BitMap {
|
||||||
|
b := BitMap(make([]byte, 0))
|
||||||
|
return &b
|
||||||
|
}
|
||||||
|
|
||||||
|
func toByteSize(bitSize int64) int64 {
|
||||||
|
if bitSize%8 == 0 {
|
||||||
|
return bitSize / 8
|
||||||
|
}
|
||||||
|
return bitSize/8 + 1
|
||||||
|
}
|
||||||
|
|
||||||
|
func (b *BitMap) grow(bitSize int64) {
|
||||||
|
byteSize := toByteSize(bitSize)
|
||||||
|
gap := byteSize - int64(len(*b))
|
||||||
|
if gap <= 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
*b = append(*b, make([]byte, gap)...)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (b *BitMap) BitSize() int {
|
||||||
|
return len(*b) * 8
|
||||||
|
}
|
||||||
|
|
||||||
|
func FromBytes(bytes []byte) *BitMap {
|
||||||
|
bm := BitMap(bytes)
|
||||||
|
return &bm
|
||||||
|
}
|
||||||
|
|
||||||
|
func (b *BitMap) ToBytes() []byte {
|
||||||
|
return *b
|
||||||
|
}
|
||||||
|
|
||||||
|
func (b *BitMap) SetBit(offset int64, val byte) {
|
||||||
|
byteIndex := offset / 8
|
||||||
|
bitOffset := offset % 8
|
||||||
|
mask := byte(1 << bitOffset)
|
||||||
|
b.grow(offset + 1)
|
||||||
|
if val > 0 {
|
||||||
|
// set bit
|
||||||
|
(*b)[byteIndex] |= mask
|
||||||
|
} else {
|
||||||
|
// clear bit
|
||||||
|
(*b)[byteIndex] &^= mask
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (b *BitMap) GetBit(offset int64) byte {
|
||||||
|
byteIndex := offset / 8
|
||||||
|
bitOffset := offset % 8
|
||||||
|
if byteIndex >= int64(len(*b)) {
|
||||||
|
return 0
|
||||||
|
}
|
||||||
|
return ((*b)[byteIndex] >> bitOffset) & 0x01
|
||||||
|
}
|
||||||
|
|
||||||
|
type Callback func(offset int64, val byte) bool
|
||||||
|
|
||||||
|
func (b *BitMap) ForEachBit(begin int64, end int64, cb Callback) {
|
||||||
|
offset := begin
|
||||||
|
byteIndex := offset / 8
|
||||||
|
bitOffset := offset % 8
|
||||||
|
for byteIndex < int64(len(*b)) {
|
||||||
|
b := (*b)[byteIndex]
|
||||||
|
for bitOffset < 8 {
|
||||||
|
bit := byte(b >> bitOffset & 0x01)
|
||||||
|
if !cb(offset, bit) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
bitOffset++
|
||||||
|
offset++
|
||||||
|
}
|
||||||
|
byteIndex++
|
||||||
|
bitOffset = 0
|
||||||
|
if end > 0 && offset == end {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (b *BitMap) ForEachByte(begin int, end int, cb Callback) {
|
||||||
|
if end == 0 {
|
||||||
|
end = len(*b)
|
||||||
|
} else if end > len(*b) {
|
||||||
|
end = len(*b)
|
||||||
|
}
|
||||||
|
for i := begin; i < end; i++ {
|
||||||
|
if !cb(int64(i), (*b)[i]) {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
168
datastruct/bitmap/bitmap_test.go
Normal file
168
datastruct/bitmap/bitmap_test.go
Normal file
@@ -0,0 +1,168 @@
|
|||||||
|
package bitmap
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"math/rand"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestSetBit(t *testing.T) {
|
||||||
|
// generate distinct rand offsets
|
||||||
|
size := 1000
|
||||||
|
offsets := make([]int64, size)
|
||||||
|
for i := 0; i < size; i++ {
|
||||||
|
offsets[i] = int64(i)
|
||||||
|
}
|
||||||
|
rand.Shuffle(size, func(i, j int) {
|
||||||
|
offsets[i], offsets[j] = offsets[j], offsets[i]
|
||||||
|
})
|
||||||
|
offsets = offsets[0 : size/5]
|
||||||
|
// set bit
|
||||||
|
offsetMap := make(map[int64]struct{})
|
||||||
|
bm := New()
|
||||||
|
for _, offset := range offsets {
|
||||||
|
offsetMap[offset] = struct{}{}
|
||||||
|
bm.SetBit(offset, 1)
|
||||||
|
}
|
||||||
|
// get bit
|
||||||
|
for i := 0; i < bm.BitSize(); i++ {
|
||||||
|
offset := int64(i)
|
||||||
|
_, expect := offsetMap[offset]
|
||||||
|
actual := bm.GetBit(offset) > 0
|
||||||
|
if expect != actual {
|
||||||
|
t.Errorf("wrong value at %d", offset)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
bm2 := New()
|
||||||
|
bm2.SetBit(15, 1)
|
||||||
|
if bm2.GetBit(15) != 1 {
|
||||||
|
t.Error("wrong value")
|
||||||
|
}
|
||||||
|
if bm2.GetBit(16) != 0 {
|
||||||
|
t.Error("wrong value")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFromBytes(t *testing.T) {
|
||||||
|
bs := []byte{0xff, 0xff}
|
||||||
|
bm := FromBytes(bs)
|
||||||
|
bm.SetBit(8, 0)
|
||||||
|
expect := []byte{0xff, 0xfe}
|
||||||
|
if !bytes.Equal(bs, expect) {
|
||||||
|
t.Error("wrong value")
|
||||||
|
}
|
||||||
|
ret := bm.ToBytes()
|
||||||
|
if !bytes.Equal(ret, expect) {
|
||||||
|
t.Error("wrong value")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestForEachBit(t *testing.T) {
|
||||||
|
bm := New()
|
||||||
|
for i := 0; i < 1000; i++ {
|
||||||
|
if i%2 == 0 {
|
||||||
|
bm.SetBit(int64(i), 1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
expectOffset := int64(100)
|
||||||
|
count := 0
|
||||||
|
bm.ForEachBit(100, 200, func(offset int64, val byte) bool {
|
||||||
|
if offset != expectOffset {
|
||||||
|
t.Error("wrong offset")
|
||||||
|
}
|
||||||
|
expectOffset++
|
||||||
|
if offset%2 == 0 && val == 0 {
|
||||||
|
t.Error("expect 1")
|
||||||
|
}
|
||||||
|
count++
|
||||||
|
return true
|
||||||
|
})
|
||||||
|
if count != 100 {
|
||||||
|
t.Error("wrong count")
|
||||||
|
}
|
||||||
|
bm = New()
|
||||||
|
size := 1000
|
||||||
|
offsets := make([]int64, size)
|
||||||
|
for i := 0; i < size; i++ {
|
||||||
|
offsets[i] = int64(i)
|
||||||
|
}
|
||||||
|
rand.Shuffle(size, func(i, j int) {
|
||||||
|
offsets[i], offsets[j] = offsets[j], offsets[i]
|
||||||
|
})
|
||||||
|
offsets = offsets[0 : size/5]
|
||||||
|
offsetMap := make(map[int64]struct{})
|
||||||
|
for _, offset := range offsets {
|
||||||
|
bm.SetBit(offset, 1)
|
||||||
|
offsetMap[offset] = struct{}{}
|
||||||
|
}
|
||||||
|
bm.ForEachBit(int64(size/20+1), 0, func(offset int64, val byte) bool {
|
||||||
|
_, expect := offsetMap[offset]
|
||||||
|
actual := bm.GetBit(offset) > 0
|
||||||
|
if expect != actual {
|
||||||
|
t.Errorf("wrong value at %d", offset)
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
})
|
||||||
|
count = 0
|
||||||
|
bm.ForEachBit(0, 0, func(offset int64, val byte) bool {
|
||||||
|
count++
|
||||||
|
return false
|
||||||
|
})
|
||||||
|
if count != 1 {
|
||||||
|
t.Error("break failed")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestBitMap_ForEachByte(t *testing.T) {
|
||||||
|
bm := New()
|
||||||
|
for i := 0; i < 1000; i++ {
|
||||||
|
if i%16 == 0 {
|
||||||
|
bm.SetBit(int64(i), 1)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
bm.ForEachByte(0, 0, func(offset int64, val byte) bool {
|
||||||
|
if offset%2 == 0 {
|
||||||
|
if val != 1 {
|
||||||
|
t.Error("wrong value")
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
if val != 0 {
|
||||||
|
t.Error("wrong value")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
})
|
||||||
|
bm.ForEachByte(0, 2000, func(offset int64, val byte) bool {
|
||||||
|
if offset%2 == 0 {
|
||||||
|
if val != 1 {
|
||||||
|
t.Error("wrong value")
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
if val != 0 {
|
||||||
|
t.Error("wrong value")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
})
|
||||||
|
bm.ForEachByte(0, 500, func(offset int64, val byte) bool {
|
||||||
|
if offset%2 == 0 {
|
||||||
|
if val != 1 {
|
||||||
|
t.Error("wrong value")
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
if val != 0 {
|
||||||
|
t.Error("wrong value")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
})
|
||||||
|
count := 0
|
||||||
|
bm.ForEachByte(0, 0, func(offset int64, val byte) bool {
|
||||||
|
count++
|
||||||
|
return false
|
||||||
|
})
|
||||||
|
if count != 1 {
|
||||||
|
t.Error("break failed")
|
||||||
|
}
|
||||||
|
}
|
@@ -57,3 +57,30 @@ func BytesEquals(a []byte, b []byte) bool {
|
|||||||
}
|
}
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ConvertRange converts redis index to go slice index
|
||||||
|
// -1 => size-1
|
||||||
|
// both inclusive [0, 10] => left inclusive right exclusive [0, 9)
|
||||||
|
// out of bound to max inbound [size, size+1] => [-1, -1]
|
||||||
|
func ConvertRange(start int64, end int64, size int64) (int, int) {
|
||||||
|
if start < -size {
|
||||||
|
return -1, -1
|
||||||
|
} else if start < 0 {
|
||||||
|
start = size + start
|
||||||
|
} else if start >= size {
|
||||||
|
return -1, -1
|
||||||
|
}
|
||||||
|
if end < -size {
|
||||||
|
return -1, -1
|
||||||
|
} else if end < 0 {
|
||||||
|
end = size + end + 1
|
||||||
|
} else if end < size {
|
||||||
|
end = end + 1
|
||||||
|
} else {
|
||||||
|
end = size
|
||||||
|
}
|
||||||
|
if start > end {
|
||||||
|
return -1, -1
|
||||||
|
}
|
||||||
|
return int(start), int(end)
|
||||||
|
}
|
||||||
|
Reference in New Issue
Block a user