src/t_set.c

   1 #include "redis.h"
   2
   3 /*-----------------------------------------------------------------------------
   4  * Set Commands
   5  *----------------------------------------------------------------------------*/
   6
   7 void saddCommand(redisClient *c) {
   8     robj *set;
   9
  10     set = lookupKeyWrite(c->db,c->argv[1]);
  11     if (set == NULL) {
  12         set = createSetObject();
  13         dbAdd(c->db,c->argv[1],set);
  14     } else {
  15         if (set->type != REDIS_SET) {
  16             addReply(c,shared.wrongtypeerr);
  17             return;
  18         }
  19     }
  20     if (dictAdd(set->ptr,c->argv[2],NULL) == DICT_OK) {
  21         incrRefCount(c->argv[2]);
  22         server.dirty++;
  23         addReply(c,shared.cone);
  24     } else {
  25         addReply(c,shared.czero);
  26     }
  27 }
  28
  29 void sremCommand(redisClient *c) {
  30     robj *set;
  31
  32     if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
  33         checkType(c,set,REDIS_SET)) return;
  34
  35     if (dictDelete(set->ptr,c->argv[2]) == DICT_OK) {
  36         server.dirty++;
  37         if (htNeedsResize(set->ptr)) dictResize(set->ptr);
  38         if (dictSize((dict*)set->ptr) == 0) dbDelete(c->db,c->argv[1]);
  39         addReply(c,shared.cone);
  40     } else {
  41         addReply(c,shared.czero);
  42     }
  43 }
  44
  45 void smoveCommand(redisClient *c) {
  46     robj *srcset, *dstset;
  47
  48     srcset = lookupKeyWrite(c->db,c->argv[1]);
  49     dstset = lookupKeyWrite(c->db,c->argv[2]);
  50
  51     /* If the source key does not exist return 0, if it's of the wrong type
  52      * raise an error */
  53     if (srcset == NULL || srcset->type != REDIS_SET) {
  54         addReply(c, srcset ? shared.wrongtypeerr : shared.czero);
  55         return;
  56     }
  57     /* Error if the destination key is not a set as well */
  58     if (dstset && dstset->type != REDIS_SET) {
  59         addReply(c,shared.wrongtypeerr);
  60         return;
  61     }
  62     /* Remove the element from the source set */
  63     if (dictDelete(srcset->ptr,c->argv[3]) == DICT_ERR) {
  64         /* Key not found in the src set! return zero */
  65         addReply(c,shared.czero);
  66         return;
  67     }
  68     if (dictSize((dict*)srcset->ptr) == 0 && srcset != dstset)
  69         dbDelete(c->db,c->argv[1]);
  70     server.dirty++;
  71     /* Add the element to the destination set */
  72     if (!dstset) {
  73         dstset = createSetObject();
  74         dbAdd(c->db,c->argv[2],dstset);
  75     }
  76     if (dictAdd(dstset->ptr,c->argv[3],NULL) == DICT_OK)
  77         incrRefCount(c->argv[3]);
  78     addReply(c,shared.cone);
  79 }
  80
  81 void sismemberCommand(redisClient *c) {
  82     robj *set;
  83
  84     if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
  85         checkType(c,set,REDIS_SET)) return;
  86
  87     if (dictFind(set->ptr,c->argv[2]))
  88         addReply(c,shared.cone);
  89     else
  90         addReply(c,shared.czero);
  91 }
  92
  93 void scardCommand(redisClient *c) {
  94     robj *o;
  95     dict *s;
  96
  97     if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
  98         checkType(c,o,REDIS_SET)) return;
  99
 100     s = o->ptr;
 101     addReplyUlong(c,dictSize(s));
 102 }
 103
 104 void spopCommand(redisClient *c) {
 105     robj *set;
 106     dictEntry *de;
 107
 108     if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
 109         checkType(c,set,REDIS_SET)) return;
 110
 111     de = dictGetRandomKey(set->ptr);
 112     if (de == NULL) {
 113         addReply(c,shared.nullbulk);
 114     } else {
 115         robj *ele = dictGetEntryKey(de);
 116
 117         addReplyBulk(c,ele);
 118         dictDelete(set->ptr,ele);
 119         if (htNeedsResize(set->ptr)) dictResize(set->ptr);
 120         if (dictSize((dict*)set->ptr) == 0) dbDelete(c->db,c->argv[1]);
 121         server.dirty++;
 122     }
 123 }
 124
 125 void srandmemberCommand(redisClient *c) {
 126     robj *set;
 127     dictEntry *de;
 128
 129     if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
 130         checkType(c,set,REDIS_SET)) return;
 131
 132     de = dictGetRandomKey(set->ptr);
 133     if (de == NULL) {
 134         addReply(c,shared.nullbulk);
 135     } else {
 136         robj *ele = dictGetEntryKey(de);
 137
 138         addReplyBulk(c,ele);
 139     }
 140 }
 141
 142 int qsortCompareSetsByCardinality(const void *s1, const void *s2) {
 143     dict **d1 = (void*) s1, **d2 = (void*) s2;
 144
 145     return dictSize(*d1)-dictSize(*d2);
 146 }
 147
 148 void sinterGenericCommand(redisClient *c, robj **setskeys, unsigned long setsnum, robj *dstkey) {
 149     dict **dv = zmalloc(sizeof(dict*)*setsnum);
 150     dictIterator *di;
 151     dictEntry *de;
 152     robj *lenobj = NULL, *dstset = NULL;
 153     unsigned long j, cardinality = 0;
 154
 155     for (j = 0; j < setsnum; j++) {
 156         robj *setobj;
 157
 158         setobj = dstkey ?
 159                     lookupKeyWrite(c->db,setskeys[j]) :
 160                     lookupKeyRead(c->db,setskeys[j]);
 161         if (!setobj) {
 162             zfree(dv);
 163             if (dstkey) {
 164                 if (dbDelete(c->db,dstkey))
 165                     server.dirty++;
 166                 addReply(c,shared.czero);
 167             } else {
 168                 addReply(c,shared.emptymultibulk);
 169             }
 170             return;
 171         }
 172         if (setobj->type != REDIS_SET) {
 173             zfree(dv);
 174             addReply(c,shared.wrongtypeerr);
 175             return;
 176         }
 177         dv[j] = setobj->ptr;
 178     }
 179     /* Sort sets from the smallest to largest, this will improve our
 180      * algorithm's performace */
 181     qsort(dv,setsnum,sizeof(dict*),qsortCompareSetsByCardinality);
 182
 183     /* The first thing we should output is the total number of elements...
 184      * since this is a multi-bulk write, but at this stage we don't know
 185      * the intersection set size, so we use a trick, append an empty object
 186      * to the output list and save the pointer to later modify it with the
 187      * right length */
 188     if (!dstkey) {
 189         lenobj = createObject(REDIS_STRING,NULL);
 190         addReply(c,lenobj);
 191         decrRefCount(lenobj);
 192     } else {
 193         /* If we have a target key where to store the resulting set
 194          * create this key with an empty set inside */
 195         dstset = createSetObject();
 196     }
 197
 198     /* Iterate all the elements of the first (smallest) set, and test
 199      * the element against all the other sets, if at least one set does
 200      * not include the element it is discarded */
 201     di = dictGetIterator(dv[0]);
 202
 203     while((de = dictNext(di)) != NULL) {
 204         robj *ele;
 205
 206         for (j = 1; j < setsnum; j++)
 207             if (dictFind(dv[j],dictGetEntryKey(de)) == NULL) break;
 208         if (j != setsnum)
 209             continue; /* at least one set does not contain the member */
 210         ele = dictGetEntryKey(de);
 211         if (!dstkey) {
 212             addReplyBulk(c,ele);
 213             cardinality++;
 214         } else {
 215             dictAdd(dstset->ptr,ele,NULL);
 216             incrRefCount(ele);
 217         }
 218     }
 219     dictReleaseIterator(di);
 220
 221     if (dstkey) {
 222         /* Store the resulting set into the target, if the intersection
 223          * is not an empty set. */
 224         dbDelete(c->db,dstkey);
 225         if (dictSize((dict*)dstset->ptr) > 0) {
 226             dbAdd(c->db,dstkey,dstset);
 227             addReplyLongLong(c,dictSize((dict*)dstset->ptr));
 228         } else {
 229             decrRefCount(dstset);
 230             addReply(c,shared.czero);
 231         }
 232         server.dirty++;
 233     } else {
 234         lenobj->ptr = sdscatprintf(sdsempty(),"*%lu\r\n",cardinality);
 235     }
 236     zfree(dv);
 237 }
 238
 239 void sinterCommand(redisClient *c) {
 240     sinterGenericCommand(c,c->argv+1,c->argc-1,NULL);
 241 }
 242
 243 void sinterstoreCommand(redisClient *c) {
 244     sinterGenericCommand(c,c->argv+2,c->argc-2,c->argv[1]);
 245 }
 246
 247 void sunionDiffGenericCommand(redisClient *c, robj **setskeys, int setsnum, robj *dstkey, int op) {
 248     dict **dv = zmalloc(sizeof(dict*)*setsnum);
 249     dictIterator *di;
 250     dictEntry *de;
 251     robj *dstset = NULL;
 252     int j, cardinality = 0;
 253
 254     for (j = 0; j < setsnum; j++) {
 255         robj *setobj;
 256
 257         setobj = dstkey ?
 258                     lookupKeyWrite(c->db,setskeys[j]) :
 259                     lookupKeyRead(c->db,setskeys[j]);
 260         if (!setobj) {
 261             dv[j] = NULL;
 262             continue;
 263         }
 264         if (setobj->type != REDIS_SET) {
 265             zfree(dv);
 266             addReply(c,shared.wrongtypeerr);
 267             return;
 268         }
 269         dv[j] = setobj->ptr;
 270     }
 271
 272     /* We need a temp set object to store our union. If the dstkey
 273      * is not NULL (that is, we are inside an SUNIONSTORE operation) then
 274      * this set object will be the resulting object to set into the target key*/
 275     dstset = createSetObject();
 276
 277     /* Iterate all the elements of all the sets, add every element a single
 278      * time to the result set */
 279     for (j = 0; j < setsnum; j++) {
 280         if (op == REDIS_OP_DIFF && j == 0 && !dv[j]) break; /* result set is empty */
 281         if (!dv[j]) continue; /* non existing keys are like empty sets */
 282
 283         di = dictGetIterator(dv[j]);
 284
 285         while((de = dictNext(di)) != NULL) {
 286             robj *ele;
 287
 288             /* dictAdd will not add the same element multiple times */
 289             ele = dictGetEntryKey(de);
 290             if (op == REDIS_OP_UNION || j == 0) {
 291                 if (dictAdd(dstset->ptr,ele,NULL) == DICT_OK) {
 292                     incrRefCount(ele);
 293                     cardinality++;
 294                 }
 295             } else if (op == REDIS_OP_DIFF) {
 296                 if (dictDelete(dstset->ptr,ele) == DICT_OK) {
 297                     cardinality--;
 298                 }
 299             }
 300         }
 301         dictReleaseIterator(di);
 302
 303         /* result set is empty? Exit asap. */
 304         if (op == REDIS_OP_DIFF && cardinality == 0) break;
 305     }
 306
 307     /* Output the content of the resulting set, if not in STORE mode */
 308     if (!dstkey) {
 309         addReplySds(c,sdscatprintf(sdsempty(),"*%d\r\n",cardinality));
 310         di = dictGetIterator(dstset->ptr);
 311         while((de = dictNext(di)) != NULL) {
 312             robj *ele;
 313
 314             ele = dictGetEntryKey(de);
 315             addReplyBulk(c,ele);
 316         }
 317         dictReleaseIterator(di);
 318         decrRefCount(dstset);
 319     } else {
 320         /* If we have a target key where to store the resulting set
 321          * create this key with the result set inside */
 322         dbDelete(c->db,dstkey);
 323         if (dictSize((dict*)dstset->ptr) > 0) {
 324             dbAdd(c->db,dstkey,dstset);
 325             addReplyLongLong(c,dictSize((dict*)dstset->ptr));
 326         } else {
 327             decrRefCount(dstset);
 328             addReply(c,shared.czero);
 329         }
 330         server.dirty++;
 331     }
 332     zfree(dv);
 333 }
 334
 335 void sunionCommand(redisClient *c) {
 336     sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_UNION);
 337 }
 338
 339 void sunionstoreCommand(redisClient *c) {
 340     sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_UNION);
 341 }
 342
 343 void sdiffCommand(redisClient *c) {
 344     sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_DIFF);
 345 }
 346
 347 void sdiffstoreCommand(redisClient *c) {
 348     sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_DIFF);
 349 }