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