Modules: zset lex iterator #3.

3 files changed, 49 insertions, 3 deletions

diff --git a/src/module.c b/src/module.c
index 039bba227..635caa7b4 100644
--- a/src/module.c
+++ b/src/module.c
@@ -520,7 +520,7 @@ int RM_ReplyWithSimpleString(RedisModuleCtx *ctx, const char *msg) {
  * of the array.
  *
  * The function always returns REDISMODULE_OK. */
-int RM_ReplyWithArray(RedisModuleCtx *ctx, int len) {
+int RM_ReplyWithArray(RedisModuleCtx *ctx, long len) {
     addReplyMultiBulkLen(ctx->client,len);
     return REDISMODULE_OK;
 }
@@ -1186,7 +1186,6 @@ int zsetInitLexRange(RedisModuleKey *key, RedisModuleString *min, RedisModuleStr
     if (!key->value || key->value->type != OBJ_ZSET) return REDISMODULE_ERR;
 
     RM_ZsetRangeStop(key);
-    key->ztype = REDISMODULE_ZSET_RANGE_LEX;
     key->zer = 0;
 
     /* Setup the range structure used by the sorted set core implementation
@@ -1194,6 +1193,10 @@ int zsetInitLexRange(RedisModuleKey *key, RedisModuleString *min, RedisModuleStr
     zlexrangespec *zlrs = &key->zlrs;
     if (zslParseLexRange(min, max, zlrs) == C_ERR) return REDISMODULE_ERR;
 
+    /* Set the range type to lex only after successfully parsing the range,
+     * otherwise we don't want the zlexrangespec to be freed. */
+    key->ztype = REDISMODULE_ZSET_RANGE_LEX;
+
     if (key->value->encoding == OBJ_ENCODING_ZIPLIST) {
         key->zcurrent = first ? zzlFirstInLexRange(key->value->ptr,zlrs) :
                                 zzlLastInLexRange(key->value->ptr,zlrs);
@@ -1206,6 +1209,7 @@ int zsetInitLexRange(RedisModuleKey *key, RedisModuleString *min, RedisModuleStr
         serverPanic("Unsupported zset encoding");
     }
     if (key->zcurrent == NULL) key->zer = 1;
+
     return REDISMODULE_OK;
 }
 
@@ -1838,6 +1842,8 @@ void moduleRegisterCoreAPI(void) {
     REGISTER_API(ZsetRangeStop);
     REGISTER_API(ZsetFirstInScoreRange);
     REGISTER_API(ZsetLastInScoreRange);
+    REGISTER_API(ZsetFirstInLexRange);
+    REGISTER_API(ZsetLastInLexRange);
     REGISTER_API(ZsetRangeCurrentElement);
     REGISTER_API(ZsetRangeNext);
     REGISTER_API(ZsetRangePrev);
diff --git a/src/modules/helloworld.c b/src/modules/helloworld.c
index 738983250..785d401b2 100644
--- a/src/modules/helloworld.c
+++ b/src/modules/helloworld.c
@@ -334,6 +334,7 @@ int HelloMoreExpire_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv,
  * should match.*/
 int HelloZsumRange_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
     double score_start, score_end;
+    if (argc != 4) return RedisModule_WrongArity(ctx);
 
     if (RedisModule_StringToDouble(argv[2],&score_start) != REDISMODULE_OK ||
         RedisModule_StringToDouble(argv[3],&score_end) != REDISMODULE_OK)
@@ -379,6 +380,41 @@ int HelloZsumRange_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, i
     return REDISMODULE_OK;
 }
 
+/* HELLO.LEXRANGE key min_lex max_lex min_age max_age
+ * This command expects a sorted set stored at key in the following form:
+ * - All the elements have score 0.
+ * - Elements are pairs of "<name>:<age>", for example "Anna:52".
+ * The command will return all the sorted set items that are lexicographically
+ * between the specified range (using the same format as ZRANGEBYLEX)
+ * and having an age between min_age and max_age. */
+int HelloZsumRange_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+    RedisModule_AutoMemory(ctx); /* Use automatic memory management. */
+
+    if (argc != 6) return RedisModule_WrongArity(ctx);
+
+    RedisModuleKey *key = RedisModule_OpenKey(ctx,argv[1],
+        REDISMODULE_READ|REDISMODULE_WRITE);
+    if (RedisModule_KeyType(key) != REDISMODULE_KEYTYPE_ZSET) {
+        return RedisModule_ReplyWithError(ctx,REDISMODULE_ERRORMSG_WRONGTYPE);
+    }
+
+    RedisModule_ReplyWithArray(ctx,REDISMODULE_POSTPONED_ARRAY_LEN);
+    RedisModule_ZsetFirstInLexRange(key,argv[2],argv[3]);
+    int arraylen = 0;
+    while(!RedisModule_ZsetRangeEndReached(key)) {
+        double score;
+        RedisModuleString *ele = RedisModule_ZsetRangeCurrentElement(key,&score);
+        RedisModule_ReplyWithString(ctx,ele);
+        RedisModule_FreeString(ctx,ele);
+        RedisModule_ZsetRangeNext(key);
+        arraylen++;
+    }
+    RedisModule_ZsetRangeStop(key);
+    RedisModule_SetArrayLength(ctx,arraylen);
+    RedisModule_CloseKey(key);
+    return REDISMODULE_OK;
+}
+
 /* This function must be present on each Redis module. It is used in order to
  * register the commands into the Redis server. */
 int RedisModule_OnLoad(RedisModuleCtx *ctx) {
diff --git a/src/redismodule.h b/src/redismodule.h
index bdc85b137..ee6223dd3 100644
--- a/src/redismodule.h
+++ b/src/redismodule.h
@@ -98,7 +98,7 @@ void REDISMODULE_API_FUNC(RedisModule_FreeString)(RedisModuleCtx *ctx, RedisModu
 const char *REDISMODULE_API_FUNC(RedisModule_StringPtrLen)(RedisModuleString *str, size_t *len);
 int REDISMODULE_API_FUNC(RedisModule_ReplyWithError)(RedisModuleCtx *ctx, const char *err);
 int REDISMODULE_API_FUNC(RedisModule_ReplyWithSimpleString)(RedisModuleCtx *ctx, const char *msg);
-int REDISMODULE_API_FUNC(RedisModule_ReplyWithArray)(RedisModuleCtx *ctx, int len);
+int REDISMODULE_API_FUNC(RedisModule_ReplyWithArray)(RedisModuleCtx *ctx, long len);
 int REDISMODULE_API_FUNC(RedisModule_ReplyWithStringBuffer)(RedisModuleCtx *ctx, const char *buf, size_t len);
 int REDISMODULE_API_FUNC(RedisModule_ReplyWithString)(RedisModuleCtx *ctx, RedisModuleString *str);
 int REDISMODULE_API_FUNC(RedisModule_ReplyWithNull)(RedisModuleCtx *ctx);
@@ -124,6 +124,8 @@ int REDISMODULE_API_FUNC(RedisModule_ZsetRem)(RedisModuleKey *key, RedisModuleSt
 void REDISMODULE_API_FUNC(RedisModule_ZsetRangeStop)(RedisModuleKey *key);
 int REDISMODULE_API_FUNC(RedisModule_ZsetFirstInScoreRange)(RedisModuleKey *key, double min, double max, int minex, int maxex);
 int REDISMODULE_API_FUNC(RedisModule_ZsetLastInScoreRange)(RedisModuleKey *key, double min, double max, int minex, int maxex);
+int REDISMODULE_API_FUNC(RedisModule_ZsetFirstInLexRange)(RedisModuleKey *key, RedisModuleString *min, RedisModuleString *max);
+int REDISMODULE_API_FUNC(RedisModule_ZsetLastInLexRange)(RedisModuleKey *key, RedisModuleString *min, RedisModuleString *max);
 RedisModuleString *REDISMODULE_API_FUNC(RedisModule_ZsetRangeCurrentElement)(RedisModuleKey *key, double *score);
 int REDISMODULE_API_FUNC(RedisModule_ZsetRangeNext)(RedisModuleKey *key);
 int REDISMODULE_API_FUNC(RedisModule_ZsetRangePrev)(RedisModuleKey *key);
@@ -184,6 +186,8 @@ static int RedisModule_Init(RedisModuleCtx *ctx, const char *name, int ver, int
     REDISMODULE_GET_API(ZsetRangeStop);
     REDISMODULE_GET_API(ZsetFirstInScoreRange);
     REDISMODULE_GET_API(ZsetLastInScoreRange);
+    REDISMODULE_GET_API(ZsetFirstInLexRange);
+    REDISMODULE_GET_API(ZsetLastInLexRange);
     REDISMODULE_GET_API(ZsetRangeCurrentElement);
     REDISMODULE_GET_API(ZsetRangeNext);
     REDISMODULE_GET_API(ZsetRangePrev);