From 9212eafd5df14d59d7bfa1bd0f90a9f8f7c75689 Mon Sep 17 00:00:00 2001
From: Pieter Noordhuis <pcnoordhuis@gmail.com>
Date: Thu, 4 Mar 2010 12:01:45 +0100
Subject: [PATCH] implemented ZREMBYRANK

---
 redis.c        | 96 ++++++++++++++++++++++++++++++++++++++++++++++++++
 test-redis.tcl | 14 ++++++--
 2 files changed, 108 insertions(+), 2 deletions(-)

diff --git a/redis.c b/redis.c
index 0464554b..cfb0123d 100644
--- a/redis.c
+++ b/redis.c
@@ -661,6 +661,7 @@ static void brpopCommand(redisClient *c);
 static void appendCommand(redisClient *c);
 static void substrCommand(redisClient *c);
 static void zrankCommand(redisClient *c);
+static void zremrangebyrankCommand(redisClient *c);
 
 /*================================= Globals ================================= */
 
@@ -708,6 +709,7 @@ static struct redisCommand cmdTable[] = {
     {"zincrby",zincrbyCommand,4,REDIS_CMD_BULK|REDIS_CMD_DENYOOM,1,1,1},
     {"zrem",zremCommand,3,REDIS_CMD_BULK,1,1,1},
     {"zremrangebyscore",zremrangebyscoreCommand,4,REDIS_CMD_INLINE,1,1,1},
+    {"zremrangebyrank",zremrangebyrankCommand,4,REDIS_CMD_INLINE,1,1,1},
     {"zrange",zrangeCommand,-4,REDIS_CMD_INLINE,1,1,1},
     {"zrangebyscore",zrangebyscoreCommand,-4,REDIS_CMD_INLINE,1,1,1},
     {"zcount",zcountCommand,4,REDIS_CMD_INLINE,1,1,1},
@@ -5052,6 +5054,61 @@ static unsigned long zslDeleteRange(zskiplist *zsl, double min, double max, dict
     return removed; /* not found */
 }
 
+/* Delete all the elements with rank between start and end from the skiplist.
+ * Start and end are inclusive. */
+static unsigned long zslDeleteRangeByRank(zskiplist *zsl, unsigned int start, unsigned int end, dict *dict) {
+    zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
+    unsigned long traversed = 0, removed = 0;
+    int i;
+
+    /* start and end are given 0-based, but zsl uses 1-based
+     * ranks internally */
+    start++; end++;
+
+    x = zsl->header;
+    for (i = zsl->level-1; i >= 0; i--) {
+        while (x->forward[i] && (traversed + (i > 0 ? x->span[i-1] : 1)) < start) {
+            traversed += i > 0 ? x->span[i-1] : 1;
+            x = x->forward[i];
+        }
+        update[i] = x;
+    }
+
+    traversed++;
+    x = x->forward[0];
+    while (x && traversed <= end) {
+        zskiplistNode *next;
+
+        for (i = 0; i < zsl->level; i++) {
+            if (update[i]->forward[i] == x) {
+                if (i > 0) {
+                    update[i]->span[i-1] += x->span[i-1] - 1;
+                }
+                update[i]->forward[i] = x->forward[i];
+            } else {
+                /* invariant: i > 0, because update[0]->forward[0]
+                 * is always equal to x */
+                update[i]->span[i-1] -= 1;
+            }
+        }
+        if (x->forward[0]) {
+            x->forward[0]->backward = x->backward;
+        } else {
+            zsl->tail = x->backward;
+        }
+        next = x->forward[0];
+        dictDelete(dict,x->obj);
+        zslFreeNode(x);
+        while(zsl->level > 1 && zsl->header->forward[zsl->level-1] == NULL)
+            zsl->level--;
+        zsl->length--;
+        removed++;
+        traversed++;
+        x = next;
+    }
+    return removed;
+}
+
 /* Find the first node having a score equal or greater than the specified one.
  * Returns NULL if there is no match. */
 static zskiplistNode *zslFirstWithScore(zskiplist *zsl, double score) {
@@ -5271,6 +5328,45 @@ static void zremrangebyscoreCommand(redisClient *c) {
     }
 }
 
+static void zremrangebyrankCommand(redisClient *c) {
+    int start = atoi(c->argv[2]->ptr);
+    int end = atoi(c->argv[3]->ptr);
+    robj *zsetobj;
+    zset *zs;
+
+    zsetobj = lookupKeyWrite(c->db,c->argv[1]);
+    if (zsetobj == NULL) {
+        addReply(c,shared.czero);
+    } else {
+        if (zsetobj->type != REDIS_ZSET) {
+            addReply(c,shared.wrongtypeerr);
+            return;
+        }
+
+        zs = zsetobj->ptr;
+        int llen = zs->zsl->length;
+        long deleted;
+
+        /* convert negative indexes */
+        if (start < 0) start = llen+start;
+        if (end < 0) end = llen+end;
+        if (start < 0) start = 0;
+        if (end < 0) end = 0;
+
+        /* indexes sanity checks */
+        if (start > end || start >= llen) {
+            addReply(c,shared.czero);
+            return;
+        }
+        if (end >= llen) end = llen-1;
+
+        deleted = zslDeleteRangeByRank(zs->zsl,start,end,zs->dict);
+        if (htNeedsResize(zs->dict)) dictResize(zs->dict);
+        server.dirty += deleted;
+        addReplyLong(c, deleted);
+    }
+}
+
 static void zrangeGenericCommand(redisClient *c, int reverse) {
     robj *o;
     int start = atoi(c->argv[2]->ptr);
diff --git a/test-redis.tcl b/test-redis.tcl
index a6acea18..b1c4da01 100644
--- a/test-redis.tcl
+++ b/test-redis.tcl
@@ -1442,7 +1442,7 @@ proc main {server port} {
         $r zrangebyscore zset 20 50 LIMIT 2 3 withscores
     } {d 40 e 50}
 
-    test {ZREMRANGE basics} {
+    test {ZREMRANGEBYSCORE basics} {
         $r del zset
         $r zadd zset 1 a
         $r zadd zset 2 b
@@ -1452,7 +1452,7 @@ proc main {server port} {
         list [$r zremrangebyscore zset 2 4] [$r zrange zset 0 -1]
     } {3 {a e}}
 
-    test {ZREMRANGE from -inf to +inf} {
+    test {ZREMRANGEBYSCORE from -inf to +inf} {
         $r del zset
         $r zadd zset 1 a
         $r zadd zset 2 b
@@ -1462,6 +1462,16 @@ proc main {server port} {
         list [$r zremrangebyscore zset -inf +inf] [$r zrange zset 0 -1]
     } {5 {}}
 
+    test {ZREMRANGEBYRANK 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 zremrangebyrank zset 1 3] [$r zrange zset 0 -1]
+    } {3 {a e}}
+
     test {SORT against sorted sets} {
         $r del zset
         $r zadd zset 1 a
-- 
2.47.2