]> git.saurik.com Git - redis.git/blob - src/t_zset.c
f1a10378cbda8e596468404fae0eb9da280d6542
[redis.git] / src / t_zset.c
1 #include "redis.h"
2
3 #include <math.h>
4
5 /*-----------------------------------------------------------------------------
6 * Sorted set API
7 *----------------------------------------------------------------------------*/
8
9 /* ZSETs are ordered sets using two data structures to hold the same elements
10 * in order to get O(log(N)) INSERT and REMOVE operations into a sorted
11 * data structure.
12 *
13 * The elements are added to an hash table mapping Redis objects to scores.
14 * At the same time the elements are added to a skip list mapping scores
15 * to Redis objects (so objects are sorted by scores in this "view"). */
16
17 /* This skiplist implementation is almost a C translation of the original
18 * algorithm described by William Pugh in "Skip Lists: A Probabilistic
19 * Alternative to Balanced Trees", modified in three ways:
20 * a) this implementation allows for repeated values.
21 * b) the comparison is not just by key (our 'score') but by satellite data.
22 * c) there is a back pointer, so it's a doubly linked list with the back
23 * pointers being only at "level 1". This allows to traverse the list
24 * from tail to head, useful for ZREVRANGE. */
25
26 zskiplistNode *zslCreateNode(int level, double score, robj *obj) {
27 zskiplistNode *zn = zmalloc(sizeof(*zn)+level*sizeof(struct zskiplistLevel));
28 zn->score = score;
29 zn->obj = obj;
30 return zn;
31 }
32
33 zskiplist *zslCreate(void) {
34 int j;
35 zskiplist *zsl;
36
37 zsl = zmalloc(sizeof(*zsl));
38 zsl->level = 1;
39 zsl->length = 0;
40 zsl->header = zslCreateNode(ZSKIPLIST_MAXLEVEL,0,NULL);
41 for (j = 0; j < ZSKIPLIST_MAXLEVEL; j++) {
42 zsl->header->level[j].forward = NULL;
43 zsl->header->level[j].span = 0;
44 }
45 zsl->header->backward = NULL;
46 zsl->tail = NULL;
47 return zsl;
48 }
49
50 void zslFreeNode(zskiplistNode *node) {
51 decrRefCount(node->obj);
52 zfree(node);
53 }
54
55 void zslFree(zskiplist *zsl) {
56 zskiplistNode *node = zsl->header->level[0].forward, *next;
57
58 zfree(zsl->header);
59 while(node) {
60 next = node->level[0].forward;
61 zslFreeNode(node);
62 node = next;
63 }
64 zfree(zsl);
65 }
66
67 int zslRandomLevel(void) {
68 int level = 1;
69 while ((random()&0xFFFF) < (ZSKIPLIST_P * 0xFFFF))
70 level += 1;
71 return (level<ZSKIPLIST_MAXLEVEL) ? level : ZSKIPLIST_MAXLEVEL;
72 }
73
74 zskiplistNode *zslInsert(zskiplist *zsl, double score, robj *obj) {
75 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
76 unsigned int rank[ZSKIPLIST_MAXLEVEL];
77 int i, level;
78
79 x = zsl->header;
80 for (i = zsl->level-1; i >= 0; i--) {
81 /* store rank that is crossed to reach the insert position */
82 rank[i] = i == (zsl->level-1) ? 0 : rank[i+1];
83 while (x->level[i].forward &&
84 (x->level[i].forward->score < score ||
85 (x->level[i].forward->score == score &&
86 compareStringObjects(x->level[i].forward->obj,obj) < 0))) {
87 rank[i] += x->level[i].span;
88 x = x->level[i].forward;
89 }
90 update[i] = x;
91 }
92 /* we assume the key is not already inside, since we allow duplicated
93 * scores, and the re-insertion of score and redis object should never
94 * happpen since the caller of zslInsert() should test in the hash table
95 * if the element is already inside or not. */
96 level = zslRandomLevel();
97 if (level > zsl->level) {
98 for (i = zsl->level; i < level; i++) {
99 rank[i] = 0;
100 update[i] = zsl->header;
101 update[i]->level[i].span = zsl->length;
102 }
103 zsl->level = level;
104 }
105 x = zslCreateNode(level,score,obj);
106 for (i = 0; i < level; i++) {
107 x->level[i].forward = update[i]->level[i].forward;
108 update[i]->level[i].forward = x;
109
110 /* update span covered by update[i] as x is inserted here */
111 x->level[i].span = update[i]->level[i].span - (rank[0] - rank[i]);
112 update[i]->level[i].span = (rank[0] - rank[i]) + 1;
113 }
114
115 /* increment span for untouched levels */
116 for (i = level; i < zsl->level; i++) {
117 update[i]->level[i].span++;
118 }
119
120 x->backward = (update[0] == zsl->header) ? NULL : update[0];
121 if (x->level[0].forward)
122 x->level[0].forward->backward = x;
123 else
124 zsl->tail = x;
125 zsl->length++;
126 return x;
127 }
128
129 /* Internal function used by zslDelete, zslDeleteByScore and zslDeleteByRank */
130 void zslDeleteNode(zskiplist *zsl, zskiplistNode *x, zskiplistNode **update) {
131 int i;
132 for (i = 0; i < zsl->level; i++) {
133 if (update[i]->level[i].forward == x) {
134 update[i]->level[i].span += x->level[i].span - 1;
135 update[i]->level[i].forward = x->level[i].forward;
136 } else {
137 update[i]->level[i].span -= 1;
138 }
139 }
140 if (x->level[0].forward) {
141 x->level[0].forward->backward = x->backward;
142 } else {
143 zsl->tail = x->backward;
144 }
145 while(zsl->level > 1 && zsl->header->level[zsl->level-1].forward == NULL)
146 zsl->level--;
147 zsl->length--;
148 }
149
150 /* Delete an element with matching score/object from the skiplist. */
151 int zslDelete(zskiplist *zsl, double score, robj *obj) {
152 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
153 int i;
154
155 x = zsl->header;
156 for (i = zsl->level-1; i >= 0; i--) {
157 while (x->level[i].forward &&
158 (x->level[i].forward->score < score ||
159 (x->level[i].forward->score == score &&
160 compareStringObjects(x->level[i].forward->obj,obj) < 0)))
161 x = x->level[i].forward;
162 update[i] = x;
163 }
164 /* We may have multiple elements with the same score, what we need
165 * is to find the element with both the right score and object. */
166 x = x->level[0].forward;
167 if (x && score == x->score && equalStringObjects(x->obj,obj)) {
168 zslDeleteNode(zsl, x, update);
169 zslFreeNode(x);
170 return 1;
171 } else {
172 return 0; /* not found */
173 }
174 return 0; /* not found */
175 }
176
177 /* Struct to hold a inclusive/exclusive range spec. */
178 typedef struct {
179 double min, max;
180 int minex, maxex; /* are min or max exclusive? */
181 } zrangespec;
182
183 /* Delete all the elements with score between min and max from the skiplist.
184 * Min and mx are inclusive, so a score >= min || score <= max is deleted.
185 * Note that this function takes the reference to the hash table view of the
186 * sorted set, in order to remove the elements from the hash table too. */
187 unsigned long zslDeleteRangeByScore(zskiplist *zsl, zrangespec range, dict *dict) {
188 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
189 unsigned long removed = 0;
190 int i;
191
192 x = zsl->header;
193 for (i = zsl->level-1; i >= 0; i--) {
194 while (x->level[i].forward && (range.minex ?
195 x->level[i].forward->score <= range.min :
196 x->level[i].forward->score < range.min))
197 x = x->level[i].forward;
198 update[i] = x;
199 }
200
201 /* Current node is the last with score < or <= min. */
202 x = x->level[0].forward;
203
204 /* Delete nodes while in range. */
205 while (x && (range.maxex ? x->score < range.max : x->score <= range.max)) {
206 zskiplistNode *next = x->level[0].forward;
207 zslDeleteNode(zsl,x,update);
208 dictDelete(dict,x->obj);
209 zslFreeNode(x);
210 removed++;
211 x = next;
212 }
213 return removed;
214 }
215
216 /* Delete all the elements with rank between start and end from the skiplist.
217 * Start and end are inclusive. Note that start and end need to be 1-based */
218 unsigned long zslDeleteRangeByRank(zskiplist *zsl, unsigned int start, unsigned int end, dict *dict) {
219 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
220 unsigned long traversed = 0, removed = 0;
221 int i;
222
223 x = zsl->header;
224 for (i = zsl->level-1; i >= 0; i--) {
225 while (x->level[i].forward && (traversed + x->level[i].span) < start) {
226 traversed += x->level[i].span;
227 x = x->level[i].forward;
228 }
229 update[i] = x;
230 }
231
232 traversed++;
233 x = x->level[0].forward;
234 while (x && traversed <= end) {
235 zskiplistNode *next = x->level[0].forward;
236 zslDeleteNode(zsl,x,update);
237 dictDelete(dict,x->obj);
238 zslFreeNode(x);
239 removed++;
240 traversed++;
241 x = next;
242 }
243 return removed;
244 }
245
246 /* Find the first node having a score equal or greater than the specified one.
247 * Returns NULL if there is no match. */
248 zskiplistNode *zslFirstWithScore(zskiplist *zsl, double score) {
249 zskiplistNode *x;
250 int i;
251
252 x = zsl->header;
253 for (i = zsl->level-1; i >= 0; i--) {
254 while (x->level[i].forward && x->level[i].forward->score < score)
255 x = x->level[i].forward;
256 }
257 /* We may have multiple elements with the same score, what we need
258 * is to find the element with both the right score and object. */
259 return x->level[0].forward;
260 }
261
262 /* Find the rank for an element by both score and key.
263 * Returns 0 when the element cannot be found, rank otherwise.
264 * Note that the rank is 1-based due to the span of zsl->header to the
265 * first element. */
266 unsigned long zslistTypeGetRank(zskiplist *zsl, double score, robj *o) {
267 zskiplistNode *x;
268 unsigned long rank = 0;
269 int i;
270
271 x = zsl->header;
272 for (i = zsl->level-1; i >= 0; i--) {
273 while (x->level[i].forward &&
274 (x->level[i].forward->score < score ||
275 (x->level[i].forward->score == score &&
276 compareStringObjects(x->level[i].forward->obj,o) <= 0))) {
277 rank += x->level[i].span;
278 x = x->level[i].forward;
279 }
280
281 /* x might be equal to zsl->header, so test if obj is non-NULL */
282 if (x->obj && equalStringObjects(x->obj,o)) {
283 return rank;
284 }
285 }
286 return 0;
287 }
288
289 /* Finds an element by its rank. The rank argument needs to be 1-based. */
290 zskiplistNode* zslistTypeGetElementByRank(zskiplist *zsl, unsigned long rank) {
291 zskiplistNode *x;
292 unsigned long traversed = 0;
293 int i;
294
295 x = zsl->header;
296 for (i = zsl->level-1; i >= 0; i--) {
297 while (x->level[i].forward && (traversed + x->level[i].span) <= rank)
298 {
299 traversed += x->level[i].span;
300 x = x->level[i].forward;
301 }
302 if (traversed == rank) {
303 return x;
304 }
305 }
306 return NULL;
307 }
308
309 /* Populate the rangespec according to the objects min and max. */
310 int zslParseRange(robj *min, robj *max, zrangespec *spec) {
311 spec->minex = spec->maxex = 0;
312
313 /* Parse the min-max interval. If one of the values is prefixed
314 * by the "(" character, it's considered "open". For instance
315 * ZRANGEBYSCORE zset (1.5 (2.5 will match min < x < max
316 * ZRANGEBYSCORE zset 1.5 2.5 will instead match min <= x <= max */
317 if (min->encoding == REDIS_ENCODING_INT) {
318 spec->min = (long)min->ptr;
319 } else {
320 if (((char*)min->ptr)[0] == '(') {
321 spec->min = strtod((char*)min->ptr+1,NULL);
322 spec->minex = 1;
323 } else {
324 spec->min = strtod((char*)min->ptr,NULL);
325 }
326 }
327 if (max->encoding == REDIS_ENCODING_INT) {
328 spec->max = (long)max->ptr;
329 } else {
330 if (((char*)max->ptr)[0] == '(') {
331 spec->max = strtod((char*)max->ptr+1,NULL);
332 spec->maxex = 1;
333 } else {
334 spec->max = strtod((char*)max->ptr,NULL);
335 }
336 }
337
338 return REDIS_OK;
339 }
340
341
342 /*-----------------------------------------------------------------------------
343 * Sorted set commands
344 *----------------------------------------------------------------------------*/
345
346 /* This generic command implements both ZADD and ZINCRBY. */
347 void zaddGenericCommand(redisClient *c, robj *key, robj *ele, double score, int incr) {
348 robj *zsetobj;
349 zset *zs;
350 zskiplistNode *znode;
351
352 zsetobj = lookupKeyWrite(c->db,key);
353 if (zsetobj == NULL) {
354 zsetobj = createZsetObject();
355 dbAdd(c->db,key,zsetobj);
356 } else {
357 if (zsetobj->type != REDIS_ZSET) {
358 addReply(c,shared.wrongtypeerr);
359 return;
360 }
361 }
362 zs = zsetobj->ptr;
363
364 /* Since both ZADD and ZINCRBY are implemented here, we need to increment
365 * the score first by the current score if ZINCRBY is called. */
366 if (incr) {
367 /* Read the old score. If the element was not present starts from 0 */
368 dictEntry *de = dictFind(zs->dict,ele);
369 if (de != NULL)
370 score += *(double*)dictGetEntryVal(de);
371
372 if (isnan(score)) {
373 addReplyError(c,"resulting score is not a number (NaN)");
374 /* Note that we don't need to check if the zset may be empty and
375 * should be removed here, as we can only obtain Nan as score if
376 * there was already an element in the sorted set. */
377 return;
378 }
379 }
380
381 /* We need to remove and re-insert the element when it was already present
382 * in the dictionary, to update the skiplist. Note that we delay adding a
383 * pointer to the score because we want to reference the score in the
384 * skiplist node. */
385 if (dictAdd(zs->dict,ele,NULL) == DICT_OK) {
386 dictEntry *de;
387
388 /* New element */
389 incrRefCount(ele); /* added to hash */
390 znode = zslInsert(zs->zsl,score,ele);
391 incrRefCount(ele); /* added to skiplist */
392
393 /* Update the score in the dict entry */
394 de = dictFind(zs->dict,ele);
395 redisAssert(de != NULL);
396 dictGetEntryVal(de) = &znode->score;
397 touchWatchedKey(c->db,c->argv[1]);
398 server.dirty++;
399 if (incr)
400 addReplyDouble(c,score);
401 else
402 addReply(c,shared.cone);
403 } else {
404 dictEntry *de;
405 robj *curobj;
406 double *curscore;
407 int deleted;
408
409 /* Update score */
410 de = dictFind(zs->dict,ele);
411 redisAssert(de != NULL);
412 curobj = dictGetEntryKey(de);
413 curscore = dictGetEntryVal(de);
414
415 /* When the score is updated, reuse the existing string object to
416 * prevent extra alloc/dealloc of strings on ZINCRBY. */
417 if (score != *curscore) {
418 deleted = zslDelete(zs->zsl,*curscore,curobj);
419 redisAssert(deleted != 0);
420 znode = zslInsert(zs->zsl,score,curobj);
421 incrRefCount(curobj);
422
423 /* Update the score in the current dict entry */
424 dictGetEntryVal(de) = &znode->score;
425 touchWatchedKey(c->db,c->argv[1]);
426 server.dirty++;
427 }
428 if (incr)
429 addReplyDouble(c,score);
430 else
431 addReply(c,shared.czero);
432 }
433 }
434
435 void zaddCommand(redisClient *c) {
436 double scoreval;
437 if (getDoubleFromObjectOrReply(c,c->argv[2],&scoreval,NULL) != REDIS_OK) return;
438 zaddGenericCommand(c,c->argv[1],c->argv[3],scoreval,0);
439 }
440
441 void zincrbyCommand(redisClient *c) {
442 double scoreval;
443 if (getDoubleFromObjectOrReply(c,c->argv[2],&scoreval,NULL) != REDIS_OK) return;
444 zaddGenericCommand(c,c->argv[1],c->argv[3],scoreval,1);
445 }
446
447 void zremCommand(redisClient *c) {
448 robj *zsetobj;
449 zset *zs;
450 dictEntry *de;
451 double curscore;
452 int deleted;
453
454 if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
455 checkType(c,zsetobj,REDIS_ZSET)) return;
456
457 zs = zsetobj->ptr;
458 de = dictFind(zs->dict,c->argv[2]);
459 if (de == NULL) {
460 addReply(c,shared.czero);
461 return;
462 }
463 /* Delete from the skiplist */
464 curscore = *(double*)dictGetEntryVal(de);
465 deleted = zslDelete(zs->zsl,curscore,c->argv[2]);
466 redisAssert(deleted != 0);
467
468 /* Delete from the hash table */
469 dictDelete(zs->dict,c->argv[2]);
470 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
471 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
472 touchWatchedKey(c->db,c->argv[1]);
473 server.dirty++;
474 addReply(c,shared.cone);
475 }
476
477 void zremrangebyscoreCommand(redisClient *c) {
478 zrangespec range;
479 long deleted;
480 robj *o;
481 zset *zs;
482
483 /* Parse the range arguments. */
484 zslParseRange(c->argv[2],c->argv[3],&range);
485
486 if ((o = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
487 checkType(c,o,REDIS_ZSET)) return;
488
489 zs = o->ptr;
490 deleted = zslDeleteRangeByScore(zs->zsl,range,zs->dict);
491 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
492 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
493 if (deleted) touchWatchedKey(c->db,c->argv[1]);
494 server.dirty += deleted;
495 addReplyLongLong(c,deleted);
496 }
497
498 void zremrangebyrankCommand(redisClient *c) {
499 long start;
500 long end;
501 int llen;
502 long deleted;
503 robj *zsetobj;
504 zset *zs;
505
506 if ((getLongFromObjectOrReply(c, c->argv[2], &start, NULL) != REDIS_OK) ||
507 (getLongFromObjectOrReply(c, c->argv[3], &end, NULL) != REDIS_OK)) return;
508
509 if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
510 checkType(c,zsetobj,REDIS_ZSET)) return;
511 zs = zsetobj->ptr;
512 llen = zs->zsl->length;
513
514 /* convert negative indexes */
515 if (start < 0) start = llen+start;
516 if (end < 0) end = llen+end;
517 if (start < 0) start = 0;
518
519 /* Invariant: start >= 0, so this test will be true when end < 0.
520 * The range is empty when start > end or start >= length. */
521 if (start > end || start >= llen) {
522 addReply(c,shared.czero);
523 return;
524 }
525 if (end >= llen) end = llen-1;
526
527 /* increment start and end because zsl*Rank functions
528 * use 1-based rank */
529 deleted = zslDeleteRangeByRank(zs->zsl,start+1,end+1,zs->dict);
530 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
531 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
532 if (deleted) touchWatchedKey(c->db,c->argv[1]);
533 server.dirty += deleted;
534 addReplyLongLong(c, deleted);
535 }
536
537 typedef struct {
538 dict *dict;
539 double weight;
540 } zsetopsrc;
541
542 int qsortCompareZsetopsrcByCardinality(const void *s1, const void *s2) {
543 zsetopsrc *d1 = (void*) s1, *d2 = (void*) s2;
544 unsigned long size1, size2;
545 size1 = d1->dict ? dictSize(d1->dict) : 0;
546 size2 = d2->dict ? dictSize(d2->dict) : 0;
547 return size1 - size2;
548 }
549
550 #define REDIS_AGGR_SUM 1
551 #define REDIS_AGGR_MIN 2
552 #define REDIS_AGGR_MAX 3
553 #define zunionInterDictValue(_e) (dictGetEntryVal(_e) == NULL ? 1.0 : *(double*)dictGetEntryVal(_e))
554
555 inline static void zunionInterAggregate(double *target, double val, int aggregate) {
556 if (aggregate == REDIS_AGGR_SUM) {
557 *target = *target + val;
558 /* The result of adding two doubles is NaN when one variable
559 * is +inf and the other is -inf. When these numbers are added,
560 * we maintain the convention of the result being 0.0. */
561 if (isnan(*target)) *target = 0.0;
562 } else if (aggregate == REDIS_AGGR_MIN) {
563 *target = val < *target ? val : *target;
564 } else if (aggregate == REDIS_AGGR_MAX) {
565 *target = val > *target ? val : *target;
566 } else {
567 /* safety net */
568 redisPanic("Unknown ZUNION/INTER aggregate type");
569 }
570 }
571
572 void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
573 int i, j, setnum;
574 int aggregate = REDIS_AGGR_SUM;
575 zsetopsrc *src;
576 robj *dstobj;
577 zset *dstzset;
578 zskiplistNode *znode;
579 dictIterator *di;
580 dictEntry *de;
581 int touched = 0;
582
583 /* expect setnum input keys to be given */
584 setnum = atoi(c->argv[2]->ptr);
585 if (setnum < 1) {
586 addReplyError(c,
587 "at least 1 input key is needed for ZUNIONSTORE/ZINTERSTORE");
588 return;
589 }
590
591 /* test if the expected number of keys would overflow */
592 if (3+setnum > c->argc) {
593 addReply(c,shared.syntaxerr);
594 return;
595 }
596
597 /* read keys to be used for input */
598 src = zmalloc(sizeof(zsetopsrc) * setnum);
599 for (i = 0, j = 3; i < setnum; i++, j++) {
600 robj *obj = lookupKeyWrite(c->db,c->argv[j]);
601 if (!obj) {
602 src[i].dict = NULL;
603 } else {
604 if (obj->type == REDIS_ZSET) {
605 src[i].dict = ((zset*)obj->ptr)->dict;
606 } else if (obj->type == REDIS_SET) {
607 src[i].dict = (obj->ptr);
608 } else {
609 zfree(src);
610 addReply(c,shared.wrongtypeerr);
611 return;
612 }
613 }
614
615 /* default all weights to 1 */
616 src[i].weight = 1.0;
617 }
618
619 /* parse optional extra arguments */
620 if (j < c->argc) {
621 int remaining = c->argc - j;
622
623 while (remaining) {
624 if (remaining >= (setnum + 1) && !strcasecmp(c->argv[j]->ptr,"weights")) {
625 j++; remaining--;
626 for (i = 0; i < setnum; i++, j++, remaining--) {
627 if (getDoubleFromObjectOrReply(c,c->argv[j],&src[i].weight,
628 "weight value is not a double") != REDIS_OK)
629 {
630 zfree(src);
631 return;
632 }
633 }
634 } else if (remaining >= 2 && !strcasecmp(c->argv[j]->ptr,"aggregate")) {
635 j++; remaining--;
636 if (!strcasecmp(c->argv[j]->ptr,"sum")) {
637 aggregate = REDIS_AGGR_SUM;
638 } else if (!strcasecmp(c->argv[j]->ptr,"min")) {
639 aggregate = REDIS_AGGR_MIN;
640 } else if (!strcasecmp(c->argv[j]->ptr,"max")) {
641 aggregate = REDIS_AGGR_MAX;
642 } else {
643 zfree(src);
644 addReply(c,shared.syntaxerr);
645 return;
646 }
647 j++; remaining--;
648 } else {
649 zfree(src);
650 addReply(c,shared.syntaxerr);
651 return;
652 }
653 }
654 }
655
656 /* sort sets from the smallest to largest, this will improve our
657 * algorithm's performance */
658 qsort(src,setnum,sizeof(zsetopsrc),qsortCompareZsetopsrcByCardinality);
659
660 dstobj = createZsetObject();
661 dstzset = dstobj->ptr;
662
663 if (op == REDIS_OP_INTER) {
664 /* skip going over all entries if the smallest zset is NULL or empty */
665 if (src[0].dict && dictSize(src[0].dict) > 0) {
666 /* precondition: as src[0].dict is non-empty and the zsets are ordered
667 * from small to large, all src[i > 0].dict are non-empty too */
668 di = dictGetIterator(src[0].dict);
669 while((de = dictNext(di)) != NULL) {
670 double score, value;
671
672 score = src[0].weight * zunionInterDictValue(de);
673 for (j = 1; j < setnum; j++) {
674 dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
675 if (other) {
676 value = src[j].weight * zunionInterDictValue(other);
677 zunionInterAggregate(&score,value,aggregate);
678 } else {
679 break;
680 }
681 }
682
683 /* Only continue when present in every source dict. */
684 if (j == setnum) {
685 robj *o = dictGetEntryKey(de);
686 znode = zslInsert(dstzset->zsl,score,o);
687 incrRefCount(o); /* added to skiplist */
688 dictAdd(dstzset->dict,o,&znode->score);
689 incrRefCount(o); /* added to dictionary */
690 }
691 }
692 dictReleaseIterator(di);
693 }
694 } else if (op == REDIS_OP_UNION) {
695 for (i = 0; i < setnum; i++) {
696 if (!src[i].dict) continue;
697
698 di = dictGetIterator(src[i].dict);
699 while((de = dictNext(di)) != NULL) {
700 double score, value;
701
702 /* skip key when already processed */
703 if (dictFind(dstzset->dict,dictGetEntryKey(de)) != NULL)
704 continue;
705
706 /* initialize score */
707 score = src[i].weight * zunionInterDictValue(de);
708
709 /* because the zsets are sorted by size, its only possible
710 * for sets at larger indices to hold this entry */
711 for (j = (i+1); j < setnum; j++) {
712 dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
713 if (other) {
714 value = src[j].weight * zunionInterDictValue(other);
715 zunionInterAggregate(&score,value,aggregate);
716 }
717 }
718
719 robj *o = dictGetEntryKey(de);
720 znode = zslInsert(dstzset->zsl,score,o);
721 incrRefCount(o); /* added to skiplist */
722 dictAdd(dstzset->dict,o,&znode->score);
723 incrRefCount(o); /* added to dictionary */
724 }
725 dictReleaseIterator(di);
726 }
727 } else {
728 /* unknown operator */
729 redisAssert(op == REDIS_OP_INTER || op == REDIS_OP_UNION);
730 }
731
732 if (dbDelete(c->db,dstkey)) {
733 touchWatchedKey(c->db,dstkey);
734 touched = 1;
735 server.dirty++;
736 }
737 if (dstzset->zsl->length) {
738 dbAdd(c->db,dstkey,dstobj);
739 addReplyLongLong(c, dstzset->zsl->length);
740 if (!touched) touchWatchedKey(c->db,dstkey);
741 server.dirty++;
742 } else {
743 decrRefCount(dstobj);
744 addReply(c, shared.czero);
745 }
746 zfree(src);
747 }
748
749 void zunionstoreCommand(redisClient *c) {
750 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_UNION);
751 }
752
753 void zinterstoreCommand(redisClient *c) {
754 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_INTER);
755 }
756
757 void zrangeGenericCommand(redisClient *c, int reverse) {
758 robj *o;
759 long start;
760 long end;
761 int withscores = 0;
762 int llen;
763 int rangelen, j;
764 zset *zsetobj;
765 zskiplist *zsl;
766 zskiplistNode *ln;
767 robj *ele;
768
769 if ((getLongFromObjectOrReply(c, c->argv[2], &start, NULL) != REDIS_OK) ||
770 (getLongFromObjectOrReply(c, c->argv[3], &end, NULL) != REDIS_OK)) return;
771
772 if (c->argc == 5 && !strcasecmp(c->argv[4]->ptr,"withscores")) {
773 withscores = 1;
774 } else if (c->argc >= 5) {
775 addReply(c,shared.syntaxerr);
776 return;
777 }
778
779 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.emptymultibulk)) == NULL
780 || checkType(c,o,REDIS_ZSET)) return;
781 zsetobj = o->ptr;
782 zsl = zsetobj->zsl;
783 llen = zsl->length;
784
785 /* convert negative indexes */
786 if (start < 0) start = llen+start;
787 if (end < 0) end = llen+end;
788 if (start < 0) start = 0;
789
790 /* Invariant: start >= 0, so this test will be true when end < 0.
791 * The range is empty when start > end or start >= length. */
792 if (start > end || start >= llen) {
793 addReply(c,shared.emptymultibulk);
794 return;
795 }
796 if (end >= llen) end = llen-1;
797 rangelen = (end-start)+1;
798
799 /* check if starting point is trivial, before searching
800 * the element in log(N) time */
801 if (reverse) {
802 ln = start == 0 ? zsl->tail : zslistTypeGetElementByRank(zsl, llen-start);
803 } else {
804 ln = start == 0 ?
805 zsl->header->level[0].forward : zslistTypeGetElementByRank(zsl, start+1);
806 }
807
808 /* Return the result in form of a multi-bulk reply */
809 addReplyMultiBulkLen(c,withscores ? (rangelen*2) : rangelen);
810 for (j = 0; j < rangelen; j++) {
811 ele = ln->obj;
812 addReplyBulk(c,ele);
813 if (withscores)
814 addReplyDouble(c,ln->score);
815 ln = reverse ? ln->backward : ln->level[0].forward;
816 }
817 }
818
819 void zrangeCommand(redisClient *c) {
820 zrangeGenericCommand(c,0);
821 }
822
823 void zrevrangeCommand(redisClient *c) {
824 zrangeGenericCommand(c,1);
825 }
826
827 /* This command implements ZRANGEBYSCORE, ZREVRANGEBYSCORE and ZCOUNT.
828 * If "justcount", only the number of elements in the range is returned. */
829 void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) {
830 zrangespec range;
831 robj *o, *emptyreply;
832 zset *zsetobj;
833 zskiplist *zsl;
834 zskiplistNode *ln;
835 int offset = 0, limit = -1;
836 int withscores = 0;
837 unsigned long rangelen = 0;
838 void *replylen = NULL;
839
840 /* Parse the range arguments. */
841 zslParseRange(c->argv[2],c->argv[3],&range);
842
843 /* Parse optional extra arguments. Note that ZCOUNT will exactly have
844 * 4 arguments, so we'll never enter the following code path. */
845 if (c->argc > 4) {
846 int remaining = c->argc - 4;
847 int pos = 4;
848
849 while (remaining) {
850 if (remaining >= 1 && !strcasecmp(c->argv[pos]->ptr,"withscores")) {
851 pos++; remaining--;
852 withscores = 1;
853 } else if (remaining >= 3 && !strcasecmp(c->argv[pos]->ptr,"limit")) {
854 offset = atoi(c->argv[pos+1]->ptr);
855 limit = atoi(c->argv[pos+2]->ptr);
856 pos += 3; remaining -= 3;
857 } else {
858 addReply(c,shared.syntaxerr);
859 return;
860 }
861 }
862 }
863
864 /* Ok, lookup the key and get the range */
865 emptyreply = justcount ? shared.czero : shared.emptymultibulk;
866 if ((o = lookupKeyReadOrReply(c,c->argv[1],emptyreply)) == NULL ||
867 checkType(c,o,REDIS_ZSET)) return;
868 zsetobj = o->ptr;
869 zsl = zsetobj->zsl;
870
871 /* If reversed, assume the elements are sorted from high to low score. */
872 ln = zslFirstWithScore(zsl,range.min);
873 if (reverse) {
874 /* If range.min is out of range, ln will be NULL and we need to use
875 * the tail of the skiplist as first node of the range. */
876 if (ln == NULL) ln = zsl->tail;
877
878 /* zslFirstWithScore returns the first element with where with
879 * score >= range.min, so backtrack to make sure the element we use
880 * here has score <= range.min. */
881 while (ln && ln->score > range.min) ln = ln->backward;
882
883 /* Move to the right element according to the range spec. */
884 if (range.minex) {
885 /* Find last element with score < range.min */
886 while (ln && ln->score == range.min) ln = ln->backward;
887 } else {
888 /* Find last element with score <= range.min */
889 while (ln && ln->level[0].forward &&
890 ln->level[0].forward->score == range.min)
891 ln = ln->level[0].forward;
892 }
893 } else {
894 if (range.minex) {
895 /* Find first element with score > range.min */
896 while (ln && ln->score == range.min) ln = ln->level[0].forward;
897 }
898 }
899
900 /* No "first" element in the specified interval. */
901 if (ln == NULL) {
902 addReply(c,emptyreply);
903 return;
904 }
905
906 /* We don't know in advance how many matching elements there
907 * are in the list, so we push this object that will represent
908 * the multi-bulk length in the output buffer, and will "fix"
909 * it later */
910 if (!justcount)
911 replylen = addDeferredMultiBulkLength(c);
912
913 /* If there is an offset, just traverse the number of elements without
914 * checking the score because that is done in the next loop. */
915 while(ln && offset--) {
916 if (reverse)
917 ln = ln->backward;
918 else
919 ln = ln->level[0].forward;
920 }
921
922 while (ln && limit--) {
923 /* Check if this this element is in range. */
924 if (reverse) {
925 if (range.maxex) {
926 /* Element should have score > range.max */
927 if (ln->score <= range.max) break;
928 } else {
929 /* Element should have score >= range.max */
930 if (ln->score < range.max) break;
931 }
932 } else {
933 if (range.maxex) {
934 /* Element should have score < range.max */
935 if (ln->score >= range.max) break;
936 } else {
937 /* Element should have score <= range.max */
938 if (ln->score > range.max) break;
939 }
940 }
941
942 /* Do our magic */
943 rangelen++;
944 if (!justcount) {
945 addReplyBulk(c,ln->obj);
946 if (withscores)
947 addReplyDouble(c,ln->score);
948 }
949
950 if (reverse)
951 ln = ln->backward;
952 else
953 ln = ln->level[0].forward;
954 }
955
956 if (justcount) {
957 addReplyLongLong(c,(long)rangelen);
958 } else {
959 setDeferredMultiBulkLength(c,replylen,
960 withscores ? (rangelen*2) : rangelen);
961 }
962 }
963
964 void zrangebyscoreCommand(redisClient *c) {
965 genericZrangebyscoreCommand(c,0,0);
966 }
967
968 void zrevrangebyscoreCommand(redisClient *c) {
969 genericZrangebyscoreCommand(c,1,0);
970 }
971
972 void zcountCommand(redisClient *c) {
973 genericZrangebyscoreCommand(c,0,1);
974 }
975
976 void zcardCommand(redisClient *c) {
977 robj *o;
978 zset *zs;
979
980 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
981 checkType(c,o,REDIS_ZSET)) return;
982
983 zs = o->ptr;
984 addReplyLongLong(c,zs->zsl->length);
985 }
986
987 void zscoreCommand(redisClient *c) {
988 robj *o;
989 zset *zs;
990 dictEntry *de;
991
992 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
993 checkType(c,o,REDIS_ZSET)) return;
994
995 zs = o->ptr;
996 de = dictFind(zs->dict,c->argv[2]);
997 if (!de) {
998 addReply(c,shared.nullbulk);
999 } else {
1000 double *score = dictGetEntryVal(de);
1001
1002 addReplyDouble(c,*score);
1003 }
1004 }
1005
1006 void zrankGenericCommand(redisClient *c, int reverse) {
1007 robj *o;
1008 zset *zs;
1009 zskiplist *zsl;
1010 dictEntry *de;
1011 unsigned long rank;
1012 double *score;
1013
1014 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
1015 checkType(c,o,REDIS_ZSET)) return;
1016
1017 zs = o->ptr;
1018 zsl = zs->zsl;
1019 de = dictFind(zs->dict,c->argv[2]);
1020 if (!de) {
1021 addReply(c,shared.nullbulk);
1022 return;
1023 }
1024
1025 score = dictGetEntryVal(de);
1026 rank = zslistTypeGetRank(zsl, *score, c->argv[2]);
1027 if (rank) {
1028 if (reverse) {
1029 addReplyLongLong(c, zsl->length - rank);
1030 } else {
1031 addReplyLongLong(c, rank-1);
1032 }
1033 } else {
1034 addReply(c,shared.nullbulk);
1035 }
1036 }
1037
1038 void zrankCommand(redisClient *c) {
1039 zrankGenericCommand(c, 0);
1040 }
1041
1042 void zrevrankCommand(redisClient *c) {
1043 zrankGenericCommand(c, 1);
1044 }