src/t_set.c

   1 #include "redis.h"
   2
   3 /*-----------------------------------------------------------------------------
   4  * Set Commands
   5  *----------------------------------------------------------------------------*/
   6
   7 /* Factory method to return a set that *can* hold "value". When the object has
   8  * an integer-encodable value, an intset will be returned. Otherwise a regular
   9  * hash table. */
  10 robj *setTypeCreate(robj *value) {
  11     if (isObjectRepresentableAsLongLong(value,NULL) == REDIS_OK)
  12         return createIntsetObject();
  13     return createSetObject();
  14 }
  15
  16 int setTypeAdd(robj *subject, robj *value) {
  17     long long llval;
  18     if (subject->encoding == REDIS_ENCODING_HT) {
  19         if (dictAdd(subject->ptr,value,NULL) == DICT_OK) {
  20             incrRefCount(value);
  21             return 1;
  22         }
  23     } else if (subject->encoding == REDIS_ENCODING_INTSET) {
  24         if (isObjectRepresentableAsLongLong(value,&llval) == REDIS_OK) {
  25             uint8_t success = 0;
  26             subject->ptr = intsetAdd(subject->ptr,llval,&success);
  27             if (success) {
  28                 /* Convert to regular set when the intset contains
  29                  * too many entries. */
  30                 if (intsetLen(subject->ptr) > server.set_max_intset_entries)
  31                     setTypeConvert(subject,REDIS_ENCODING_HT);
  32                 return 1;
  33             }
  34         } else {
  35             /* Failed to get integer from object, convert to regular set. */
  36             setTypeConvert(subject,REDIS_ENCODING_HT);
  37
  38             /* The set *was* an intset and this value is not integer
  39              * encodable, so dictAdd should always work. */
  40             redisAssert(dictAdd(subject->ptr,value,NULL) == DICT_OK);
  41             incrRefCount(value);
  42             return 1;
  43         }
  44     } else {
  45         redisPanic("Unknown set encoding");
  46     }
  47     return 0;
  48 }
  49
  50 int setTypeRemove(robj *subject, robj *value) {
  51     long long llval;
  52     if (subject->encoding == REDIS_ENCODING_HT) {
  53         if (dictDelete(subject->ptr,value) == DICT_OK) {
  54             if (htNeedsResize(subject->ptr)) dictResize(subject->ptr);
  55             return 1;
  56         }
  57     } else if (subject->encoding == REDIS_ENCODING_INTSET) {
  58         if (isObjectRepresentableAsLongLong(value,&llval) == REDIS_OK) {
  59             uint8_t success;
  60             subject->ptr = intsetRemove(subject->ptr,llval,&success);
  61             if (success) return 1;
  62         }
  63     } else {
  64         redisPanic("Unknown set encoding");
  65     }
  66     return 0;
  67 }
  68
  69 int setTypeIsMember(robj *subject, robj *value) {
  70     long long llval;
  71     if (subject->encoding == REDIS_ENCODING_HT) {
  72         return dictFind((dict*)subject->ptr,value) != NULL;
  73     } else if (subject->encoding == REDIS_ENCODING_INTSET) {
  74         if (isObjectRepresentableAsLongLong(value,&llval) == REDIS_OK) {
  75             return intsetFind((intset*)subject->ptr,llval);
  76         }
  77     } else {
  78         redisPanic("Unknown set encoding");
  79     }
  80     return 0;
  81 }
  82
  83 setTypeIterator *setTypeInitIterator(robj *subject) {
  84     setTypeIterator *si = zmalloc(sizeof(setTypeIterator));
  85     si->subject = subject;
  86     si->encoding = subject->encoding;
  87     if (si->encoding == REDIS_ENCODING_HT) {
  88         si->di = dictGetIterator(subject->ptr);
  89     } else if (si->encoding == REDIS_ENCODING_INTSET) {
  90         si->ii = 0;
  91     } else {
  92         redisPanic("Unknown set encoding");
  93     }
  94     return si;
  95 }
  96
  97 void setTypeReleaseIterator(setTypeIterator *si) {
  98     if (si->encoding == REDIS_ENCODING_HT)
  99         dictReleaseIterator(si->di);
 100     zfree(si);
 101 }
 102
 103 /* Move to the next entry in the set. Returns the object at the current
 104  * position, or NULL when the end is reached. This object will have its
 105  * refcount incremented, so the caller needs to take care of this. */
 106 robj *setTypeNext(setTypeIterator *si) {
 107     robj *ret = NULL;
 108     if (si->encoding == REDIS_ENCODING_HT) {
 109         dictEntry *de = dictNext(si->di);
 110         if (de != NULL) {
 111             ret = dictGetEntryKey(de);
 112             incrRefCount(ret);
 113         }
 114     } else if (si->encoding == REDIS_ENCODING_INTSET) {
 115         int64_t llval;
 116         if (intsetGet(si->subject->ptr,si->ii++,&llval))
 117             ret = createStringObjectFromLongLong(llval);
 118     }
 119     return ret;
 120 }
 121
 122
 123 /* Return random element from set. The returned object will always have
 124  * an incremented refcount. */
 125 robj *setTypeRandomElement(robj *subject) {
 126     robj *ret = NULL;
 127     if (subject->encoding == REDIS_ENCODING_HT) {
 128         dictEntry *de = dictGetRandomKey(subject->ptr);
 129         ret = dictGetEntryKey(de);
 130         incrRefCount(ret);
 131     } else if (subject->encoding == REDIS_ENCODING_INTSET) {
 132         long long llval = intsetRandom(subject->ptr);
 133         ret = createStringObjectFromLongLong(llval);
 134     } else {
 135         redisPanic("Unknown set encoding");
 136     }
 137     return ret;
 138 }
 139
 140 unsigned long setTypeSize(robj *subject) {
 141     if (subject->encoding == REDIS_ENCODING_HT) {
 142         return dictSize((dict*)subject->ptr);
 143     } else if (subject->encoding == REDIS_ENCODING_INTSET) {
 144         return intsetLen((intset*)subject->ptr);
 145     } else {
 146         redisPanic("Unknown set encoding");
 147     }
 148 }
 149
 150 /* Convert the set to specified encoding. The resulting dict (when converting
 151  * to a hashtable) is presized to hold the number of elements in the original
 152  * set. */
 153 void setTypeConvert(robj *subject, int enc) {
 154     setTypeIterator *si;
 155     robj *element;
 156     redisAssert(subject->type == REDIS_SET);
 157
 158     if (enc == REDIS_ENCODING_HT) {
 159         dict *d = dictCreate(&setDictType,NULL);
 160         /* Presize the dict to avoid rehashing */
 161         dictExpand(d,intsetLen(subject->ptr));
 162
 163         /* setTypeGet returns a robj with incremented refcount */
 164         si = setTypeInitIterator(subject);
 165         while ((element = setTypeNext(si)) != NULL)
 166             redisAssert(dictAdd(d,element,NULL) == DICT_OK);
 167         setTypeReleaseIterator(si);
 168
 169         subject->encoding = REDIS_ENCODING_HT;
 170         zfree(subject->ptr);
 171         subject->ptr = d;
 172     } else {
 173         redisPanic("Unsupported set conversion");
 174     }
 175 }
 176
 177 void saddCommand(redisClient *c) {
 178     robj *set;
 179
 180     set = lookupKeyWrite(c->db,c->argv[1]);
 181     c->argv[2] = tryObjectEncoding(c->argv[2]);
 182     if (set == NULL) {
 183         set = setTypeCreate(c->argv[2]);
 184         dbAdd(c->db,c->argv[1],set);
 185     } else {
 186         if (set->type != REDIS_SET) {
 187             addReply(c,shared.wrongtypeerr);
 188             return;
 189         }
 190     }
 191     if (setTypeAdd(set,c->argv[2])) {
 192         touchWatchedKey(c->db,c->argv[1]);
 193         server.dirty++;
 194         addReply(c,shared.cone);
 195     } else {
 196         addReply(c,shared.czero);
 197     }
 198 }
 199
 200 void sremCommand(redisClient *c) {
 201     robj *set;
 202
 203     if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
 204         checkType(c,set,REDIS_SET)) return;
 205
 206     c->argv[2] = tryObjectEncoding(c->argv[2]);
 207     if (setTypeRemove(set,c->argv[2])) {
 208         if (setTypeSize(set) == 0) dbDelete(c->db,c->argv[1]);
 209         touchWatchedKey(c->db,c->argv[1]);
 210         server.dirty++;
 211         addReply(c,shared.cone);
 212     } else {
 213         addReply(c,shared.czero);
 214     }
 215 }
 216
 217 void smoveCommand(redisClient *c) {
 218     robj *srcset, *dstset, *ele;
 219     srcset = lookupKeyWrite(c->db,c->argv[1]);
 220     dstset = lookupKeyWrite(c->db,c->argv[2]);
 221     ele = c->argv[3] = tryObjectEncoding(c->argv[3]);
 222
 223     /* If the source key does not exist return 0 */
 224     if (srcset == NULL) {
 225         addReply(c,shared.czero);
 226         return;
 227     }
 228
 229     /* If the source key has the wrong type, or the destination key
 230      * is set and has the wrong type, return with an error. */
 231     if (checkType(c,srcset,REDIS_SET) ||
 232         (dstset && checkType(c,dstset,REDIS_SET))) return;
 233
 234     /* If srcset and dstset are equal, SMOVE is a no-op */
 235     if (srcset == dstset) {
 236         addReply(c,shared.cone);
 237         return;
 238     }
 239
 240     /* If the element cannot be removed from the src set, return 0. */
 241     if (!setTypeRemove(srcset,ele)) {
 242         addReply(c,shared.czero);
 243         return;
 244     }
 245
 246     /* Remove the src set from the database when empty */
 247     if (setTypeSize(srcset) == 0) dbDelete(c->db,c->argv[1]);
 248     touchWatchedKey(c->db,c->argv[1]);
 249     touchWatchedKey(c->db,c->argv[2]);
 250     server.dirty++;
 251
 252     /* Create the destination set when it doesn't exist */
 253     if (!dstset) {
 254         dstset = setTypeCreate(ele);
 255         dbAdd(c->db,c->argv[2],dstset);
 256     }
 257
 258     /* An extra key has changed when ele was successfully added to dstset */
 259     if (setTypeAdd(dstset,ele)) server.dirty++;
 260     addReply(c,shared.cone);
 261 }
 262
 263 void sismemberCommand(redisClient *c) {
 264     robj *set;
 265
 266     if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
 267         checkType(c,set,REDIS_SET)) return;
 268
 269     c->argv[2] = tryObjectEncoding(c->argv[2]);
 270     if (setTypeIsMember(set,c->argv[2]))
 271         addReply(c,shared.cone);
 272     else
 273         addReply(c,shared.czero);
 274 }
 275
 276 void scardCommand(redisClient *c) {
 277     robj *o;
 278
 279     if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
 280         checkType(c,o,REDIS_SET)) return;
 281
 282     addReplyLongLong(c,setTypeSize(o));
 283 }
 284
 285 void spopCommand(redisClient *c) {
 286     robj *set, *ele;
 287
 288     if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
 289         checkType(c,set,REDIS_SET)) return;
 290
 291     ele = setTypeRandomElement(set);
 292     if (ele == NULL) {
 293         addReply(c,shared.nullbulk);
 294     } else {
 295         setTypeRemove(set,ele);
 296         addReplyBulk(c,ele);
 297         decrRefCount(ele);
 298         if (setTypeSize(set) == 0) dbDelete(c->db,c->argv[1]);
 299         touchWatchedKey(c->db,c->argv[1]);
 300         server.dirty++;
 301     }
 302 }
 303
 304 void srandmemberCommand(redisClient *c) {
 305     robj *set, *ele;
 306
 307     if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
 308         checkType(c,set,REDIS_SET)) return;
 309
 310     ele = setTypeRandomElement(set);
 311     if (ele == NULL) {
 312         addReply(c,shared.nullbulk);
 313     } else {
 314         addReplyBulk(c,ele);
 315         decrRefCount(ele);
 316     }
 317 }
 318
 319 int qsortCompareSetsByCardinality(const void *s1, const void *s2) {
 320     return setTypeSize(*(robj**)s1)-setTypeSize(*(robj**)s2);
 321 }
 322
 323 void sinterGenericCommand(redisClient *c, robj **setkeys, unsigned long setnum, robj *dstkey) {
 324     robj **sets = zmalloc(sizeof(robj*)*setnum);
 325     setTypeIterator *si;
 326     robj *ele, *dstset = NULL;
 327     void *replylen = NULL;
 328     unsigned long j, cardinality = 0;
 329
 330     for (j = 0; j < setnum; j++) {
 331         robj *setobj = dstkey ?
 332             lookupKeyWrite(c->db,setkeys[j]) :
 333             lookupKeyRead(c->db,setkeys[j]);
 334         if (!setobj) {
 335             zfree(sets);
 336             if (dstkey) {
 337                 if (dbDelete(c->db,dstkey)) {
 338                     touchWatchedKey(c->db,dstkey);
 339                     server.dirty++;
 340                 }
 341                 addReply(c,shared.czero);
 342             } else {
 343                 addReply(c,shared.emptymultibulk);
 344             }
 345             return;
 346         }
 347         if (checkType(c,setobj,REDIS_SET)) {
 348             zfree(sets);
 349             return;
 350         }
 351         sets[j] = setobj;
 352     }
 353     /* Sort sets from the smallest to largest, this will improve our
 354      * algorithm's performace */
 355     qsort(sets,setnum,sizeof(robj*),qsortCompareSetsByCardinality);
 356
 357     /* The first thing we should output is the total number of elements...
 358      * since this is a multi-bulk write, but at this stage we don't know
 359      * the intersection set size, so we use a trick, append an empty object
 360      * to the output list and save the pointer to later modify it with the
 361      * right length */
 362     if (!dstkey) {
 363         replylen = addDeferredMultiBulkLength(c);
 364     } else {
 365         /* If we have a target key where to store the resulting set
 366          * create this key with an empty set inside */
 367         dstset = createIntsetObject();
 368     }
 369
 370     /* Iterate all the elements of the first (smallest) set, and test
 371      * the element against all the other sets, if at least one set does
 372      * not include the element it is discarded */
 373     si = setTypeInitIterator(sets[0]);
 374     while((ele = setTypeNext(si)) != NULL) {
 375         for (j = 1; j < setnum; j++)
 376             if (!setTypeIsMember(sets[j],ele)) break;
 377
 378         /* Only take action when all sets contain the member */
 379         if (j == setnum) {
 380             if (!dstkey) {
 381                 addReplyBulk(c,ele);
 382                 cardinality++;
 383             } else {
 384                 setTypeAdd(dstset,ele);
 385             }
 386         }
 387         decrRefCount(ele);
 388     }
 389     setTypeReleaseIterator(si);
 390
 391     if (dstkey) {
 392         /* Store the resulting set into the target, if the intersection
 393          * is not an empty set. */
 394         dbDelete(c->db,dstkey);
 395         if (setTypeSize(dstset) > 0) {
 396             dbAdd(c->db,dstkey,dstset);
 397             addReplyLongLong(c,setTypeSize(dstset));
 398         } else {
 399             decrRefCount(dstset);
 400             addReply(c,shared.czero);
 401         }
 402         touchWatchedKey(c->db,dstkey);
 403         server.dirty++;
 404     } else {
 405         setDeferredMultiBulkLength(c,replylen,cardinality);
 406     }
 407     zfree(sets);
 408 }
 409
 410 void sinterCommand(redisClient *c) {
 411     sinterGenericCommand(c,c->argv+1,c->argc-1,NULL);
 412 }
 413
 414 void sinterstoreCommand(redisClient *c) {
 415     sinterGenericCommand(c,c->argv+2,c->argc-2,c->argv[1]);
 416 }
 417
 418 #define REDIS_OP_UNION 0
 419 #define REDIS_OP_DIFF 1
 420 #define REDIS_OP_INTER 2
 421
 422 void sunionDiffGenericCommand(redisClient *c, robj **setkeys, int setnum, robj *dstkey, int op) {
 423     robj **sets = zmalloc(sizeof(robj*)*setnum);
 424     setTypeIterator *si;
 425     robj *ele, *dstset = NULL;
 426     int j, cardinality = 0;
 427
 428     for (j = 0; j < setnum; j++) {
 429         robj *setobj = dstkey ?
 430             lookupKeyWrite(c->db,setkeys[j]) :
 431             lookupKeyRead(c->db,setkeys[j]);
 432         if (!setobj) {
 433             sets[j] = NULL;
 434             continue;
 435         }
 436         if (checkType(c,setobj,REDIS_SET)) {
 437             zfree(sets);
 438             return;
 439         }
 440         sets[j] = setobj;
 441     }
 442
 443     /* We need a temp set object to store our union. If the dstkey
 444      * is not NULL (that is, we are inside an SUNIONSTORE operation) then
 445      * this set object will be the resulting object to set into the target key*/
 446     dstset = createIntsetObject();
 447
 448     /* Iterate all the elements of all the sets, add every element a single
 449      * time to the result set */
 450     for (j = 0; j < setnum; j++) {
 451         if (op == REDIS_OP_DIFF && j == 0 && !sets[j]) break; /* result set is empty */
 452         if (!sets[j]) continue; /* non existing keys are like empty sets */
 453
 454         si = setTypeInitIterator(sets[j]);
 455         while((ele = setTypeNext(si)) != NULL) {
 456             if (op == REDIS_OP_UNION || j == 0) {
 457                 if (setTypeAdd(dstset,ele)) {
 458                     cardinality++;
 459                 }
 460             } else if (op == REDIS_OP_DIFF) {
 461                 if (setTypeRemove(dstset,ele)) {
 462                     cardinality--;
 463                 }
 464             }
 465             decrRefCount(ele);
 466         }
 467         setTypeReleaseIterator(si);
 468
 469         /* Exit when result set is empty. */
 470         if (op == REDIS_OP_DIFF && cardinality == 0) break;
 471     }
 472
 473     /* Output the content of the resulting set, if not in STORE mode */
 474     if (!dstkey) {
 475         addReplyMultiBulkLen(c,cardinality);
 476         si = setTypeInitIterator(dstset);
 477         while((ele = setTypeNext(si)) != NULL) {
 478             addReplyBulk(c,ele);
 479             decrRefCount(ele);
 480         }
 481         setTypeReleaseIterator(si);
 482         decrRefCount(dstset);
 483     } else {
 484         /* If we have a target key where to store the resulting set
 485          * create this key with the result set inside */
 486         dbDelete(c->db,dstkey);
 487         if (setTypeSize(dstset) > 0) {
 488             dbAdd(c->db,dstkey,dstset);
 489             addReplyLongLong(c,setTypeSize(dstset));
 490         } else {
 491             decrRefCount(dstset);
 492             addReply(c,shared.czero);
 493         }
 494         touchWatchedKey(c->db,dstkey);
 495         server.dirty++;
 496     }
 497     zfree(sets);
 498 }
 499
 500 void sunionCommand(redisClient *c) {
 501     sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_UNION);
 502 }
 503
 504 void sunionstoreCommand(redisClient *c) {
 505     sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_UNION);
 506 }
 507
 508 void sdiffCommand(redisClient *c) {
 509     sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_DIFF);
 510 }
 511
 512 void sdiffstoreCommand(redisClient *c) {
 513     sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_DIFF);
 514 }