ZLEXCOUNT implemented.

Like ZCOUNT for lexicographical ranges.

3 files changed, 75 insertions, 0 deletions

diff --git a/src/redis.c b/src/redis.c
index cd0410f16..fe4618925 100644
--- a/src/redis.c
+++ b/src/redis.c
@@ -176,6 +176,7 @@ struct redisCommand redisCommandTable[] = {
     {"zrangebylex",zrangebylexCommand,-4,"r",0,NULL,1,1,1,0,0},
     {"zrevrangebylex",zrevrangebylexCommand,-4,"r",0,NULL,1,1,1,0,0},
     {"zcount",zcountCommand,4,"r",0,NULL,1,1,1,0,0},
+    {"zlexcount",zlexcountCommand,4,"r",0,NULL,1,1,1,0,0},
     {"zrevrange",zrevrangeCommand,-4,"r",0,NULL,1,1,1,0,0},
     {"zcard",zcardCommand,2,"r",0,NULL,1,1,1,0,0},
     {"zscore",zscoreCommand,3,"r",0,NULL,1,1,1,0,0},
diff --git a/src/redis.h b/src/redis.h
index bb1ad2ab6..344ac117a 100644
--- a/src/redis.h
+++ b/src/redis.h
@@ -1308,6 +1308,7 @@ void zrevrangebyscoreCommand(redisClient *c);
 void zrangebylexCommand(redisClient *c);
 void zrevrangebylexCommand(redisClient *c);
 void zcountCommand(redisClient *c);
+void zlexcountCommand(redisClient *c);
 void zrevrangeCommand(redisClient *c);
 void zcardCommand(redisClient *c);
 void zremCommand(redisClient *c);
diff --git a/src/t_zset.c b/src/t_zset.c
index de1e1103b..6f0f5ad79 100644
--- a/src/t_zset.c
+++ b/src/t_zset.c
@@ -2357,6 +2357,79 @@ void zcountCommand(redisClient *c) {
     addReplyLongLong(c, count);
 }
 
+void zlexcountCommand(redisClient *c) {
+    robj *key = c->argv[1];
+    robj *zobj;
+    zlexrangespec range;
+    int count = 0;
+
+    /* Parse the range arguments */
+    if (zslParseLexRange(c->argv[2],c->argv[3],&range) != REDIS_OK) {
+        addReplyError(c,"min or max not valid string range item");
+        return;
+    }
+
+    /* Lookup the sorted set */
+    if ((zobj = lookupKeyReadOrReply(c, key, shared.czero)) == NULL ||
+        checkType(c, zobj, REDIS_ZSET)) return;
+
+    if (zobj->encoding == REDIS_ENCODING_ZIPLIST) {
+        unsigned char *zl = zobj->ptr;
+        unsigned char *eptr, *sptr;
+
+        /* Use the first element in range as the starting point */
+        eptr = zzlFirstInLexRange(zl,range);
+
+        /* No "first" element */
+        if (eptr == NULL) {
+            addReply(c, shared.czero);
+            return;
+        }
+
+        /* First element is in range */
+        sptr = ziplistNext(zl,eptr);
+        redisAssertWithInfo(c,zobj,zzlLexValueLteMax(eptr,&range));
+
+        /* Iterate over elements in range */
+        while (eptr) {
+            /* Abort when the node is no longer in range. */
+            if (!zzlLexValueLteMax(eptr,&range)) {
+                break;
+            } else {
+                count++;
+                zzlNext(zl,&eptr,&sptr);
+            }
+        }
+    } else if (zobj->encoding == REDIS_ENCODING_SKIPLIST) {
+        zset *zs = zobj->ptr;
+        zskiplist *zsl = zs->zsl;
+        zskiplistNode *zn;
+        unsigned long rank;
+
+        /* Find first element in range */
+        zn = zslFirstInLexRange(zsl, range);
+
+        /* Use rank of first element, if any, to determine preliminary count */
+        if (zn != NULL) {
+            rank = zslGetRank(zsl, zn->score, zn->obj);
+            count = (zsl->length - (rank - 1));
+
+            /* Find last element in range */
+            zn = zslLastInLexRange(zsl, range);
+
+            /* Use rank of last element, if any, to determine the actual count */
+            if (zn != NULL) {
+                rank = zslGetRank(zsl, zn->score, zn->obj);
+                count -= (zsl->length - rank);
+            }
+        }
+    } else {
+        redisPanic("Unknown sorted set encoding");
+    }
+
+    addReplyLongLong(c, count);
+}
+
 /* This command implements ZRANGEBYLEX, ZREVRANGEBYLEX. */
 void genericZrangebylexCommand(redisClient *c, int reverse) {
     zlexrangespec range;