]> git.saurik.com Git - redis.git/blame - src/t_set.c
Use existing reply functions where possible
[redis.git] / src / t_set.c
CommitLineData
e2641e09 1#include "redis.h"
2
3/*-----------------------------------------------------------------------------
4 * Set Commands
5 *----------------------------------------------------------------------------*/
6
96ffb2fe
PN
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. */
10robj *setTypeCreate(robj *value) {
ec7e1389 11 if (isObjectRepresentableAsLongLong(value,NULL) == REDIS_OK)
96ffb2fe
PN
12 return createIntsetObject();
13 return createSetObject();
14}
15
16int 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) {
ec7e1389 24 if (isObjectRepresentableAsLongLong(value,&llval) == REDIS_OK) {
96ffb2fe
PN
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
50int 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) {
ec7e1389 58 if (isObjectRepresentableAsLongLong(value,&llval) == REDIS_OK) {
96ffb2fe
PN
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
69int 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) {
ec7e1389 74 if (isObjectRepresentableAsLongLong(value,&llval) == REDIS_OK) {
96ffb2fe
PN
75 return intsetFind((intset*)subject->ptr,llval);
76 }
77 } else {
78 redisPanic("Unknown set encoding");
79 }
80 return 0;
81}
82
cb72d0f1 83setTypeIterator *setTypeInitIterator(robj *subject) {
740eee1c 84 setTypeIterator *si = zmalloc(sizeof(setTypeIterator));
96ffb2fe
PN
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
cb72d0f1 97void setTypeReleaseIterator(setTypeIterator *si) {
96ffb2fe
PN
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. */
cb72d0f1 106robj *setTypeNext(setTypeIterator *si) {
96ffb2fe
PN
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) {
23c64fe5 115 int64_t llval;
96ffb2fe
PN
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. */
125robj *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
140unsigned 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. */
153void setTypeConvert(robj *subject, int enc) {
cb72d0f1 154 setTypeIterator *si;
96ffb2fe
PN
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
e2641e09 177void saddCommand(redisClient *c) {
178 robj *set;
179
180 set = lookupKeyWrite(c->db,c->argv[1]);
181 if (set == NULL) {
96ffb2fe 182 set = setTypeCreate(c->argv[2]);
e2641e09 183 dbAdd(c->db,c->argv[1],set);
184 } else {
185 if (set->type != REDIS_SET) {
186 addReply(c,shared.wrongtypeerr);
187 return;
188 }
189 }
96ffb2fe 190 if (setTypeAdd(set,c->argv[2])) {
5b4bff9c 191 touchWatchedKey(c->db,c->argv[1]);
e2641e09 192 server.dirty++;
193 addReply(c,shared.cone);
194 } else {
195 addReply(c,shared.czero);
196 }
197}
198
199void sremCommand(redisClient *c) {
200 robj *set;
201
202 if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
203 checkType(c,set,REDIS_SET)) return;
204
96ffb2fe
PN
205 if (setTypeRemove(set,c->argv[2])) {
206 if (setTypeSize(set) == 0) dbDelete(c->db,c->argv[1]);
5b4bff9c 207 touchWatchedKey(c->db,c->argv[1]);
e2641e09 208 server.dirty++;
e2641e09 209 addReply(c,shared.cone);
210 } else {
211 addReply(c,shared.czero);
212 }
213}
214
215void smoveCommand(redisClient *c) {
96ffb2fe 216 robj *srcset, *dstset, *ele;
e2641e09 217 srcset = lookupKeyWrite(c->db,c->argv[1]);
218 dstset = lookupKeyWrite(c->db,c->argv[2]);
96ffb2fe 219 ele = c->argv[3];
e2641e09 220
96ffb2fe
PN
221 /* If the source key does not exist return 0 */
222 if (srcset == NULL) {
223 addReply(c,shared.czero);
e2641e09 224 return;
225 }
96ffb2fe
PN
226
227 /* If the source key has the wrong type, or the destination key
228 * is set and has the wrong type, return with an error. */
229 if (checkType(c,srcset,REDIS_SET) ||
230 (dstset && checkType(c,dstset,REDIS_SET))) return;
231
232 /* If srcset and dstset are equal, SMOVE is a no-op */
233 if (srcset == dstset) {
234 addReply(c,shared.cone);
e2641e09 235 return;
236 }
96ffb2fe
PN
237
238 /* If the element cannot be removed from the src set, return 0. */
239 if (!setTypeRemove(srcset,ele)) {
e2641e09 240 addReply(c,shared.czero);
241 return;
242 }
96ffb2fe
PN
243
244 /* Remove the src set from the database when empty */
245 if (setTypeSize(srcset) == 0) dbDelete(c->db,c->argv[1]);
5b4bff9c 246 touchWatchedKey(c->db,c->argv[1]);
247 touchWatchedKey(c->db,c->argv[2]);
e2641e09 248 server.dirty++;
96ffb2fe
PN
249
250 /* Create the destination set when it doesn't exist */
e2641e09 251 if (!dstset) {
96ffb2fe 252 dstset = setTypeCreate(ele);
e2641e09 253 dbAdd(c->db,c->argv[2],dstset);
254 }
96ffb2fe
PN
255
256 /* An extra key has changed when ele was successfully added to dstset */
257 if (setTypeAdd(dstset,ele)) server.dirty++;
e2641e09 258 addReply(c,shared.cone);
259}
260
261void sismemberCommand(redisClient *c) {
262 robj *set;
263
264 if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
265 checkType(c,set,REDIS_SET)) return;
266
96ffb2fe 267 if (setTypeIsMember(set,c->argv[2]))
e2641e09 268 addReply(c,shared.cone);
269 else
270 addReply(c,shared.czero);
271}
272
273void scardCommand(redisClient *c) {
274 robj *o;
e2641e09 275
276 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
277 checkType(c,o,REDIS_SET)) return;
278
b70d3555 279 addReplyLongLong(c,setTypeSize(o));
e2641e09 280}
281
282void spopCommand(redisClient *c) {
96ffb2fe 283 robj *set, *ele;
e2641e09 284
285 if ((set = lookupKeyWriteOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
286 checkType(c,set,REDIS_SET)) return;
287
96ffb2fe
PN
288 ele = setTypeRandomElement(set);
289 if (ele == NULL) {
e2641e09 290 addReply(c,shared.nullbulk);
291 } else {
96ffb2fe 292 setTypeRemove(set,ele);
e2641e09 293 addReplyBulk(c,ele);
96ffb2fe
PN
294 decrRefCount(ele);
295 if (setTypeSize(set) == 0) dbDelete(c->db,c->argv[1]);
5b4bff9c 296 touchWatchedKey(c->db,c->argv[1]);
e2641e09 297 server.dirty++;
298 }
299}
300
301void srandmemberCommand(redisClient *c) {
96ffb2fe 302 robj *set, *ele;
e2641e09 303
304 if ((set = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
305 checkType(c,set,REDIS_SET)) return;
306
96ffb2fe
PN
307 ele = setTypeRandomElement(set);
308 if (ele == NULL) {
e2641e09 309 addReply(c,shared.nullbulk);
310 } else {
e2641e09 311 addReplyBulk(c,ele);
96ffb2fe 312 decrRefCount(ele);
e2641e09 313 }
314}
315
316int qsortCompareSetsByCardinality(const void *s1, const void *s2) {
96ffb2fe 317 return setTypeSize(*(robj**)s1)-setTypeSize(*(robj**)s2);
e2641e09 318}
319
96ffb2fe
PN
320void sinterGenericCommand(redisClient *c, robj **setkeys, unsigned long setnum, robj *dstkey) {
321 robj **sets = zmalloc(sizeof(robj*)*setnum);
cb72d0f1 322 setTypeIterator *si;
b301c1fc
PN
323 robj *ele, *dstset = NULL;
324 void *replylen = NULL;
e2641e09 325 unsigned long j, cardinality = 0;
326
96ffb2fe
PN
327 for (j = 0; j < setnum; j++) {
328 robj *setobj = dstkey ?
329 lookupKeyWrite(c->db,setkeys[j]) :
330 lookupKeyRead(c->db,setkeys[j]);
e2641e09 331 if (!setobj) {
96ffb2fe 332 zfree(sets);
e2641e09 333 if (dstkey) {
5b4bff9c 334 if (dbDelete(c->db,dstkey)) {
335 touchWatchedKey(c->db,dstkey);
e2641e09 336 server.dirty++;
5b4bff9c 337 }
e2641e09 338 addReply(c,shared.czero);
339 } else {
340 addReply(c,shared.emptymultibulk);
341 }
342 return;
343 }
96ffb2fe
PN
344 if (checkType(c,setobj,REDIS_SET)) {
345 zfree(sets);
e2641e09 346 return;
347 }
96ffb2fe 348 sets[j] = setobj;
e2641e09 349 }
350 /* Sort sets from the smallest to largest, this will improve our
351 * algorithm's performace */
96ffb2fe 352 qsort(sets,setnum,sizeof(robj*),qsortCompareSetsByCardinality);
e2641e09 353
354 /* The first thing we should output is the total number of elements...
355 * since this is a multi-bulk write, but at this stage we don't know
356 * the intersection set size, so we use a trick, append an empty object
357 * to the output list and save the pointer to later modify it with the
358 * right length */
359 if (!dstkey) {
b301c1fc 360 replylen = addDeferredMultiBulkLength(c);
e2641e09 361 } else {
362 /* If we have a target key where to store the resulting set
363 * create this key with an empty set inside */
96ffb2fe 364 dstset = createIntsetObject();
e2641e09 365 }
366
367 /* Iterate all the elements of the first (smallest) set, and test
368 * the element against all the other sets, if at least one set does
369 * not include the element it is discarded */
96ffb2fe
PN
370 si = setTypeInitIterator(sets[0]);
371 while((ele = setTypeNext(si)) != NULL) {
372 for (j = 1; j < setnum; j++)
373 if (!setTypeIsMember(sets[j],ele)) break;
374
375 /* Only take action when all sets contain the member */
376 if (j == setnum) {
377 if (!dstkey) {
378 addReplyBulk(c,ele);
379 cardinality++;
380 } else {
381 setTypeAdd(dstset,ele);
382 }
e2641e09 383 }
96ffb2fe 384 decrRefCount(ele);
e2641e09 385 }
96ffb2fe 386 setTypeReleaseIterator(si);
e2641e09 387
388 if (dstkey) {
389 /* Store the resulting set into the target, if the intersection
390 * is not an empty set. */
391 dbDelete(c->db,dstkey);
96ffb2fe 392 if (setTypeSize(dstset) > 0) {
e2641e09 393 dbAdd(c->db,dstkey,dstset);
96ffb2fe 394 addReplyLongLong(c,setTypeSize(dstset));
e2641e09 395 } else {
396 decrRefCount(dstset);
397 addReply(c,shared.czero);
398 }
5b4bff9c 399 touchWatchedKey(c->db,dstkey);
e2641e09 400 server.dirty++;
401 } else {
b301c1fc 402 setDeferredMultiBulkLength(c,replylen,cardinality);
e2641e09 403 }
96ffb2fe 404 zfree(sets);
e2641e09 405}
406
407void sinterCommand(redisClient *c) {
408 sinterGenericCommand(c,c->argv+1,c->argc-1,NULL);
409}
410
411void sinterstoreCommand(redisClient *c) {
412 sinterGenericCommand(c,c->argv+2,c->argc-2,c->argv[1]);
413}
414
96ffb2fe
PN
415#define REDIS_OP_UNION 0
416#define REDIS_OP_DIFF 1
417#define REDIS_OP_INTER 2
e2641e09 418
96ffb2fe
PN
419void sunionDiffGenericCommand(redisClient *c, robj **setkeys, int setnum, robj *dstkey, int op) {
420 robj **sets = zmalloc(sizeof(robj*)*setnum);
cb72d0f1 421 setTypeIterator *si;
96ffb2fe
PN
422 robj *ele, *dstset = NULL;
423 int j, cardinality = 0;
e2641e09 424
96ffb2fe
PN
425 for (j = 0; j < setnum; j++) {
426 robj *setobj = dstkey ?
427 lookupKeyWrite(c->db,setkeys[j]) :
428 lookupKeyRead(c->db,setkeys[j]);
e2641e09 429 if (!setobj) {
96ffb2fe 430 sets[j] = NULL;
e2641e09 431 continue;
432 }
96ffb2fe
PN
433 if (checkType(c,setobj,REDIS_SET)) {
434 zfree(sets);
e2641e09 435 return;
436 }
96ffb2fe 437 sets[j] = setobj;
e2641e09 438 }
439
440 /* We need a temp set object to store our union. If the dstkey
441 * is not NULL (that is, we are inside an SUNIONSTORE operation) then
442 * this set object will be the resulting object to set into the target key*/
96ffb2fe 443 dstset = createIntsetObject();
e2641e09 444
445 /* Iterate all the elements of all the sets, add every element a single
446 * time to the result set */
96ffb2fe
PN
447 for (j = 0; j < setnum; j++) {
448 if (op == REDIS_OP_DIFF && j == 0 && !sets[j]) break; /* result set is empty */
449 if (!sets[j]) continue; /* non existing keys are like empty sets */
e2641e09 450
96ffb2fe
PN
451 si = setTypeInitIterator(sets[j]);
452 while((ele = setTypeNext(si)) != NULL) {
e2641e09 453 if (op == REDIS_OP_UNION || j == 0) {
96ffb2fe 454 if (setTypeAdd(dstset,ele)) {
e2641e09 455 cardinality++;
456 }
457 } else if (op == REDIS_OP_DIFF) {
96ffb2fe 458 if (setTypeRemove(dstset,ele)) {
e2641e09 459 cardinality--;
460 }
461 }
96ffb2fe 462 decrRefCount(ele);
e2641e09 463 }
96ffb2fe 464 setTypeReleaseIterator(si);
e2641e09 465
96ffb2fe 466 /* Exit when result set is empty. */
e2641e09 467 if (op == REDIS_OP_DIFF && cardinality == 0) break;
468 }
469
470 /* Output the content of the resulting set, if not in STORE mode */
471 if (!dstkey) {
0537e7bf 472 addReplyMultiBulkLen(c,cardinality);
96ffb2fe
PN
473 si = setTypeInitIterator(dstset);
474 while((ele = setTypeNext(si)) != NULL) {
e2641e09 475 addReplyBulk(c,ele);
96ffb2fe 476 decrRefCount(ele);
e2641e09 477 }
96ffb2fe 478 setTypeReleaseIterator(si);
e2641e09 479 decrRefCount(dstset);
480 } else {
481 /* If we have a target key where to store the resulting set
482 * create this key with the result set inside */
483 dbDelete(c->db,dstkey);
96ffb2fe 484 if (setTypeSize(dstset) > 0) {
e2641e09 485 dbAdd(c->db,dstkey,dstset);
96ffb2fe 486 addReplyLongLong(c,setTypeSize(dstset));
e2641e09 487 } else {
488 decrRefCount(dstset);
489 addReply(c,shared.czero);
490 }
5b4bff9c 491 touchWatchedKey(c->db,dstkey);
e2641e09 492 server.dirty++;
493 }
96ffb2fe 494 zfree(sets);
e2641e09 495}
496
497void sunionCommand(redisClient *c) {
498 sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_UNION);
499}
500
501void sunionstoreCommand(redisClient *c) {
502 sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_UNION);
503}
504
505void sdiffCommand(redisClient *c) {
506 sunionDiffGenericCommand(c,c->argv+1,c->argc-1,NULL,REDIS_OP_DIFF);
507}
508
509void sdiffstoreCommand(redisClient *c) {
510 sunionDiffGenericCommand(c,c->argv+2,c->argc-2,c->argv[1],REDIS_OP_DIFF);
511}