feat: add zipzset

internal/zset/bench_test.go

@@ -39,12 +39,11 @@ func benchZSetI(name string, newf func() iface.ZSetI, b *testing.B) {
 			m.Set(genKey(i%N), float64(i%N))
 		}
 	})
-	b.Run(name+"/remove", func(b *testing.B) {
-		m := genZSet(newf(), N)
+	b.Run(name+"/popmin", func(b *testing.B) {
+		m := genZSet(newf(), b.N)
 		b.ResetTimer()
-		for i := 0; i < N; i++ {
-			k := genKey(i)
-			m.Remove(k)
+		for i := 0; i < b.N; i++ {
+			m.PopMin()
 		}
 	})
 	b.Run(name+"/scan", func(b *testing.B) {

internal/zset/fuzz_test.go

@@ -3,54 +3,64 @@ package zset
 import (
 	"fmt"
 	"github.com/stretchr/testify/assert"
+	"github.com/xgzlucario/rotom/internal/resp"
 	"testing"
+	"time"
 )
 
-const MAX = 10000
-
 func FuzzTestZSet(f *testing.F) {
 	zs := New()
 	zzs := NewZipZSet()
 
-	f.Fuzz(func(t *testing.T, op int, key string, score float64) {
+	f.Fuzz(func(t *testing.T, key string, score float64) {
 		ast := assert.New(t)
-		switch op % 10 {
-		case 0, 1, 2: // Set
-			if zs.Len() > MAX {
-				return
-			}
+		ts := time.Now().Second()
+		switch ts % 10 {
+		case 0, 1: // Set
 			ast.Equal(zs.Set(key, score), zzs.Set(key, score))
+			ast.Equal(zs.Len(), zzs.Len())
 
-		case 3, 4: // Get
+		case 2, 3: // Get
 			sc1, ok1 := zs.Get(key)
 			sc2, ok2 := zzs.Get(key)
 			ast.Equal(sc1, sc2)
 			ast.Equal(ok1, ok2)
-
-		case 5: // Rank
 			ast.Equal(zs.Rank(key), zzs.Rank(key))
 
-		case 6: // PopMin
+		case 4, 5: // PopMin
 			k1, s1 := zs.PopMin()
 			k2, s2 := zzs.PopMin()
 			ast.Equal(k1, k2)
 			ast.Equal(s1, s2)
 
-		case 7: // Delete
+		case 6, 7: // Remove
 			ast.Equal(zs.Remove(key), zzs.Remove(key))
+			ast.Equal(zs.Len(), zzs.Len())
 
 		case 8: // Scan
-			kv1 := make([]string, 0)
-			kv2 := make([]string, 0)
+			kv1 := make([]string, 0, zs.Len())
+			kv2 := make([]string, 0, zs.Len())
 			zs.Scan(func(k string, v float64) {
 				kv1 = append(kv1, fmt.Sprintf("%s->%v", k, v))
 			})
 			zzs.Scan(func(k string, v float64) {
 				kv2 = append(kv2, fmt.Sprintf("%s->%v", k, v))
 			})
-			ast.ElementsMatch(kv1, kv2)
+			ast.Equal(kv1, kv2)
 
 		case 9: // Encode
+			writer := resp.NewWriter(1024)
+
+			// zset
+			ast.Nil(zs.Encode(writer))
+			zs = New()
+			ast.Nil(zs.Decode(resp.NewReader(writer.Bytes())))
+			writer.Reset()
+
+			// zipzset
+			ast.Nil(zzs.Encode(writer))
+			zzs = NewZipZSet()
+			ast.Nil(zzs.Decode(resp.NewReader(writer.Bytes())))
 		}
 	})
 }

internal/zset/zipset.go

@@ -12,7 +12,7 @@ import (
 var (
 	_ iface.ZSetI = (*ZipZSet)(nil)
 
-	order = binary.BigEndian
+	order = binary.LittleEndian
 )
 
 // ZipZSet store data as [score1, key1, score2, key2...] in listpack.
@@ -37,9 +37,9 @@ func (*ZipZSet) decode(entry []byte) (string, float64) {
 }
 
 func (zs *ZipZSet) Get(key string) (float64, bool) {
-	it, _, val := zs.seek(key)
+	it, _, score := zs.rank(key)
 	if it != nil {
-		return b2f(val), true
+		return score, true
 	}
 	return 0, false
 }
@@ -51,36 +51,40 @@ func (zs *ZipZSet) Set(key string, score float64) bool {
 }
 
 func (zs *ZipZSet) Remove(key string) bool {
-	it, _, _ := zs.seek(key)
+	it, _, _ := zs.rank(key)
 	if it != nil {
 		it.RemoveNext()
 	}
 	return it != nil
 }
 
-func (zs *ZipZSet) seek(key string) (it *list.LpIterator, index int, entry []byte) {
-	it = zs.data.Iterator().SeekLast()
-	index = -1
+func (zs *ZipZSet) rank(key string) (*list.LpIterator, int, float64) {
+	it := zs.data.Iterator().SeekLast()
+	index := -1
 	for !it.IsFirst() {
-		entry = it.Prev()
+		prevKey, prevScore := zs.decode(it.Prev())
 		index++
-		if key == b2s(entry[8:]) {
-			return it, index, entry
+		if key == prevKey {
+			return it, index, prevScore
 		}
 	}
-	return nil, -1, nil
+	return nil, -1, 0
 }
 
 func (zs *ZipZSet) insert(key string, score float64) {
 	it := zs.data.Iterator().SeekLast()
 	for !it.IsFirst() {
-		ek, es := zs.decode(it.Prev())
-		if es < score || (es == score && ek < key) {
-			continue
-		} else {
+		prevKey, prevScore := zs.decode(it.Prev())
+		if score < prevScore {
 			it.Next()
 			goto DO
 		}
+		if score == prevScore {
+			if key < prevKey {
+				it.Next()
+				goto DO
+			}
+		}
 	}
 DO:
 	it.Insert(b2s(zs.encode(key, score)))
@@ -89,13 +93,13 @@ DO:
 func (zs *ZipZSet) PopMin() (string, float64) {
 	entry, ok := zs.data.RPop()
 	if ok {
-		return zs.decode([]byte(entry))
+		return zs.decode(s2b(entry))
 	}
 	return "", 0
 }
 
 func (zs *ZipZSet) Rank(key string) int {
-	_, index, _ := zs.seek(key)
+	_, index, _ := zs.rank(key)
 	return index
 }
 
@@ -132,6 +136,10 @@ func b2s(b []byte) string {
 	return *(*string)(unsafe.Pointer(&b))
 }
 
+func s2b(s string) []byte {
+	return *(*[]byte)(unsafe.Pointer(&s))
+}
+
 func b2f(b []byte) float64 {
 	return math.Float64frombits(order.Uint64(b))
 }