X-Git-Url: https://git.saurik.com/redis.git/blobdiff_plain/5d1b4fb6983f5acd9cfb6ee5f5715547688448d2..cc4c964b33ac108c9541e31eb2e2420addf9a82e:/src/t_zset.c

diff --git a/src/t_zset.c b/src/t_zset.c
index 6b2b6d9f..3c9ede1c 100644
--- a/src/t_zset.c
+++ b/src/t_zset.c
@@ -451,9 +451,48 @@ int zzlCompareElements(unsigned char *eptr, unsigned char *cstr, unsigned int cl
 
 unsigned int zzlLength(robj *zobj) {
     unsigned char *zl = zobj->ptr;
+    redisAssert(zobj->encoding == REDIS_ENCODING_ZIPLIST);
     return ziplistLen(zl)/2;
 }
 
+/* Move to next entry based on the values in eptr and sptr. Both are set to
+ * NULL when there is no next entry. */
+void zzlNext(unsigned char *zl, unsigned char **eptr, unsigned char **sptr) {
+    unsigned char *_eptr, *_sptr;
+    redisAssert(*eptr != NULL && *sptr != NULL);
+
+    _eptr = ziplistNext(zl,*sptr);
+    if (_eptr != NULL) {
+        _sptr = ziplistNext(zl,_eptr);
+        redisAssert(_sptr != NULL);
+    } else {
+        /* No next entry. */
+        _sptr = NULL;
+    }
+
+    *eptr = _eptr;
+    *sptr = _sptr;
+}
+
+/* Move to the previous entry based on the values in eptr and sptr. Both are
+ * set to NULL when there is no next entry. */
+void zzlPrev(unsigned char *zl, unsigned char **eptr, unsigned char **sptr) {
+    unsigned char *_eptr, *_sptr;
+    redisAssert(*eptr != NULL && *sptr != NULL);
+
+    _sptr = ziplistPrev(zl,*eptr);
+    if (_sptr != NULL) {
+        _eptr = ziplistPrev(zl,_sptr);
+        redisAssert(_eptr != NULL);
+    } else {
+        /* No previous entry. */
+        _eptr = NULL;
+    }
+
+    *eptr = _eptr;
+    *sptr = _sptr;
+}
+
 /* Returns if there is a part of the zset is in range. Should only be used
  * internally by zzlFirstInRange and zzlLastInRange. */
 int zzlIsInRange(unsigned char *zl, zrangespec *range) {
@@ -620,7 +659,7 @@ int zzlInsert(robj *zobj, robj *ele, double score) {
             break;
         } else if (s == score) {
             /* Ensure lexicographical ordering for elements. */
-            if (zzlCompareElements(eptr,ele->ptr,sdslen(ele->ptr)) < 0) {
+            if (zzlCompareElements(eptr,ele->ptr,sdslen(ele->ptr)) > 0) {
                 zzlInsertAt(zobj,ele,score,eptr);
                 break;
             }
@@ -666,6 +705,14 @@ unsigned long zzlDeleteRangeByScore(robj *zobj, zrangespec range) {
     return deleted;
 }
 
+/* Delete all the elements with rank between start and end from the skiplist.
+ * Start and end are inclusive. Note that start and end need to be 1-based */
+unsigned long zzlDeleteRangeByRank(robj *zobj, unsigned int start, unsigned int end) {
+    unsigned int num = (end-start)+1;
+    zobj->ptr = ziplistDeleteRange(zobj->ptr,2*(start-1),2*num);
+    return num;
+}
+
 /*-----------------------------------------------------------------------------
  * Common sorted set API
  *----------------------------------------------------------------------------*/
@@ -682,6 +729,85 @@ int zsLength(robj *zobj) {
     return length;
 }
 
+void zsConvert(robj *zobj, int encoding) {
+    zset *zs;
+    zskiplistNode *node, *next;
+    robj *ele;
+    double score;
+
+    if (zobj->encoding == encoding) return;
+    if (zobj->encoding == REDIS_ENCODING_ZIPLIST) {
+        unsigned char *zl = zobj->ptr;
+        unsigned char *eptr, *sptr;
+        unsigned char *vstr;
+        unsigned int vlen;
+        long long vlong;
+
+        if (encoding != REDIS_ENCODING_RAW)
+            redisPanic("Unknown target encoding");
+
+        zs = zmalloc(sizeof(*zs));
+        zs->dict = dictCreate(&zsetDictType,NULL);
+        zs->zsl = zslCreate();
+
+        eptr = ziplistIndex(zl,0);
+        redisAssert(eptr != NULL);
+        sptr = ziplistNext(zl,eptr);
+        redisAssert(sptr != NULL);
+
+        while (eptr != NULL) {
+            score = zzlGetScore(sptr);
+            redisAssert(ziplistGet(eptr,&vstr,&vlen,&vlong));
+            if (vstr == NULL)
+                ele = createStringObjectFromLongLong(vlong);
+            else
+                ele = createStringObject((char*)vstr,vlen);
+
+            /* Has incremented refcount since it was just created. */
+            node = zslInsert(zs->zsl,score,ele);
+            redisAssert(dictAdd(zs->dict,ele,&node->score) == DICT_OK);
+            incrRefCount(ele); /* Added to dictionary. */
+            zzlNext(zl,&eptr,&sptr);
+        }
+
+        zfree(zobj->ptr);
+        zobj->ptr = zs;
+        zobj->encoding = REDIS_ENCODING_RAW;
+    } else if (zobj->encoding == REDIS_ENCODING_RAW) {
+        unsigned char *zl = ziplistNew();
+
+        if (encoding != REDIS_ENCODING_ZIPLIST)
+            redisPanic("Unknown target encoding");
+
+        /* Approach similar to zslFree(), since we want to free the skiplist at
+         * the same time as creating the ziplist. */
+        zs = zobj->ptr;
+        dictRelease(zs->dict);
+        node = zs->zsl->header->level[0].forward;
+        zfree(zs->zsl->header);
+        zfree(zs->zsl);
+
+        /* Immediately store pointer to ziplist in object because it will
+         * change because of reallocations when pushing to the ziplist. */
+        zobj->ptr = zl;
+
+        while (node) {
+            ele = getDecodedObject(node->obj);
+            redisAssert(zzlInsertAt(zobj,ele,node->score,NULL) == REDIS_OK);
+            decrRefCount(ele);
+
+            next = node->level[0].forward;
+            zslFreeNode(node);
+            node = next;
+        }
+
+        zfree(zs);
+        zobj->encoding = REDIS_ENCODING_ZIPLIST;
+    } else {
+        redisPanic("Unknown sorted set encoding");
+    }
+}
+
 /*-----------------------------------------------------------------------------
  * Sorted set commands 
  *----------------------------------------------------------------------------*/
@@ -700,7 +826,13 @@ void zaddGenericCommand(redisClient *c, int incr) {
 
     zobj = lookupKeyWrite(c->db,key);
     if (zobj == NULL) {
-        zobj = createZsetZiplistObject();
+        if (server.zset_max_ziplist_entries == 0 ||
+            server.zset_max_ziplist_value < sdslen(c->argv[3]->ptr))
+        {
+            zobj = createZsetObject();
+        } else {
+            zobj = createZsetZiplistObject();
+        }
         dbAdd(c->db,key,zobj);
     } else {
         if (zobj->type != REDIS_ZSET) {
@@ -739,7 +871,13 @@ void zaddGenericCommand(redisClient *c, int incr) {
             else /* ZADD */
                 addReply(c,shared.czero);
         } else {
+            /* Optimize: check if the element is too large or the list becomes
+             * too long *before* executing zzlInsert. */
             redisAssert(zzlInsert(zobj,ele,score) == REDIS_OK);
+            if (zzlLength(zobj) > server.zset_max_ziplist_entries)
+                zsConvert(zobj,REDIS_ENCODING_RAW);
+            if (sdslen(ele->ptr) > server.zset_max_ziplist_value)
+                zsConvert(zobj,REDIS_ENCODING_RAW);
 
             signalModifiedKey(c->db,key);
             server.dirty++;
@@ -892,22 +1030,21 @@ void zremrangebyscoreCommand(redisClient *c) {
 }
 
 void zremrangebyrankCommand(redisClient *c) {
+    robj *key = c->argv[1];
+    robj *zobj;
     long start;
     long end;
     int llen;
-    long deleted;
-    robj *zsetobj;
-    zset *zs;
+    unsigned long deleted;
 
     if ((getLongFromObjectOrReply(c, c->argv[2], &start, NULL) != REDIS_OK) ||
         (getLongFromObjectOrReply(c, c->argv[3], &end, NULL) != REDIS_OK)) return;
 
-    if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
-        checkType(c,zsetobj,REDIS_ZSET)) return;
-    zs = zsetobj->ptr;
-    llen = zs->zsl->length;
+    if ((zobj = lookupKeyWriteOrReply(c,key,shared.czero)) == NULL ||
+        checkType(c,zobj,REDIS_ZSET)) return;
 
-    /* convert negative indexes */
+    /* Sanitize indexes. */
+    llen = zsLength(zobj);
     if (start < 0) start = llen+start;
     if (end < 0) end = llen+end;
     if (start < 0) start = 0;
@@ -920,14 +1057,23 @@ void zremrangebyrankCommand(redisClient *c) {
     }
     if (end >= llen) end = llen-1;
 
-    /* increment start and end because zsl*Rank functions
-     * use 1-based rank */
-    deleted = zslDeleteRangeByRank(zs->zsl,start+1,end+1,zs->dict);
-    if (htNeedsResize(zs->dict)) dictResize(zs->dict);
-    if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
-    if (deleted) signalModifiedKey(c->db,c->argv[1]);
+    if (zobj->encoding == REDIS_ENCODING_ZIPLIST) {
+        /* Correct for 1-based rank. */
+        deleted = zzlDeleteRangeByRank(zobj,start+1,end+1);
+    } else if (zobj->encoding == REDIS_ENCODING_RAW) {
+        zset *zs = zobj->ptr;
+
+        /* Correct for 1-based rank. */
+        deleted = zslDeleteRangeByRank(zs->zsl,start+1,end+1,zs->dict);
+        if (htNeedsResize(zs->dict)) dictResize(zs->dict);
+        if (dictSize(zs->dict) == 0) dbDelete(c->db,key);
+    } else {
+        redisPanic("Unknown sorted set encoding");
+    }
+
+    if (deleted) signalModifiedKey(c->db,key);
     server.dirty += deleted;
-    addReplyLongLong(c, deleted);
+    addReplyLongLong(c,deleted);
 }
 
 typedef struct {
@@ -1202,35 +1348,24 @@ void zrangeGenericCommand(redisClient *c, int reverse) {
         else
             eptr = ziplistIndex(zl,2*start);
 
+        redisAssert(eptr != NULL);
+        sptr = ziplistNext(zl,eptr);
+
         while (rangelen--) {
-            redisAssert(eptr != NULL);
+            redisAssert(eptr != NULL && sptr != NULL);
             redisAssert(ziplistGet(eptr,&vstr,&vlen,&vlong));
             if (vstr == NULL)
                 addReplyBulkLongLong(c,vlong);
             else
                 addReplyBulkCBuffer(c,vstr,vlen);
 
-            if (withscores) {
-                sptr = ziplistNext(zl,eptr);
-                redisAssert(sptr != NULL);
+            if (withscores)
                 addReplyDouble(c,zzlGetScore(sptr));
-            }
 
-            if (reverse) {
-                /* Move to previous element by moving to the score of previous
-                 * element. When NULL, we know there also is no element. */
-                sptr = ziplistPrev(zl,eptr);
-                if (sptr != NULL) {
-                    eptr = ziplistPrev(zl,sptr);
-                    redisAssert(eptr != NULL);
-                } else {
-                    eptr = NULL;
-                }
-            } else {
-                sptr = ziplistNext(zl,eptr);
-                redisAssert(sptr != NULL);
-                eptr = ziplistNext(zl,sptr);
-            }
+            if (reverse)
+                zzlPrev(zl,&eptr,&sptr);
+            else
+                zzlNext(zl,&eptr,&sptr);
         }
 
     } else if (zobj->encoding == REDIS_ENCODING_RAW) {
@@ -1275,10 +1410,8 @@ void zrevrangeCommand(redisClient *c) {
  * If "justcount", only the number of elements in the range is returned. */
 void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) {
     zrangespec range;
-    robj *o, *emptyreply;
-    zset *zsetobj;
-    zskiplist *zsl;
-    zskiplistNode *ln;
+    robj *key = c->argv[1];
+    robj *emptyreply, *zobj;
     int offset = 0, limit = -1;
     int withscores = 0;
     unsigned long rangelen = 0;
@@ -1322,61 +1455,132 @@ void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) {
 
     /* Ok, lookup the key and get the range */
     emptyreply = justcount ? shared.czero : shared.emptymultibulk;
-    if ((o = lookupKeyReadOrReply(c,c->argv[1],emptyreply)) == NULL ||
-        checkType(c,o,REDIS_ZSET)) return;
-    zsetobj = o->ptr;
-    zsl = zsetobj->zsl;
+    if ((zobj = lookupKeyReadOrReply(c,key,emptyreply)) == NULL ||
+        checkType(c,zobj,REDIS_ZSET)) return;
 
-    /* If reversed, get the last node in range as starting point. */
-    if (reverse) {
-        ln = zslLastInRange(zsl,range);
-    } else {
-        ln = zslFirstInRange(zsl,range);
-    }
+    if (zobj->encoding == REDIS_ENCODING_ZIPLIST) {
+        unsigned char *zl = zobj->ptr;
+        unsigned char *eptr, *sptr;
+        unsigned char *vstr;
+        unsigned int vlen;
+        long long vlong;
+        double score;
 
-    /* No "first" element in the specified interval. */
-    if (ln == NULL) {
-        addReply(c,emptyreply);
-        return;
-    }
+        /* If reversed, get the last node in range as starting point. */
+        if (reverse)
+            eptr = zzlLastInRange(zobj,range);
+        else
+            eptr = zzlFirstInRange(zobj,range);
 
-    /* We don't know in advance how many matching elements there are in the
-     * list, so we push this object that will represent the multi-bulk length
-     * in the output buffer, and will "fix" it later */
-    if (!justcount)
-        replylen = addDeferredMultiBulkLength(c);
+        /* No "first" element in the specified interval. */
+        if (eptr == NULL) {
+            addReply(c,emptyreply);
+            return;
+        }
 
-    /* If there is an offset, just traverse the number of elements without
-     * checking the score because that is done in the next loop. */
-    while(ln && offset--) {
-        ln = reverse ? ln->backward : ln->level[0].forward;
-    }
+        /* Get score pointer for the first element. */
+        redisAssert(eptr != NULL);
+        sptr = ziplistNext(zl,eptr);
 
-    while (ln && limit--) {
-        /* Abort when the node is no longer in range. */
-        if (reverse) {
-            if (!zslValueGteMin(ln->score,&range)) break;
-        } else {
-            if (!zslValueLteMax(ln->score,&range)) break;
+        /* We don't know in advance how many matching elements there are in the
+         * list, so we push this object that will represent the multi-bulk
+         * length in the output buffer, and will "fix" it later */
+        if (!justcount)
+            replylen = addDeferredMultiBulkLength(c);
+
+        /* If there is an offset, just traverse the number of elements without
+         * checking the score because that is done in the next loop. */
+        while (eptr && offset--)
+            if (reverse)
+                zzlPrev(zl,&eptr,&sptr);
+            else
+                zzlNext(zl,&eptr,&sptr);
+
+        while (eptr && limit--) {
+            score = zzlGetScore(sptr);
+
+            /* Abort when the node is no longer in range. */
+            if (reverse) {
+                if (!zslValueGteMin(score,&range)) break;
+            } else {
+                if (!zslValueLteMax(score,&range)) break;
+            }
+
+            /* Do our magic */
+            rangelen++;
+            if (!justcount) {
+                redisAssert(ziplistGet(eptr,&vstr,&vlen,&vlong));
+                if (vstr == NULL)
+                    addReplyBulkLongLong(c,vlong);
+                else
+                    addReplyBulkCBuffer(c,vstr,vlen);
+
+                if (withscores)
+                    addReplyDouble(c,score);
+            }
+
+            /* Move to next node */
+            if (reverse)
+                zzlPrev(zl,&eptr,&sptr);
+            else
+                zzlNext(zl,&eptr,&sptr);
         }
+    } else if (zobj->encoding == REDIS_ENCODING_RAW) {
+        zset *zs = zobj->ptr;
+        zskiplist *zsl = zs->zsl;
+        zskiplistNode *ln;
 
-        /* Do our magic */
-        rangelen++;
-        if (!justcount) {
-            addReplyBulk(c,ln->obj);
-            if (withscores)
-                addReplyDouble(c,ln->score);
+        /* If reversed, get the last node in range as starting point. */
+        if (reverse)
+            ln = zslLastInRange(zsl,range);
+        else
+            ln = zslFirstInRange(zsl,range);
+
+        /* No "first" element in the specified interval. */
+        if (ln == NULL) {
+            addReply(c,emptyreply);
+            return;
         }
 
-        /* Move to next node */
-        ln = reverse ? ln->backward : ln->level[0].forward;
+        /* We don't know in advance how many matching elements there are in the
+         * list, so we push this object that will represent the multi-bulk
+         * length in the output buffer, and will "fix" it later */
+        if (!justcount)
+            replylen = addDeferredMultiBulkLength(c);
+
+        /* If there is an offset, just traverse the number of elements without
+         * checking the score because that is done in the next loop. */
+        while (ln && offset--)
+            ln = reverse ? ln->backward : ln->level[0].forward;
+
+        while (ln && limit--) {
+            /* Abort when the node is no longer in range. */
+            if (reverse) {
+                if (!zslValueGteMin(ln->score,&range)) break;
+            } else {
+                if (!zslValueLteMax(ln->score,&range)) break;
+            }
+
+            /* Do our magic */
+            rangelen++;
+            if (!justcount) {
+                addReplyBulk(c,ln->obj);
+                if (withscores)
+                    addReplyDouble(c,ln->score);
+            }
+
+            /* Move to next node */
+            ln = reverse ? ln->backward : ln->level[0].forward;
+        }
+    } else {
+        redisPanic("Unknown sorted set encoding");
     }
 
     if (justcount) {
         addReplyLongLong(c,(long)rangelen);
     } else {
-        setDeferredMultiBulkLength(c,replylen,
-             withscores ? (rangelen*2) : rangelen);
+        if (withscores) rangelen *= 2;
+        setDeferredMultiBulkLength(c,replylen,rangelen);
     }
 }
 
@@ -1393,66 +1597,103 @@ void zcountCommand(redisClient *c) {
 }
 
 void zcardCommand(redisClient *c) {
-    robj *o;
-    zset *zs;
+    robj *key = c->argv[1];
+    robj *zobj;
 
-    if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
-        checkType(c,o,REDIS_ZSET)) return;
+    if ((zobj = lookupKeyReadOrReply(c,key,shared.czero)) == NULL ||
+        checkType(c,zobj,REDIS_ZSET)) return;
 
-    zs = o->ptr;
-    addReplyLongLong(c,zs->zsl->length);
+    addReplyLongLong(c,zsLength(zobj));
 }
 
 void zscoreCommand(redisClient *c) {
-    robj *o;
-    zset *zs;
-    dictEntry *de;
+    robj *key = c->argv[1];
+    robj *zobj;
+    double score;
 
-    if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
-        checkType(c,o,REDIS_ZSET)) return;
+    if ((zobj = lookupKeyReadOrReply(c,key,shared.nullbulk)) == NULL ||
+        checkType(c,zobj,REDIS_ZSET)) return;
 
-    zs = o->ptr;
-    c->argv[2] = tryObjectEncoding(c->argv[2]);
-    de = dictFind(zs->dict,c->argv[2]);
-    if (!de) {
-        addReply(c,shared.nullbulk);
-    } else {
-        double *score = dictGetEntryVal(de);
+    if (zobj->encoding == REDIS_ENCODING_ZIPLIST) {
+        if (zzlFind(zobj,c->argv[2],&score) != NULL)
+            addReplyDouble(c,score);
+        else
+            addReply(c,shared.nullbulk);
+    } else if (zobj->encoding == REDIS_ENCODING_RAW) {
+        zset *zs = zobj->ptr;
+        dictEntry *de;
 
-        addReplyDouble(c,*score);
+        c->argv[2] = tryObjectEncoding(c->argv[2]);
+        de = dictFind(zs->dict,c->argv[2]);
+        if (de != NULL) {
+            score = *(double*)dictGetEntryVal(de);
+            addReplyDouble(c,score);
+        } else {
+            addReply(c,shared.nullbulk);
+        }
+    } else {
+        redisPanic("Unknown sorted set encoding");
     }
 }
 
 void zrankGenericCommand(redisClient *c, int reverse) {
-    robj *o;
-    zset *zs;
-    zskiplist *zsl;
-    dictEntry *de;
+    robj *key = c->argv[1];
+    robj *ele = c->argv[2];
+    robj *zobj;
+    unsigned long llen;
     unsigned long rank;
-    double *score;
 
-    if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
-        checkType(c,o,REDIS_ZSET)) return;
+    if ((zobj = lookupKeyReadOrReply(c,key,shared.nullbulk)) == NULL ||
+        checkType(c,zobj,REDIS_ZSET)) return;
+    llen = zsLength(zobj);
 
-    zs = o->ptr;
-    zsl = zs->zsl;
-    c->argv[2] = tryObjectEncoding(c->argv[2]);
-    de = dictFind(zs->dict,c->argv[2]);
-    if (!de) {
-        addReply(c,shared.nullbulk);
-        return;
-    }
+    redisAssert(ele->encoding == REDIS_ENCODING_RAW);
+    if (zobj->encoding == REDIS_ENCODING_ZIPLIST) {
+        unsigned char *zl = zobj->ptr;
+        unsigned char *eptr, *sptr;
 
-    score = dictGetEntryVal(de);
-    rank = zslGetRank(zsl, *score, c->argv[2]);
-    if (rank) {
-        if (reverse) {
-            addReplyLongLong(c, zsl->length - rank);
+        eptr = ziplistIndex(zl,0);
+        redisAssert(eptr != NULL);
+        sptr = ziplistNext(zl,eptr);
+        redisAssert(sptr != NULL);
+
+        rank = 1;
+        while(eptr != NULL) {
+            if (ziplistCompare(eptr,ele->ptr,sdslen(ele->ptr)))
+                break;
+            rank++;
+            zzlNext(zl,&eptr,&sptr);
+        }
+
+        if (eptr != NULL) {
+            if (reverse)
+                addReplyLongLong(c,llen-rank);
+            else
+                addReplyLongLong(c,rank-1);
+        } else {
+            addReply(c,shared.nullbulk);
+        }
+    } else if (zobj->encoding == REDIS_ENCODING_RAW) {
+        zset *zs = zobj->ptr;
+        zskiplist *zsl = zs->zsl;
+        dictEntry *de;
+        double score;
+
+        ele = c->argv[2] = tryObjectEncoding(c->argv[2]);
+        de = dictFind(zs->dict,ele);
+        if (de != NULL) {
+            score = *(double*)dictGetEntryVal(de);
+            rank = zslGetRank(zsl,score,ele);
+            redisAssert(rank); /* Existing elements always have a rank. */
+            if (reverse)
+                addReplyLongLong(c,llen-rank);
+            else
+                addReplyLongLong(c,rank-1);
         } else {
-            addReplyLongLong(c, rank-1);
+            addReply(c,shared.nullbulk);
         }
     } else {
-        addReply(c,shared.nullbulk);
+        redisPanic("Unknown sorted set encoding");
     }
 }