fix for ZADD in score update mode

[redis.git] / redis.c

diff --git a/redis.c b/redis.c
index 7158de60f6a2625e7281cc54a3fc489fc50a0e16..ee99792bbd084d91e6be38e9083733e1c9cb6535 100644 (file)
--- a/redis.c
+++ b/redis.c
@@ -434,6 +434,8 @@ static void debugCommand(redisClient *c);
 static void msetCommand(redisClient *c);
 static void msetnxCommand(redisClient *c);
 static void zaddCommand(redisClient *c);
+static void zrangeCommand(redisClient *c);
+static void zlenCommand(redisClient *c);
 
 /*================================= Globals ================================= */
 
@@ -473,6 +475,8 @@ static struct redisCommand cmdTable[] = {
     {"sdiffstore",sdiffstoreCommand,-3,REDIS_CMD_INLINE|REDIS_CMD_DENYOOM},
     {"smembers",sinterCommand,2,REDIS_CMD_INLINE},
     {"zadd",zaddCommand,4,REDIS_CMD_BULK|REDIS_CMD_DENYOOM},
+    {"zrange",zrangeCommand,4,REDIS_CMD_INLINE},
+    {"zlen",zlenCommand,2,REDIS_CMD_INLINE},
     {"incrby",incrbyCommand,3,REDIS_CMD_INLINE|REDIS_CMD_DENYOOM},
     {"decrby",decrbyCommand,3,REDIS_CMD_INLINE|REDIS_CMD_DENYOOM},
     {"getset",getsetCommand,3,REDIS_CMD_BULK|REDIS_CMD_DENYOOM},
@@ -2050,6 +2054,7 @@ static robj *getDecodedObject(const robj *o) {
 static int compareStringObjects(robj *a, robj *b) {
     assert(a->type == REDIS_STRING && b->type == REDIS_STRING);
 
+    if (a == b) return 0;
     if (a->encoding == REDIS_ENCODING_INT && b->encoding == REDIS_ENCODING_INT){
         return (long)a->ptr - (long)b->ptr;
     } else {
@@ -3710,6 +3715,7 @@ static zskiplist *zslCreate(void) {
     
     zsl = zmalloc(sizeof(*zsl));
     zsl->level = 1;
+    zsl->length = 0;
     zsl->header = zslCreateNode(ZSKIPLIST_MAXLEVEL,0,NULL);
     for (j = 0; j < ZSKIPLIST_MAXLEVEL; j++)
         zsl->header->forward[j] = NULL;
@@ -3748,7 +3754,6 @@ static void zslInsert(zskiplist *zsl, double score, robj *obj) {
             x = x->forward[i];
         update[i] = x;
     }
-    x = x->forward[1];
     /* we assume the key is not already inside, since we allow duplicated
      * scores, and the re-insertion of score and redis object should never
      * happpen since the caller of zslInsert() should test in the hash table
@@ -3764,10 +3769,39 @@ static void zslInsert(zskiplist *zsl, double score, robj *obj) {
         x->forward[i] = update[i]->forward[i];
         update[i]->forward[i] = x;
     }
+    zsl->length++;
 }
 
 static int zslDelete(zskiplist *zsl, double score, robj *obj) {
-    return 1;
+    zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
+    int i;
+
+    x = zsl->header;
+    for (i = zsl->level-1; i >= 0; i--) {
+        while (x->forward[i] && x->forward[i]->score < score)
+            x = x->forward[i];
+        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. */
+    x = x->forward[0];
+    while(x->score == score) {
+        if (compareStringObjects(x->obj,obj) == 0) {
+            for (i = 0; i < zsl->level; i++) {
+                if (update[i]->forward[i] != x) break;
+                update[i]->forward[i] = x->forward[i];
+            }
+            zslFreeNode(x);
+            while(zsl->level > 1 && zsl->header->forward[zsl->level-1] == NULL)
+                zsl->level--;
+            zsl->length--;
+            return 1;
+        } else {
+            x = x->forward[0];
+            if (!x) return 0; /* end of the list reached, not found */
+        }
+    }
+    return 0; /* not found */
 }
 
 /* The actual Z-commands implementations */
@@ -3809,7 +3843,7 @@ static void zaddCommand(redisClient *c) {
         if (*score != *oldscore) {
             int deleted;
 
-            deleted = zslDelete(zs->zsl,*score,c->argv[3]);
+            deleted = zslDelete(zs->zsl,*oldscore,c->argv[3]);
             assert(deleted != 0);
             zslInsert(zs->zsl,*score,c->argv[3]);
             incrRefCount(c->argv[3]);
@@ -3819,6 +3853,76 @@ static void zaddCommand(redisClient *c) {
     }
 }
 
+static void zrangeCommand(redisClient *c) {
+    robj *o;
+    int start = atoi(c->argv[2]->ptr);
+    int end = atoi(c->argv[3]->ptr);
+
+    o = lookupKeyRead(c->db,c->argv[1]);
+    if (o == NULL) {
+        addReply(c,shared.nullmultibulk);
+    } else {
+        if (o->type != REDIS_ZSET) {
+            addReply(c,shared.wrongtypeerr);
+        } else {
+            zset *zsetobj = o->ptr;
+            zskiplist *zsl = zsetobj->zsl;
+            zskiplistNode *ln;
+
+            int llen = zsl->length;
+            int rangelen, j;
+            robj *ele;
+
+            /* 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) {
+                /* Out of range start or start > end result in empty list */
+                addReply(c,shared.emptymultibulk);
+                return;
+            }
+            if (end >= llen) end = llen-1;
+            rangelen = (end-start)+1;
+
+            /* Return the result in form of a multi-bulk reply */
+            ln = zsl->header->forward[0];
+            while (start--)
+                ln = ln->forward[0];
+
+            addReplySds(c,sdscatprintf(sdsempty(),"*%d\r\n",rangelen));
+            for (j = 0; j < rangelen; j++) {
+                ele = ln->obj;
+                addReplyBulkLen(c,ele);
+                addReply(c,ele);
+                addReply(c,shared.crlf);
+                ln = ln->forward[0];
+            }
+        }
+    }
+}
+
+static void zlenCommand(redisClient *c) {
+    robj *o;
+    zset *zs;
+    
+    o = lookupKeyRead(c->db,c->argv[1]);
+    if (o == NULL) {
+        addReply(c,shared.czero);
+        return;
+    } else {
+        if (o->type != REDIS_ZSET) {
+            addReply(c,shared.wrongtypeerr);
+        } else {
+            zs = o->ptr;
+            addReplySds(c,sdscatprintf(sdsempty(),":%d\r\n",zs->zsl->length));
+        }
+    }
+}
+
 /* ========================= Non type-specific commands  ==================== */
 
 static void flushdbCommand(redisClient *c) {
@@ -4875,6 +4979,7 @@ static struct redisFunctionSym symsTable[] = {
 {"zslDelete",(unsigned long)zslDelete},
 {"createZsetObject",(unsigned long)createZsetObject},
 {"zaddCommand",(unsigned long)zaddCommand},
+{"zrangeCommand",(unsigned long)zrangeCommand},
 {NULL,0}
 };