Add new operations to set.go and zset.go, and fix bugs in key.go and …

…ds.go
diiyw · May 1, 2024 · 55bc0c7 · 55bc0c7
1 parent aec0e01
commit 55bc0c7
Show file tree

Hide file tree

Showing 14 changed files with 330 additions and 90 deletions.
diff --git a/README.md b/README.md
@@ -45,7 +45,7 @@ Simple way to embed in your application.
 |                     |                 | SCAN            | SETEX               | SREM            | HICRBYFLOAT    	| RPUSHX           | ZRANGEBYSCORE         |
 |                     |                 | RANDOMKEY       | INCRBY              | SMOVE           | HSETNX          | LREM             | ZREVRANGEBYSCORE      |
 |                     |                 | RENAMEEX        | DECRBY              | SRANDMEMBER     | HMGET           | LSET             | ZREM                  |
-|                     |                 |                 | SETNX               | SINTERSTORE     | HMSET           | LRANGE           | ZREMRANGEBYRANK       |
+|                     |                 |  PERSIST        | SETNX               | SINTERSTORE     | HMSET           | LRANGE           | ZREMRANGEBYRANK       |
 |                     |                 |                 | INCRBYFLOAT         | SUNIONSTORE     | HCLEAR          | LPOPRPUSH        | ZREMRANGEBYSCORE      |
 |                     |                 |                 | APPEND              |                 | HSCAN           | RPOPLPUSH        | ZCLEAR                |
 |                     |                 |                 | GETRANGE            |                 | HVALS           | BLPOP            | ZEXISTS               |

diff --git a/ds/list/list.go b/ds/list/list.go
@@ -117,6 +117,12 @@ func (l *DoublyLinkedList) LRange(start, end int64) [][]byte {
 	var result [][]byte
 	currentNode := l.head
 	var index int64 = 0
+	if start != 0 && start >= end {
+		return result
+	}
+	if start < 0 {
+		start = l.size() + start
+	}
 	if end < 0 {
 		end = l.size() + end
 	}

diff --git a/ds/set/set.go b/ds/set/set.go
@@ -220,6 +220,12 @@ func (s *Set) SClear() {
 	s.data.Clear()
 }
 
+// Iter returns an iterator for the set.
+func (s *Set) Iter(fn func(member string) bool) {
+	s.data.Scan(func(member string, _ struct{}) bool {
+		return fn(member)
+	})
+}
 
 // Type returns the type of the data structure
 func (s *Set) Type() ds.DataType {

diff --git a/ds/zset/skiplist.go b/ds/zset/skiplist.go
@@ -212,7 +212,7 @@ func (skiplist *skiplist) getByRank(rank int64) *node {
 }
 
 func (skiplist *skiplist) hasInRange(min float64, max float64) bool {
-	if min > max || min == max {
+	if min > max {
 		// empty range
 		return false
 	}

diff --git a/ds/zset/sorted_set.go b/ds/zset/sorted_set.go
@@ -93,14 +93,17 @@ func (sortedSet *SortedSet) ZCard() int64 {
 }
 
 // ZRem removes the given member from set
-func (sortedSet *SortedSet) ZRem(member string) bool {
-	v, ok := sortedSet.dict.Get(member)
-	if ok {
-		sortedSet.skiplist.remove(member, v.Score)
-		sortedSet.dict.Delete(member)
-		return true
+func (sortedSet *SortedSet) ZRem(members ...string) int64 {
+	var count int64
+	for _, member := range members {
+		v, ok := sortedSet.dict.Get(member)
+		if ok {
+			sortedSet.skiplist.remove(member, v.Score)
+			sortedSet.dict.Delete(member)
+			count++
+		}
 	}
-	return false
+	return count
 }
 
 // getRank returns the rank of the given member, sort by ascending order, rank starts from 0
@@ -115,7 +118,7 @@ func (sortedSet *SortedSet) getRank(member string, desc bool) (rank int64) {
 	} else {
 		r--
 	}
-	return r + 1
+	return r
 }
 
 // ZRank returns the rank of the given member, sort by ascending order, rank starts from 0
@@ -158,20 +161,25 @@ func (sortedSet *SortedSet) ZScore(member string) float64 {
 // forEachByRank visits each member which rank within [start, stop), sort by ascending order, rank starts from 1
 func (sortedSet *SortedSet) forEachByRank(start int64, stop int64, desc bool, consumer func(item *Item) bool) {
 	size := sortedSet.ZCard()
-	if start < 0 || start > size {
+	if start > size {
 		return
 	}
+	if start == 0 {
+		start = 1
+	}
+	if stop < 0 {
+		stop = size + stop + 1
+	}
 	if stop < start {
 		return
 	}
+	if start < 0 {
+		start = size + start
+	}
 	// stop max is size
 	if stop > size {
 		stop = size
 	}
-	if start == 0 {
-		start = 1
-	}
-
 	// find start node
 	var node *node
 	if desc {
@@ -201,8 +209,7 @@ func (sortedSet *SortedSet) forEachByRank(start int64, stop int64, desc bool, co
 
 // rangeByRank returns members which rank within [start, stop), sort by ascending order, rank starts from 0
 func (sortedSet *SortedSet) rangeByRank(start int64, stop int64, desc bool) []*Item {
-	sliceSize := int(stop - start + 1)
-	slice := make([]*Item, 0, sliceSize)
+	slice := make([]*Item, 0)
 	sortedSet.forEachByRank(start, stop, desc, func(item *Item) bool {
 		slice = append(slice, item)
 		return true
@@ -225,18 +232,9 @@ func (sortedSet *SortedSet) rangeCount(min float64, max float64) int64 {
 	var i int64 = 0
 	// ascending order
 	sortedSet.forEachByRank(0, sortedSet.ZCard(), false, func(element *Item) bool {
-		gtMin := min < element.Score // greater than min
-		if !gtMin {
-			// has not into range, continue foreach
-			return true
-		}
-		ltMax := max > element.Score // less than max
-		if !ltMax {
-			// break through score border, break foreach
-			return false
+		if element.Score >= min && element.Score <= max {
+			i++
 		}
-		// gtMin && ltMax
-		i++
 		return true
 	})
 	return i
@@ -344,11 +342,16 @@ func (sortedSet *SortedSet) ZRevRangeByScore(min float64, max float64, offset, c
 // ZIncrBy increases the score of the given member
 func (sortedSet *SortedSet) ZIncrBy(member string, score float64) float64 {
 	element, ok := sortedSet.dict.Get(member)
-	if !ok {
-		return 0
+	if ok {
+		score += element.Score
 	}
-	sortedSet.zAdd(member, element.Score+score)
-	return element.Score + score
+	sortedSet.zAdd(member, score)
+	return score
+}
+
+// GetMax returns the member with the highest score
+func (sortedSet *SortedSet) ZMax() *Item {
+	return &sortedSet.skiplist.tail.Item
 }
 
 func (sortedSet *SortedSet) GetValue() []*pb.KeyScore {