SDIFF/SDIFFSTORE implemnted unifying it with the implementation of SUNION/SUNIONSTORE

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},

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) {
+    if (dstkey) {
-        lenobj = createObject(REDIS_STRING,NULL);
-        addReply(c,lenobj);
-        decrRefCount(lenobj);
-    } else {
         /* 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 (op == REDIS_OP_UNION || j == 0) {
                 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);
                     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) {

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