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, 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 a non empty set.
 124  * The returned element can be a long long value if the set is encoded
 125  * as an "intset" blob of integers, or a redis object if the set
 126  * is a regular set.
 127  *
 128  * The caller provides both pointers to be populated with the right
 129  * object. The return value of the function is the object->encoding
 130  * field of the object and is used by the caller to check if the
 131  * long long pointer or the redis object pointere was populated.
 132  *
 133  * When an object is returned (the set was a real set) the ref count
 134  * of the object is not incremented so this function can be considered
 135  * copy-on-write friendly. */
 136 int setTypeRandomElement(robj *setobj, robj **objele, long long *llele) {
 137     if (setobj->encoding == REDIS_ENCODING_HT) {
 138         dictEntry *de = dictGetRandomKey(setobj->ptr);
 139         *objele = dictGetEntryKey(de);
 140     } else if (setobj->encoding == REDIS_ENCODING_INTSET) {
 141         *llele = intsetRandom(setobj->ptr);
 142     } else {
 143         redisPanic("Unknown set encoding");
 144     }
 145     return setobj->encoding;
 146 }
 147
 148 unsigned long setTypeSize(robj *subject) {
 149     if (subject->encoding == REDIS_ENCODING_HT) {
 150         return dictSize((dict*)subject->ptr);
 151     } else if (subject->encoding == REDIS_ENCODING_INTSET) {
 152         return intsetLen((intset*)subject->ptr);
 153     } else {
 154         redisPanic("Unknown set encoding");
 155     }
 156 }
 157
 158 /* Convert the set to specified encoding. The resulting dict (when converting
 159  * to a hashtable) is presized to hold the number of elements in the original
 160  * set. */
 161 void setTypeConvert(robj *subject, int enc) {
 162     setTypeIterator *si;
 163     robj *element;
 164     redisAssert(subject->type == REDIS_SET);
 165
 166     if (enc == REDIS_ENCODING_HT) {
 167         dict *d = dictCreate(&setDictType,NULL);
 168         /* Presize the dict to avoid rehashing */
 169         dictExpand(d,intsetLen(subject->ptr));
 170
 171         /* setTypeGet returns a robj with incremented refcount */
 172         si = setTypeInitIterator(subject);
 173         while ((element = setTypeNext(si)) != NULL)
 174             redisAssert(dictAdd(d,element,NULL) == DICT_OK);
 175         setTypeReleaseIterator(si);
 176
 177         subject->encoding = REDIS_ENCODING_HT;
 178         zfree(subject->ptr);
 179         subject->ptr = d;
 180     } else {
 181         redisPanic("Unsupported set conversion");
 182     }
 183 }
 184
 185 void saddCommand(redisClient *c) {
 186     robj *set;
 187
 188     set = lookupKeyWrite(c->db,c->argv[1]);
 189     c->argv[2] = tryObjectEncoding(c->argv[2]);
 190     if (set == NULL) {
 191         set = setTypeCreate(c->argv[2]);
 192         dbAdd(c->db,c->argv[1],set);
 193     } else {
 194         if (set->type != REDIS_SET) {
 195             addReply(c,shared.wrongtypeerr);
 196             return;
 197         }
 198     }
 199     if (setTypeAdd(set,c->argv[2])) {
 200         touchWatchedKey(c->db,c->argv[1]);
 201         server.dirty++;
 202         addReply(c,shared.cone);
 203     } else {
 204         addReply(c,shared.czero);
 205     }
 206 }
 207
 208 void sremCommand(redisClient *c) {
 209     robj *set;
 210
 211     if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
 212         checkType(c,set,REDIS_SET)) return;
 213
 214     c->argv[2] = tryObjectEncoding(c->argv[2]);
 215     if (setTypeRemove(set,c->argv[2])) {
 216         if (setTypeSize(set) == 0) dbDelete(c->db,c->argv[1]);
 217         touchWatchedKey(c->db,c->argv[1]);
 218         server.dirty++;
 219         addReply(c,shared.cone);
 220     } else {
 221         addReply(c,shared.czero);
 222     }
 223 }
 224
 225 void smoveCommand(redisClient *c) {
 226     robj *srcset, *dstset, *ele;
 227     srcset = lookupKeyWrite(c->db,c->argv[1]);
 228     dstset = lookupKeyWrite(c->db,c->argv[2]);
 229     ele = c->argv[3] = tryObjectEncoding(c->argv[3]);
 230
 231     /* If the source key does not exist return 0 */
 232     if (srcset == NULL) {
 233         addReply(c,shared.czero);
 234         return;
 235     }
 236
 237     /* If the source key has the wrong type, or the destination key
 238      * is set and has the wrong type, return with an error. */
 239     if (checkType(c,srcset,REDIS_SET) ||
 240         (dstset && checkType(c,dstset,REDIS_SET))) return;
 241
 242     /* If srcset and dstset are equal, SMOVE is a no-op */
 243     if (srcset == dstset) {
 244         addReply(c,shared.cone);
 245         return;
 246     }
 247
 248     /* If the element cannot be removed from the src set, return 0. */
 249     if (!setTypeRemove(srcset,ele)) {
 250         addReply(c,shared.czero);
 251         return;
 252     }
 253
 254     /* Remove the src set from the database when empty */
 255     if (setTypeSize(srcset) == 0) dbDelete(c->db,c->argv[1]);
 256     touchWatchedKey(c->db,c->argv[1]);
 257     touchWatchedKey(c->db,c->argv[2]);
 258     server.dirty++;
 259
 260     /* Create the destination set when it doesn't exist */
 261     if (!dstset) {
 262         dstset = setTypeCreate(ele);
 263         dbAdd(c->db,c->argv[2],dstset);
 264     }
 265
 266     /* An extra key has changed when ele was successfully added to dstset */
 267     if (setTypeAdd(dstset,ele)) server.dirty++;
 268     addReply(c,shared.cone);
 269 }
 270
 271 void sismemberCommand(redisClient *c) {
 272     robj *set;
 273
 274     if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
 275         checkType(c,set,REDIS_SET)) return;
 276
 277     c->argv[2] = tryObjectEncoding(c->argv[2]);
 278     if (setTypeIsMember(set,c->argv[2]))
 279         addReply(c,shared.cone);
 280     else
 281         addReply(c,shared.czero);
 282 }
 283
 284 void scardCommand(redisClient *c) {
 285     robj *o;
 286
 287     if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
 288         checkType(c,o,REDIS_SET)) return;
 289
 290     addReplyLongLong(c,setTypeSize(o));
 291 }
 292
 293 void spopCommand(redisClient *c) {
 294     robj *set, *ele;
 295     long long llele;
 296     int encoding;
 297
 298     if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
 299         checkType(c,set,REDIS_SET)) return;
 300
 301     encoding = setTypeRandomElement(set,&ele,&llele);
 302     if (encoding == REDIS_ENCODING_INTSET) {
 303         addReplyBulkLongLong(c,llele);
 304         set->ptr = intsetRemove(set->ptr,llele,NULL);
 305     } else {
 306         addReplyBulk(c,ele);
 307         setTypeRemove(set,ele);
 308     }
 309     if (setTypeSize(set) == 0) dbDelete(c->db,c->argv[1]);
 310     touchWatchedKey(c->db,c->argv[1]);
 311     server.dirty++;
 312 }
 313
 314 void srandmemberCommand(redisClient *c) {
 315     robj *set, *ele;
 316     long long llele;
 317     int encoding;
 318
 319     if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
 320         checkType(c,set,REDIS_SET)) return;
 321
 322     encoding = setTypeRandomElement(set,&ele,&llele);
 323     if (encoding == REDIS_ENCODING_INTSET) {
 324         addReplyBulkLongLong(c,llele);
 325     } else {
 326         addReplyBulk(c,ele);
 327     }
 328 }
 329
 330 int qsortCompareSetsByCardinality(const void *s1, const void *s2) {
 331     return setTypeSize(*(robj**)s1)-setTypeSize(*(robj**)s2);
 332 }
 333
 334 void sinterGenericCommand(redisClient *c, robj **setkeys, unsigned long setnum, robj *dstkey) {
 335     robj **sets = zmalloc(sizeof(robj*)*setnum);
 336     setTypeIterator *si;
 337     robj *ele, *dstset = NULL;
 338     void *replylen = NULL;
 339     unsigned long j, cardinality = 0;
 340
 341     for (j = 0; j < setnum; j++) {
 342         robj *setobj = dstkey ?
 343             lookupKeyWrite(c->db,setkeys[j]) :
 344             lookupKeyRead(c->db,setkeys[j]);
 345         if (!setobj) {
 346             zfree(sets);
 347             if (dstkey) {
 348                 if (dbDelete(c->db,dstkey)) {
 349                     touchWatchedKey(c->db,dstkey);
 350                     server.dirty++;
 351                 }
 352                 addReply(c,shared.czero);
 353             } else {
 354                 addReply(c,shared.emptymultibulk);
 355             }
 356             return;
 357         }
 358         if (checkType(c,setobj,REDIS_SET)) {
 359             zfree(sets);
 360             return;
 361         }
 362         sets[j] = setobj;
 363     }
 364     /* Sort sets from the smallest to largest, this will improve our
 365      * algorithm's performace */
 366     qsort(sets,setnum,sizeof(robj*),qsortCompareSetsByCardinality);
 367
 368     /* The first thing we should output is the total number of elements...
 369      * since this is a multi-bulk write, but at this stage we don't know
 370      * the intersection set size, so we use a trick, append an empty object
 371      * to the output list and save the pointer to later modify it with the
 372      * right length */
 373     if (!dstkey) {
 374         replylen = addDeferredMultiBulkLength(c);
 375     } else {
 376         /* If we have a target key where to store the resulting set
 377          * create this key with an empty set inside */
 378         dstset = createIntsetObject();
 379     }
 380
 381     /* Iterate all the elements of the first (smallest) set, and test
 382      * the element against all the other sets, if at least one set does
 383      * not include the element it is discarded */
 384     si = setTypeInitIterator(sets[0]);
 385     while((ele = setTypeNext(si)) != NULL) {
 386         for (j = 1; j < setnum; j++)
 387             if (!setTypeIsMember(sets[j],ele)) break;
 388
 389         /* Only take action when all sets contain the member */
 390         if (j == setnum) {
 391             if (!dstkey) {
 392                 addReplyBulk(c,ele);
 393                 cardinality++;
 394             } else {
 395                 setTypeAdd(dstset,ele);
 396             }
 397         }
 398         decrRefCount(ele);
 399     }
 400     setTypeReleaseIterator(si);
 401
 402     if (dstkey) {
 403         /* Store the resulting set into the target, if the intersection
 404          * is not an empty set. */
 405         dbDelete(c->db,dstkey);
 406         if (setTypeSize(dstset) > 0) {
 407             dbAdd(c->db,dstkey,dstset);
 408             addReplyLongLong(c,setTypeSize(dstset));
 409         } else {
 410             decrRefCount(dstset);
 411             addReply(c,shared.czero);
 412         }
 413         touchWatchedKey(c->db,dstkey);
 414         server.dirty++;
 415     } else {
 416         setDeferredMultiBulkLength(c,replylen,cardinality);
 417     }
 418     zfree(sets);
 419 }
 420
 421 void sinterCommand(redisClient *c) {
 422     sinterGenericCommand(c,c->argv+1,c->argc-1,NULL);
 423 }
 424
 425 void sinterstoreCommand(redisClient *c) {
 426     sinterGenericCommand(c,c->argv+2,c->argc-2,c->argv[1]);
 427 }
 428
 429 #define REDIS_OP_UNION 0
 430 #define REDIS_OP_DIFF 1
 431 #define REDIS_OP_INTER 2
 432
 433 void sunionDiffGenericCommand(redisClient *c, robj **setkeys, int setnum, robj *dstkey, int op) {
 434     robj **sets = zmalloc(sizeof(robj*)*setnum);
 435     setTypeIterator *si;
 436     robj *ele, *dstset = NULL;
 437     int j, cardinality = 0;
 438
 439     for (j = 0; j < setnum; j++) {
 440         robj *setobj = dstkey ?
 441             lookupKeyWrite(c->db,setkeys[j]) :
 442             lookupKeyRead(c->db,setkeys[j]);
 443         if (!setobj) {
 444             sets[j] = NULL;
 445             continue;
 446         }
 447         if (checkType(c,setobj,REDIS_SET)) {
 448             zfree(sets);
 449             return;
 450         }
 451         sets[j] = setobj;
 452     }
 453
 454     /* We need a temp set object to store our union. If the dstkey
 455      * is not NULL (that is, we are inside an SUNIONSTORE operation) then
 456      * this set object will be the resulting object to set into the target key*/
 457     dstset = createIntsetObject();
 458
 459     /* Iterate all the elements of all the sets, add every element a single
 460      * time to the result set */
 461     for (j = 0; j < setnum; j++) {
 462         if (op == REDIS_OP_DIFF && j == 0 && !sets[j]) break; /* result set is empty */
 463         if (!sets[j]) continue; /* non existing keys are like empty sets */
 464
 465         si = setTypeInitIterator(sets[j]);
 466         while((ele = setTypeNext(si)) != NULL) {
 467             if (op == REDIS_OP_UNION || j == 0) {
 468                 if (setTypeAdd(dstset,ele)) {
 469                     cardinality++;
 470                 }
 471             } else if (op == REDIS_OP_DIFF) {
 472                 if (setTypeRemove(dstset,ele)) {
 473                     cardinality--;
 474                 }
 475             }
 476             decrRefCount(ele);
 477         }
 478         setTypeReleaseIterator(si);
 479
 480         /* Exit when result set is empty. */
 481         if (op == REDIS_OP_DIFF && cardinality == 0) break;
 482     }
 483
 484     /* Output the content of the resulting set, if not in STORE mode */
 485     if (!dstkey) {
 486         addReplyMultiBulkLen(c,cardinality);
 487         si = setTypeInitIterator(dstset);
 488         while((ele = setTypeNext(si)) != NULL) {
 489             addReplyBulk(c,ele);
 490             decrRefCount(ele);
 491         }
 492         setTypeReleaseIterator(si);
 493         decrRefCount(dstset);
 494     } else {
 495         /* If we have a target key where to store the resulting set
 496          * create this key with the result set inside */
 497         dbDelete(c->db,dstkey);
 498         if (setTypeSize(dstset) > 0) {
 499             dbAdd(c->db,dstkey,dstset);
 500             addReplyLongLong(c,setTypeSize(dstset));
 501         } else {
 502             decrRefCount(dstset);
 503             addReply(c,shared.czero);
 504         }
 505         touchWatchedKey(c->db,dstkey);
 506         server.dirty++;
 507     }
 508     zfree(sets);
 509 }
 510
 511 void sunionCommand(redisClient *c) {
 512     sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_UNION);
 513 }
 514
 515 void sunionstoreCommand(redisClient *c) {
 516     sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_UNION);
 517 }
 518
 519 void sdiffCommand(redisClient *c) {
 520     sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_DIFF);
 521 }
 522
 523 void sdiffstoreCommand(redisClient *c) {
 524     sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_DIFF);
 525 }