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 *setobj, robj *value) {
  51     long long llval;
  52     if (setobj->encoding == REDIS_ENCODING_HT) {
  53         if (dictDelete(setobj->ptr,value) == DICT_OK) {
  54             if (htNeedsResize(setobj->ptr)) dictResize(setobj->ptr);
  55             return 1;
  56         }
  57     } else if (setobj->encoding == REDIS_ENCODING_INTSET) {
  58         if (isObjectRepresentableAsLongLong(value,&llval) == REDIS_OK) {
  59             int success;
  60             setobj->ptr = intsetRemove(setobj->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.
 105  *
 106  * Since set elements can be internally be stored as redis objects or
 107  * simple arrays of integers, setTypeNext returns the encoding of the
 108  * set object you are iterating, and will populate the appropriate pointer
 109  * (eobj) or (llobj) accordingly.
 110  *
 111  * When there are no longer elements -1 is returned.
 112  * Returned objects ref count is not incremented, so this function is
 113  * copy on write friendly. */
 114 int setTypeNext(setTypeIterator *si, robj **objele, int64_t *llele) {
 115     if (si->encoding == REDIS_ENCODING_HT) {
 116         dictEntry *de = dictNext(si->di);
 117         if (de == NULL) return -1;
 118         *objele = dictGetEntryKey(de);
 119     } else if (si->encoding == REDIS_ENCODING_INTSET) {
 120         if (!intsetGet(si->subject->ptr,si->ii++,llele))
 121             return -1;
 122     }
 123     return si->encoding;
 124 }
 125
 126 /* The not copy on write friendly version but easy to use version
 127  * of setTypeNext() is setTypeNextObject(), returning new objects
 128  * or incrementing the ref count of returned objects. So if you don't
 129  * retain a pointer to this object you should call decrRefCount() against it.
 130  *
 131  * This function is the way to go for write operations where COW is not
 132  * an issue as the result will be anyway of incrementing the ref count. */
 133 robj *setTypeNextObject(setTypeIterator *si) {
 134     int64_t intele;
 135     robj *objele;
 136     int encoding;
 137
 138     encoding = setTypeNext(si,&objele,&intele);
 139     switch(encoding) {
 140         case -1:    return NULL;
 141         case REDIS_ENCODING_INTSET:
 142             return createStringObjectFromLongLong(intele);
 143         case REDIS_ENCODING_HT:
 144             incrRefCount(objele);
 145             return objele;
 146         default:
 147             redisPanic("Unsupported encoding");
 148     }
 149     return NULL; /* just to suppress warnings */
 150 }
 151
 152 /* Return random element from a non empty set.
 153  * The returned element can be a int64_t value if the set is encoded
 154  * as an "intset" blob of integers, or a redis object if the set
 155  * is a regular set.
 156  *
 157  * The caller provides both pointers to be populated with the right
 158  * object. The return value of the function is the object->encoding
 159  * field of the object and is used by the caller to check if the
 160  * int64_t pointer or the redis object pointere was populated.
 161  *
 162  * When an object is returned (the set was a real set) the ref count
 163  * of the object is not incremented so this function can be considered
 164  * copy on write friendly. */
 165 int setTypeRandomElement(robj *setobj, robj **objele, int64_t *llele) {
 166     if (setobj->encoding == REDIS_ENCODING_HT) {
 167         dictEntry *de = dictGetRandomKey(setobj->ptr);
 168         *objele = dictGetEntryKey(de);
 169     } else if (setobj->encoding == REDIS_ENCODING_INTSET) {
 170         *llele = intsetRandom(setobj->ptr);
 171     } else {
 172         redisPanic("Unknown set encoding");
 173     }
 174     return setobj->encoding;
 175 }
 176
 177 unsigned long setTypeSize(robj *subject) {
 178     if (subject->encoding == REDIS_ENCODING_HT) {
 179         return dictSize((dict*)subject->ptr);
 180     } else if (subject->encoding == REDIS_ENCODING_INTSET) {
 181         return intsetLen((intset*)subject->ptr);
 182     } else {
 183         redisPanic("Unknown set encoding");
 184     }
 185 }
 186
 187 /* Convert the set to specified encoding. The resulting dict (when converting
 188  * to a hashtable) is presized to hold the number of elements in the original
 189  * set. */
 190 void setTypeConvert(robj *setobj, int enc) {
 191     setTypeIterator *si;
 192     redisAssert(setobj->type == REDIS_SET &&
 193                 setobj->encoding == REDIS_ENCODING_INTSET);
 194
 195     if (enc == REDIS_ENCODING_HT) {
 196         int64_t intele;
 197         dict *d = dictCreate(&setDictType,NULL);
 198         robj *element;
 199
 200         /* Presize the dict to avoid rehashing */
 201         dictExpand(d,intsetLen(setobj->ptr));
 202
 203         /* To add the elements we extract integers and create redis objects */
 204         si = setTypeInitIterator(setobj);
 205         while (setTypeNext(si,NULL,&intele) != -1) {
 206             element = createStringObjectFromLongLong(intele);
 207             redisAssert(dictAdd(d,element,NULL) == DICT_OK);
 208         }
 209         setTypeReleaseIterator(si);
 210
 211         setobj->encoding = REDIS_ENCODING_HT;
 212         zfree(setobj->ptr);
 213         setobj->ptr = d;
 214     } else {
 215         redisPanic("Unsupported set conversion");
 216     }
 217 }
 218
 219 void saddCommand(redisClient *c) {
 220     robj *set;
 221     int j, added = 0;
 222
 223     set = lookupKeyWrite(c->db,c->argv[1]);
 224     if (set == NULL) {
 225         set = setTypeCreate(c->argv[2]);
 226         dbAdd(c->db,c->argv[1],set);
 227     } else {
 228         if (set->type != REDIS_SET) {
 229             addReply(c,shared.wrongtypeerr);
 230             return;
 231         }
 232     }
 233
 234     for (j = 2; j < c->argc; j++) {
 235         c->argv[j] = tryObjectEncoding(c->argv[j]);
 236         if (setTypeAdd(set,c->argv[j])) added++;
 237     }
 238     if (added) signalModifiedKey(c->db,c->argv[1]);
 239     server.dirty += added;
 240     addReplyLongLong(c,added);
 241 }
 242
 243 void sremCommand(redisClient *c) {
 244     robj *set;
 245     int j, deleted = 0;
 246
 247     if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
 248         checkType(c,set,REDIS_SET)) return;
 249
 250     for (j = 2; j < c->argc; j++) {
 251         if (setTypeRemove(set,c->argv[j])) {
 252             deleted++;
 253             if (setTypeSize(set) == 0) {
 254                 dbDelete(c->db,c->argv[1]);
 255                 break;
 256             }
 257         }
 258     }
 259     if (deleted) {
 260         signalModifiedKey(c->db,c->argv[1]);
 261         server.dirty += deleted;
 262     }
 263     addReplyLongLong(c,deleted);
 264 }
 265
 266 void smoveCommand(redisClient *c) {
 267     robj *srcset, *dstset, *ele;
 268     srcset = lookupKeyWrite(c->db,c->argv[1]);
 269     dstset = lookupKeyWrite(c->db,c->argv[2]);
 270     ele = c->argv[3] = tryObjectEncoding(c->argv[3]);
 271
 272     /* If the source key does not exist return 0 */
 273     if (srcset == NULL) {
 274         addReply(c,shared.czero);
 275         return;
 276     }
 277
 278     /* If the source key has the wrong type, or the destination key
 279      * is set and has the wrong type, return with an error. */
 280     if (checkType(c,srcset,REDIS_SET) ||
 281         (dstset && checkType(c,dstset,REDIS_SET))) return;
 282
 283     /* If srcset and dstset are equal, SMOVE is a no-op */
 284     if (srcset == dstset) {
 285         addReply(c,shared.cone);
 286         return;
 287     }
 288
 289     /* If the element cannot be removed from the src set, return 0. */
 290     if (!setTypeRemove(srcset,ele)) {
 291         addReply(c,shared.czero);
 292         return;
 293     }
 294
 295     /* Remove the src set from the database when empty */
 296     if (setTypeSize(srcset) == 0) dbDelete(c->db,c->argv[1]);
 297     signalModifiedKey(c->db,c->argv[1]);
 298     signalModifiedKey(c->db,c->argv[2]);
 299     server.dirty++;
 300
 301     /* Create the destination set when it doesn't exist */
 302     if (!dstset) {
 303         dstset = setTypeCreate(ele);
 304         dbAdd(c->db,c->argv[2],dstset);
 305     }
 306
 307     /* An extra key has changed when ele was successfully added to dstset */
 308     if (setTypeAdd(dstset,ele)) server.dirty++;
 309     addReply(c,shared.cone);
 310 }
 311
 312 void sismemberCommand(redisClient *c) {
 313     robj *set;
 314
 315     if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
 316         checkType(c,set,REDIS_SET)) return;
 317
 318     c->argv[2] = tryObjectEncoding(c->argv[2]);
 319     if (setTypeIsMember(set,c->argv[2]))
 320         addReply(c,shared.cone);
 321     else
 322         addReply(c,shared.czero);
 323 }
 324
 325 void scardCommand(redisClient *c) {
 326     robj *o;
 327
 328     if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
 329         checkType(c,o,REDIS_SET)) return;
 330
 331     addReplyLongLong(c,setTypeSize(o));
 332 }
 333
 334 void spopCommand(redisClient *c) {
 335     robj *set, *ele;
 336     int64_t llele;
 337     int encoding;
 338
 339     if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
 340         checkType(c,set,REDIS_SET)) return;
 341
 342     encoding = setTypeRandomElement(set,&ele,&llele);
 343     if (encoding == REDIS_ENCODING_INTSET) {
 344         ele = createStringObjectFromLongLong(llele);
 345         set->ptr = intsetRemove(set->ptr,llele,NULL);
 346     } else {
 347         incrRefCount(ele);
 348         setTypeRemove(set,ele);
 349     }
 350
 351     /* Change argv to replicate as SREM */
 352     c->argc = 3;
 353     c->argv = zrealloc(c->argv,sizeof(robj*)*(c->argc));
 354
 355     /* Overwrite SREM with SPOP (same length) */
 356     redisAssert(sdslen(c->argv[0]->ptr) == 4);
 357     memcpy(c->argv[0]->ptr, "SREM", 4);
 358
 359     /* Popped element already has incremented refcount */
 360     c->argv[2] = ele;
 361
 362     addReplyBulk(c,ele);
 363     if (setTypeSize(set) == 0) dbDelete(c->db,c->argv[1]);
 364     signalModifiedKey(c->db,c->argv[1]);
 365     server.dirty++;
 366 }
 367
 368 void srandmemberCommand(redisClient *c) {
 369     robj *set, *ele;
 370     int64_t llele;
 371     int encoding;
 372
 373     if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
 374         checkType(c,set,REDIS_SET)) return;
 375
 376     encoding = setTypeRandomElement(set,&ele,&llele);
 377     if (encoding == REDIS_ENCODING_INTSET) {
 378         addReplyBulkLongLong(c,llele);
 379     } else {
 380         addReplyBulk(c,ele);
 381     }
 382 }
 383
 384 int qsortCompareSetsByCardinality(const void *s1, const void *s2) {
 385     return setTypeSize(*(robj**)s1)-setTypeSize(*(robj**)s2);
 386 }
 387
 388 void sinterGenericCommand(redisClient *c, robj **setkeys, unsigned long setnum, robj *dstkey) {
 389     robj **sets = zmalloc(sizeof(robj*)*setnum);
 390     setTypeIterator *si;
 391     robj *eleobj, *dstset = NULL;
 392     int64_t intobj;
 393     void *replylen = NULL;
 394     unsigned long j, cardinality = 0;
 395     int encoding;
 396
 397     for (j = 0; j < setnum; j++) {
 398         robj *setobj = dstkey ?
 399             lookupKeyWrite(c->db,setkeys[j]) :
 400             lookupKeyRead(c->db,setkeys[j]);
 401         if (!setobj) {
 402             zfree(sets);
 403             if (dstkey) {
 404                 if (dbDelete(c->db,dstkey)) {
 405                     signalModifiedKey(c->db,dstkey);
 406                     server.dirty++;
 407                 }
 408                 addReply(c,shared.czero);
 409             } else {
 410                 addReply(c,shared.emptymultibulk);
 411             }
 412             return;
 413         }
 414         if (checkType(c,setobj,REDIS_SET)) {
 415             zfree(sets);
 416             return;
 417         }
 418         sets[j] = setobj;
 419     }
 420     /* Sort sets from the smallest to largest, this will improve our
 421      * algorithm's performace */
 422     qsort(sets,setnum,sizeof(robj*),qsortCompareSetsByCardinality);
 423
 424     /* The first thing we should output is the total number of elements...
 425      * since this is a multi-bulk write, but at this stage we don't know
 426      * the intersection set size, so we use a trick, append an empty object
 427      * to the output list and save the pointer to later modify it with the
 428      * right length */
 429     if (!dstkey) {
 430         replylen = addDeferredMultiBulkLength(c);
 431     } else {
 432         /* If we have a target key where to store the resulting set
 433          * create this key with an empty set inside */
 434         dstset = createIntsetObject();
 435     }
 436
 437     /* Iterate all the elements of the first (smallest) set, and test
 438      * the element against all the other sets, if at least one set does
 439      * not include the element it is discarded */
 440     si = setTypeInitIterator(sets[0]);
 441     while((encoding = setTypeNext(si,&eleobj,&intobj)) != -1) {
 442         for (j = 1; j < setnum; j++) {
 443             if (sets[j] == sets[0]) continue;
 444             if (encoding == REDIS_ENCODING_INTSET) {
 445                 /* intset with intset is simple... and fast */
 446                 if (sets[j]->encoding == REDIS_ENCODING_INTSET &&
 447                     !intsetFind((intset*)sets[j]->ptr,intobj))
 448                 {
 449                     break;
 450                 /* in order to compare an integer with an object we
 451                  * have to use the generic function, creating an object
 452                  * for this */
 453                 } else if (sets[j]->encoding == REDIS_ENCODING_HT) {
 454                     eleobj = createStringObjectFromLongLong(intobj);
 455                     if (!setTypeIsMember(sets[j],eleobj)) {
 456                         decrRefCount(eleobj);
 457                         break;
 458                     }
 459                     decrRefCount(eleobj);
 460                 }
 461             } else if (encoding == REDIS_ENCODING_HT) {
 462                 /* Optimization... if the source object is integer
 463                  * encoded AND the target set is an intset, we can get
 464                  * a much faster path. */
 465                 if (eleobj->encoding == REDIS_ENCODING_INT &&
 466                     sets[j]->encoding == REDIS_ENCODING_INTSET &&
 467                     !intsetFind((intset*)sets[j]->ptr,(long)eleobj->ptr))
 468                 {
 469                     break;
 470                 /* else... object to object check is easy as we use the
 471                  * type agnostic API here. */
 472                 } else if (!setTypeIsMember(sets[j],eleobj)) {
 473                     break;
 474                 }
 475             }
 476         }
 477
 478         /* Only take action when all sets contain the member */
 479         if (j == setnum) {
 480             if (!dstkey) {
 481                 if (encoding == REDIS_ENCODING_HT)
 482                     addReplyBulk(c,eleobj);
 483                 else
 484                     addReplyBulkLongLong(c,intobj);
 485                 cardinality++;
 486             } else {
 487                 if (encoding == REDIS_ENCODING_INTSET) {
 488                     eleobj = createStringObjectFromLongLong(intobj);
 489                     setTypeAdd(dstset,eleobj);
 490                     decrRefCount(eleobj);
 491                 } else {
 492                     setTypeAdd(dstset,eleobj);
 493                 }
 494             }
 495         }
 496     }
 497     setTypeReleaseIterator(si);
 498
 499     if (dstkey) {
 500         /* Store the resulting set into the target, if the intersection
 501          * is not an empty set. */
 502         dbDelete(c->db,dstkey);
 503         if (setTypeSize(dstset) > 0) {
 504             dbAdd(c->db,dstkey,dstset);
 505             addReplyLongLong(c,setTypeSize(dstset));
 506         } else {
 507             decrRefCount(dstset);
 508             addReply(c,shared.czero);
 509         }
 510         signalModifiedKey(c->db,dstkey);
 511         server.dirty++;
 512     } else {
 513         setDeferredMultiBulkLength(c,replylen,cardinality);
 514     }
 515     zfree(sets);
 516 }
 517
 518 void sinterCommand(redisClient *c) {
 519     sinterGenericCommand(c,c->argv+1,c->argc-1,NULL);
 520 }
 521
 522 void sinterstoreCommand(redisClient *c) {
 523     sinterGenericCommand(c,c->argv+2,c->argc-2,c->argv[1]);
 524 }
 525
 526 #define REDIS_OP_UNION 0
 527 #define REDIS_OP_DIFF 1
 528 #define REDIS_OP_INTER 2
 529
 530 void sunionDiffGenericCommand(redisClient *c, robj **setkeys, int setnum, robj *dstkey, int op) {
 531     robj **sets = zmalloc(sizeof(robj*)*setnum);
 532     setTypeIterator *si;
 533     robj *ele, *dstset = NULL;
 534     int j, cardinality = 0;
 535
 536     for (j = 0; j < setnum; j++) {
 537         robj *setobj = dstkey ?
 538             lookupKeyWrite(c->db,setkeys[j]) :
 539             lookupKeyRead(c->db,setkeys[j]);
 540         if (!setobj) {
 541             sets[j] = NULL;
 542             continue;
 543         }
 544         if (checkType(c,setobj,REDIS_SET)) {
 545             zfree(sets);
 546             return;
 547         }
 548         sets[j] = setobj;
 549     }
 550
 551     /* We need a temp set object to store our union. If the dstkey
 552      * is not NULL (that is, we are inside an SUNIONSTORE operation) then
 553      * this set object will be the resulting object to set into the target key*/
 554     dstset = createIntsetObject();
 555
 556     /* Iterate all the elements of all the sets, add every element a single
 557      * time to the result set */
 558     for (j = 0; j < setnum; j++) {
 559         if (op == REDIS_OP_DIFF && j == 0 && !sets[j]) break; /* result set is empty */
 560         if (!sets[j]) continue; /* non existing keys are like empty sets */
 561
 562         si = setTypeInitIterator(sets[j]);
 563         while((ele = setTypeNextObject(si)) != NULL) {
 564             if (op == REDIS_OP_UNION || j == 0) {
 565                 if (setTypeAdd(dstset,ele)) {
 566                     cardinality++;
 567                 }
 568             } else if (op == REDIS_OP_DIFF) {
 569                 if (setTypeRemove(dstset,ele)) {
 570                     cardinality--;
 571                 }
 572             }
 573             decrRefCount(ele);
 574         }
 575         setTypeReleaseIterator(si);
 576
 577         /* Exit when result set is empty. */
 578         if (op == REDIS_OP_DIFF && cardinality == 0) break;
 579     }
 580
 581     /* Output the content of the resulting set, if not in STORE mode */
 582     if (!dstkey) {
 583         addReplyMultiBulkLen(c,cardinality);
 584         si = setTypeInitIterator(dstset);
 585         while((ele = setTypeNextObject(si)) != NULL) {
 586             addReplyBulk(c,ele);
 587             decrRefCount(ele);
 588         }
 589         setTypeReleaseIterator(si);
 590         decrRefCount(dstset);
 591     } else {
 592         /* If we have a target key where to store the resulting set
 593          * create this key with the result set inside */
 594         dbDelete(c->db,dstkey);
 595         if (setTypeSize(dstset) > 0) {
 596             dbAdd(c->db,dstkey,dstset);
 597             addReplyLongLong(c,setTypeSize(dstset));
 598         } else {
 599             decrRefCount(dstset);
 600             addReply(c,shared.czero);
 601         }
 602         signalModifiedKey(c->db,dstkey);
 603         server.dirty++;
 604     }
 605     zfree(sets);
 606 }
 607
 608 void sunionCommand(redisClient *c) {
 609     sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_UNION);
 610 }
 611
 612 void sunionstoreCommand(redisClient *c) {
 613     sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_UNION);
 614 }
 615
 616 void sdiffCommand(redisClient *c) {
 617     sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_DIFF);
 618 }
 619
 620 void sdiffstoreCommand(redisClient *c) {
 621     sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_DIFF);
 622 }