diff --git a/redis-cli.c b/redis-cli.c index 524157cd60911c4e565bb36b0a36a8980eee5f06..e977368b19c7dc6faefe492cc7bfffd047572223 100644 --- a/redis-cli.c +++ b/redis-cli.c @@ -83,6 +83,8 @@ static struct redisCommand cmdTable[] = { {"sinterstore",-3,REDIS_CMD_INLINE}, {"sunion",-2,REDIS_CMD_INLINE}, {"sunionstore",-3,REDIS_CMD_INLINE}, + {"sdiff",-2,REDIS_CMD_INLINE}, + {"sdiffstore",-3,REDIS_CMD_INLINE}, {"smembers",2,REDIS_CMD_INLINE}, {"incrby",3,REDIS_CMD_INLINE}, {"decrby",3,REDIS_CMD_INLINE}, diff --git a/redis.c b/redis.c index 39e85124beefa486a859f7904a863d65936f9a48..6be82b4b008a7e1e30b2f263b5f6ef70e224cc4e 100644 --- a/redis.c +++ b/redis.c @@ -348,6 +348,8 @@ static void sinterCommand(redisClient *c); static void sinterstoreCommand(redisClient *c); static void sunionCommand(redisClient *c); static void sunionstoreCommand(redisClient *c); +static void sdiffCommand(redisClient *c); +static void sdiffstoreCommand(redisClient *c); static void syncCommand(redisClient *c); static void flushdbCommand(redisClient *c); static void flushallCommand(redisClient *c); @@ -391,6 +393,8 @@ static struct redisCommand cmdTable[] = { {"sinterstore",sinterstoreCommand,-3,REDIS_CMD_INLINE}, {"sunion",sunionCommand,-2,REDIS_CMD_INLINE}, {"sunionstore",sunionstoreCommand,-3,REDIS_CMD_INLINE}, + {"sdiff",sdiffCommand,-2,REDIS_CMD_INLINE}, + {"sdiffstore",sdiffstoreCommand,-3,REDIS_CMD_INLINE}, {"smembers",sinterCommand,2,REDIS_CMD_INLINE}, {"incrby",incrbyCommand,3,REDIS_CMD_INLINE}, {"decrby",decrbyCommand,3,REDIS_CMD_INLINE}, @@ -3057,14 +3061,17 @@ static void sinterstoreCommand(redisClient *c) { sinterGenericCommand(c,c->argv+2,c->argc-2,c->argv[1]); } -static void sunionGenericCommand(redisClient *c, robj **setskeys, int setsnum, robj *dstkey) { +#define REDIS_OP_UNION 0 +#define REDIS_OP_DIFF 1 + +static void sunionDiffGenericCommand(redisClient *c, robj **setskeys, int setsnum, robj *dstkey, int op) { dict **dv = zmalloc(sizeof(dict*)*setsnum); dictIterator *di; dictEntry *de; - robj *lenobj = NULL, *dstset = NULL; + robj *dstset = NULL; int j, cardinality = 0; - if (!dv) oom("sunionCommand"); + if (!dv) oom("sunionDiffGenericCommand"); for (j = 0; j < setsnum; j++) { robj *setobj; @@ -3093,11 +3100,7 @@ static void sunionGenericCommand(redisClient *c, robj **setskeys, int setsnum, r * the intersection set size, so we use a trick, append an empty object * to the output list and save the pointer to later modify it with the * right length */ - if (!dstkey) { - lenobj = createObject(REDIS_STRING,NULL); - addReply(c,lenobj); - decrRefCount(lenobj); - } else { + if (dstkey) { /* If we have a target key where to store the resulting set * create this key with an empty set inside */ deleteKey(c->db,dstkey); @@ -3119,22 +3122,39 @@ static void sunionGenericCommand(redisClient *c, robj **setskeys, int setsnum, r /* dictAdd will not add the same element multiple times */ ele = dictGetEntryKey(de); - if (dictAdd(dstset->ptr,ele,NULL) == DICT_OK) { - incrRefCount(ele); - if (!dstkey) { - addReplySds(c,sdscatprintf(sdsempty(), - "$%d\r\n",sdslen(ele->ptr))); - addReply(c,ele); - addReply(c,shared.crlf); + if (op == REDIS_OP_UNION || j == 0) { + if (dictAdd(dstset->ptr,ele,NULL) == DICT_OK) { + incrRefCount(ele); cardinality++; } + } else if (op == REDIS_OP_DIFF) { + if (dictDelete(dstset->ptr,ele) == DICT_OK) { + cardinality--; + } } } dictReleaseIterator(di); } + /* Output the content of the resulting set, if not in STORE mode */ + if (!dstkey) { + addReplySds(c,sdscatprintf(sdsempty(),"*%d\r\n",cardinality)); + di = dictGetIterator(dstset->ptr); + if (!di) oom("dictGetIterator"); + while((de = dictNext(di)) != NULL) { + robj *ele; + + ele = dictGetEntryKey(de); + addReplySds(c,sdscatprintf(sdsempty(), + "$%d\r\n",sdslen(ele->ptr))); + addReply(c,ele); + addReply(c,shared.crlf); + } + dictReleaseIterator(di); + } + + /* Cleanup */ if (!dstkey) { - lenobj->ptr = sdscatprintf(sdsempty(),"*%d\r\n",cardinality); decrRefCount(dstset); } else { addReply(c,shared.ok); @@ -3144,11 +3164,19 @@ static void sunionGenericCommand(redisClient *c, robj **setskeys, int setsnum, r } static void sunionCommand(redisClient *c) { - sunionGenericCommand(c,c->argv+1,c->argc-1,NULL); + sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_UNION); } static void sunionstoreCommand(redisClient *c) { - sunionGenericCommand(c,c->argv+2,c->argc-2,c->argv[1]); + sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_UNION); +} + +static void sdiffCommand(redisClient *c) { + sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_DIFF); +} + +static void sdiffstoreCommand(redisClient *c) { + sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_DIFF); } static void flushdbCommand(redisClient *c) { diff --git a/test-redis.tcl b/test-redis.tcl index 9e69acda0bc7e058cd9db0da059798572b69a71e..bd58cb2e7183bea9fc62334ecad0f474ceeaf7c6 100644 --- a/test-redis.tcl +++ b/test-redis.tcl @@ -498,6 +498,23 @@ proc main {server port} { lsort [$r sunion nokey1 set1 set2 nokey2] } [lsort -uniq "[$r smembers set1] [$r smembers set2]"] + test {SDIFF with two sets} { + for {set i 5} {$i < 1000} {incr i} { + $r sadd set4 $i + } + lsort [$r sdiff set1 set4] + } {0 1 2 3 4} + + test {SDIFF with three sets} { + $r sadd set5 0 + lsort [$r sdiff set1 set4 set5] + } {1 2 3 4} + + test {SDIFFSTORE with three sets} { + $r sdiffstore sres set1 set4 set5 + lsort [$r smembers sres] + } {1 2 3 4} + test {SAVE - make sure there are all the types as values} { $r lpush mysavelist hello $r lpush mysavelist world