Make ZREMRANGEBYSCORE accept the same range spec as ZRANGEBYSCORE

This allows to use inclusive/exclusive bounds for min and max when
deleting a range of scores from a sorted set.

diff --git a/src/t_zset.c b/src/t_zset.c
index 1a67febcc4846bcb84f164483dc1c69d523fab8b..f1a10378cbda8e596468404fae0eb9da280d6542 100644 (file)
--- a/src/t_zset.c
+++ b/src/t_zset.c
@@ -174,25 +174,35 @@ int zslDelete(zskiplist *zsl, double score, robj *obj) {
     return 0; /* not found */
 }
 
+/* Struct to hold a inclusive/exclusive range spec. */
+typedef struct {
+    double min, max;
+    int minex, maxex; /* are min or max exclusive? */
+} zrangespec;
+
 /* Delete all the elements with score between min and max from the skiplist.
  * Min and mx are inclusive, so a score >= min || score <= max is deleted.
  * Note that this function takes the reference to the hash table view of the
  * sorted set, in order to remove the elements from the hash table too. */
-unsigned long zslDeleteRangeByScore(zskiplist *zsl, double min, double max, dict *dict) {
+unsigned long zslDeleteRangeByScore(zskiplist *zsl, zrangespec range, dict *dict) {
     zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
     unsigned long removed = 0;
     int i;
 
     x = zsl->header;
     for (i = zsl->level-1; i >= 0; i--) {
-        while (x->level[i].forward && x->level[i].forward->score < min)
-            x = x->level[i].forward;
+        while (x->level[i].forward && (range.minex ?
+            x->level[i].forward->score <= range.min :
+            x->level[i].forward->score < range.min))
+                x = x->level[i].forward;
         update[i] = x;
     }
-    /* We may have multiple elements with the same score, what we need
-     * is to find the element with both the right score and object. */
+
+    /* Current node is the last with score < or <= min. */
     x = x->level[0].forward;
-    while (x && x->score <= max) {
+
+    /* Delete nodes while in range. */
+    while (x && (range.maxex ? x->score < range.max : x->score <= range.max)) {
         zskiplistNode *next = x->level[0].forward;
         zslDeleteNode(zsl,x,update);
         dictDelete(dict,x->obj);
@@ -200,7 +210,7 @@ unsigned long zslDeleteRangeByScore(zskiplist *zsl, double min, double max, dict
         removed++;
         x = next;
     }
-    return removed; /* not found */
+    return removed;
 }
 
 /* Delete all the elements with rank between start and end from the skiplist.
@@ -296,11 +306,6 @@ zskiplistNode* zslistTypeGetElementByRank(zskiplist *zsl, unsigned long rank) {
     return NULL;
 }
 
-typedef struct {
-    double min, max;
-    int minex, maxex; /* are min or max exclusive? */
-} zrangespec;
-
 /* Populate the rangespec according to the objects min and max. */
 int zslParseRange(robj *min, robj *max, zrangespec *spec) {
     spec->minex = spec->maxex = 0;
@@ -470,20 +475,19 @@ void zremCommand(redisClient *c) {
 }
 
 void zremrangebyscoreCommand(redisClient *c) {
-    double min;
-    double max;
+    zrangespec range;
     long deleted;
-    robj *zsetobj;
+    robj *o;
     zset *zs;
 
-    if ((getDoubleFromObjectOrReply(c, c->argv[2], &min, NULL) != REDIS_OK) ||
-        (getDoubleFromObjectOrReply(c, c->argv[3], &max, NULL) != REDIS_OK)) return;
+    /* Parse the range arguments. */
+    zslParseRange(c->argv[2],c->argv[3],&range);
 
-    if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
-        checkType(c,zsetobj,REDIS_ZSET)) return;
+    if ((o = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
+        checkType(c,o,REDIS_ZSET)) return;
 
-    zs = zsetobj->ptr;
-    deleted = zslDeleteRangeByScore(zs->zsl,min,max,zs->dict);
+    zs = o->ptr;
+    deleted = zslDeleteRangeByScore(zs->zsl,range,zs->dict);
     if (htNeedsResize(zs->dict)) dictResize(zs->dict);
     if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
     if (deleted) touchWatchedKey(c->db,c->argv[1]);

diff --git a/tests/unit/type/zset.tcl b/tests/unit/type/zset.tcl
index 949681eb47b1884ac9722436df768d724e6cf7db..47c056073bcce9e9d2171b089f091619c6c9f1fe 100644 (file)
--- a/tests/unit/type/zset.tcl
+++ b/tests/unit/type/zset.tcl
@@ -335,25 +335,56 @@ start_server {tags {"zset"}} {
         } {}
     }
 
-    test {ZREMRANGEBYSCORE basics} {
-        r del zset
-        r zadd zset 1 a
-        r zadd zset 2 b
-        r zadd zset 3 c
-        r zadd zset 4 d
-        r zadd zset 5 e
-        list [r zremrangebyscore zset 2 4] [r zrange zset 0 -1]
-    } {3 {a e}}
-
-    test {ZREMRANGEBYSCORE from -inf to +inf} {
-        r del zset
-        r zadd zset 1 a
-        r zadd zset 2 b
-        r zadd zset 3 c
-        r zadd zset 4 d
-        r zadd zset 5 e
-        list [r zremrangebyscore zset -inf +inf] [r zrange zset 0 -1]
-    } {5 {}}
+    test "ZREMRANGEBYSCORE basics" {
+        proc remrangebyscore {min max} {
+            create_zset zset {1 a 2 b 3 c 4 d 5 e}
+            r zremrangebyscore zset $min $max
+        }
+
+        # inner range
+        assert_equal 3 [remrangebyscore 2 4]
+        assert_equal {a e} [r zrange zset 0 -1]
+
+        # start underflow
+        assert_equal 1 [remrangebyscore -10 1]
+        assert_equal {b c d e} [r zrange zset 0 -1]
+
+        # end overflow
+        assert_equal 1 [remrangebyscore 5 10]
+        assert_equal {a b c d} [r zrange zset 0 -1]
+
+        # switch min and max
+        assert_equal 0 [remrangebyscore 4 2]
+        assert_equal {a b c d e} [r zrange zset 0 -1]
+
+        # -inf to mid
+        assert_equal 3 [remrangebyscore -inf 3]
+        assert_equal {d e} [r zrange zset 0 -1]
+
+        # mid to +inf
+        assert_equal 3 [remrangebyscore 3 +inf]
+        assert_equal {a b} [r zrange zset 0 -1]
+
+        # -inf to +inf
+        assert_equal 5 [remrangebyscore -inf +inf]
+        assert_equal {} [r zrange zset 0 -1]
+
+        # exclusive min
+        assert_equal 4 [remrangebyscore (1 5]
+        assert_equal {a} [r zrange zset 0 -1]
+        assert_equal 3 [remrangebyscore (2 5]
+        assert_equal {a b} [r zrange zset 0 -1]
+
+        # exclusive max
+        assert_equal 4 [remrangebyscore 1 (5]
+        assert_equal {e} [r zrange zset 0 -1]
+        assert_equal 3 [remrangebyscore 1 (4]
+        assert_equal {d e} [r zrange zset 0 -1]
+
+        # exclusive min and max
+        assert_equal 3 [remrangebyscore (1 (5]
+        assert_equal {a e} [r zrange zset 0 -1]
+    }
 
     test "ZREMRANGEBYRANK basics" {
         proc remrangebyrank {min max} {