diff --git a/src/t_zset.c b/src/t_zset.c index 1a67febc..f1a10378 100644 --- 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 949681eb..47c05607 100644 --- 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 basics" { + proc remrangebyscore {min max} { + create_zset zset {1 a 2 b 3 c 4 d 5 e} + r zremrangebyscore zset $min $max + } - 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 {}} + # 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} {