]> git.saurik.com Git - redis.git/blob - src/t_zset.c
790fb573c37b1ba0559d0d6f8927b715fea49a02
[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 static int zslValueGteMin(double value, zrangespec *spec) {
184 return spec->minex ? (value > spec->min) : (value >= spec->min);
185 }
186
187 static int zslValueLteMax(double value, zrangespec *spec) {
188 return spec->maxex ? (value < spec->max) : (value <= spec->max);
189 }
190
191 static int zslValueInRange(double value, zrangespec *spec) {
192 return zslValueGteMin(value,spec) && zslValueLteMax(value,spec);
193 }
194
195 /* Returns if there is a part of the zset is in range. */
196 int zslIsInRange(zskiplist *zsl, zrangespec *range) {
197 zskiplistNode *x;
198
199 /* Test for ranges that will always be empty. */
200 if (range->min > range->max ||
201 (range->min == range->max && (range->minex || range->maxex)))
202 return 0;
203 x = zsl->tail;
204 if (x == NULL || !zslValueGteMin(x->score,range))
205 return 0;
206 x = zsl->header->level[0].forward;
207 if (x == NULL || !zslValueLteMax(x->score,range))
208 return 0;
209 return 1;
210 }
211
212 /* Find the first node that is contained in the specified range.
213 * Returns NULL when no element is contained in the range. */
214 zskiplistNode *zslFirstInRange(zskiplist *zsl, zrangespec range) {
215 zskiplistNode *x;
216 int i;
217
218 /* If everything is out of range, return early. */
219 if (!zslIsInRange(zsl,&range)) return NULL;
220
221 x = zsl->header;
222 for (i = zsl->level-1; i >= 0; i--) {
223 /* Go forward while *OUT* of range. */
224 while (x->level[i].forward &&
225 !zslValueGteMin(x->level[i].forward->score,&range))
226 x = x->level[i].forward;
227 }
228
229 /* The tail is in range, so the previous block should always return a
230 * node that is non-NULL and the last one to be out of range. */
231 x = x->level[0].forward;
232 redisAssert(x != NULL && zslValueInRange(x->score,&range));
233 return x;
234 }
235
236 /* Find the last node that is contained in the specified range.
237 * Returns NULL when no element is contained in the range. */
238 zskiplistNode *zslLastInRange(zskiplist *zsl, zrangespec range) {
239 zskiplistNode *x;
240 int i;
241
242 /* If everything is out of range, return early. */
243 if (!zslIsInRange(zsl,&range)) return NULL;
244
245 x = zsl->header;
246 for (i = zsl->level-1; i >= 0; i--) {
247 /* Go forward while *IN* range. */
248 while (x->level[i].forward &&
249 zslValueLteMax(x->level[i].forward->score,&range))
250 x = x->level[i].forward;
251 }
252
253 /* The header is in range, so the previous block should always return a
254 * node that is non-NULL and in range. */
255 redisAssert(x != NULL && zslValueInRange(x->score,&range));
256 return x;
257 }
258
259 /* Delete all the elements with score between min and max from the skiplist.
260 * Min and mx are inclusive, so a score >= min || score <= max is deleted.
261 * Note that this function takes the reference to the hash table view of the
262 * sorted set, in order to remove the elements from the hash table too. */
263 unsigned long zslDeleteRangeByScore(zskiplist *zsl, zrangespec range, dict *dict) {
264 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
265 unsigned long removed = 0;
266 int i;
267
268 x = zsl->header;
269 for (i = zsl->level-1; i >= 0; i--) {
270 while (x->level[i].forward && (range.minex ?
271 x->level[i].forward->score <= range.min :
272 x->level[i].forward->score < range.min))
273 x = x->level[i].forward;
274 update[i] = x;
275 }
276
277 /* Current node is the last with score < or <= min. */
278 x = x->level[0].forward;
279
280 /* Delete nodes while in range. */
281 while (x && (range.maxex ? x->score < range.max : x->score <= range.max)) {
282 zskiplistNode *next = x->level[0].forward;
283 zslDeleteNode(zsl,x,update);
284 dictDelete(dict,x->obj);
285 zslFreeNode(x);
286 removed++;
287 x = next;
288 }
289 return removed;
290 }
291
292 /* Delete all the elements with rank between start and end from the skiplist.
293 * Start and end are inclusive. Note that start and end need to be 1-based */
294 unsigned long zslDeleteRangeByRank(zskiplist *zsl, unsigned int start, unsigned int end, dict *dict) {
295 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x;
296 unsigned long traversed = 0, removed = 0;
297 int i;
298
299 x = zsl->header;
300 for (i = zsl->level-1; i >= 0; i--) {
301 while (x->level[i].forward && (traversed + x->level[i].span) < start) {
302 traversed += x->level[i].span;
303 x = x->level[i].forward;
304 }
305 update[i] = x;
306 }
307
308 traversed++;
309 x = x->level[0].forward;
310 while (x && traversed <= end) {
311 zskiplistNode *next = x->level[0].forward;
312 zslDeleteNode(zsl,x,update);
313 dictDelete(dict,x->obj);
314 zslFreeNode(x);
315 removed++;
316 traversed++;
317 x = next;
318 }
319 return removed;
320 }
321
322 /* Find the rank for an element by both score and key.
323 * Returns 0 when the element cannot be found, rank otherwise.
324 * Note that the rank is 1-based due to the span of zsl->header to the
325 * first element. */
326 unsigned long zslGetRank(zskiplist *zsl, double score, robj *o) {
327 zskiplistNode *x;
328 unsigned long rank = 0;
329 int i;
330
331 x = zsl->header;
332 for (i = zsl->level-1; i >= 0; i--) {
333 while (x->level[i].forward &&
334 (x->level[i].forward->score < score ||
335 (x->level[i].forward->score == score &&
336 compareStringObjects(x->level[i].forward->obj,o) <= 0))) {
337 rank += x->level[i].span;
338 x = x->level[i].forward;
339 }
340
341 /* x might be equal to zsl->header, so test if obj is non-NULL */
342 if (x->obj && equalStringObjects(x->obj,o)) {
343 return rank;
344 }
345 }
346 return 0;
347 }
348
349 /* Finds an element by its rank. The rank argument needs to be 1-based. */
350 zskiplistNode* zslGetElementByRank(zskiplist *zsl, unsigned long rank) {
351 zskiplistNode *x;
352 unsigned long traversed = 0;
353 int i;
354
355 x = zsl->header;
356 for (i = zsl->level-1; i >= 0; i--) {
357 while (x->level[i].forward && (traversed + x->level[i].span) <= rank)
358 {
359 traversed += x->level[i].span;
360 x = x->level[i].forward;
361 }
362 if (traversed == rank) {
363 return x;
364 }
365 }
366 return NULL;
367 }
368
369 /* Populate the rangespec according to the objects min and max. */
370 static int zslParseRange(robj *min, robj *max, zrangespec *spec) {
371 char *eptr;
372 spec->minex = spec->maxex = 0;
373
374 /* Parse the min-max interval. If one of the values is prefixed
375 * by the "(" character, it's considered "open". For instance
376 * ZRANGEBYSCORE zset (1.5 (2.5 will match min < x < max
377 * ZRANGEBYSCORE zset 1.5 2.5 will instead match min <= x <= max */
378 if (min->encoding == REDIS_ENCODING_INT) {
379 spec->min = (long)min->ptr;
380 } else {
381 if (((char*)min->ptr)[0] == '(') {
382 spec->min = strtod((char*)min->ptr+1,&eptr);
383 if (eptr[0] != '\0' || isnan(spec->min)) return REDIS_ERR;
384 spec->minex = 1;
385 } else {
386 spec->min = strtod((char*)min->ptr,&eptr);
387 if (eptr[0] != '\0' || isnan(spec->min)) return REDIS_ERR;
388 }
389 }
390 if (max->encoding == REDIS_ENCODING_INT) {
391 spec->max = (long)max->ptr;
392 } else {
393 if (((char*)max->ptr)[0] == '(') {
394 spec->max = strtod((char*)max->ptr+1,&eptr);
395 if (eptr[0] != '\0' || isnan(spec->max)) return REDIS_ERR;
396 spec->maxex = 1;
397 } else {
398 spec->max = strtod((char*)max->ptr,&eptr);
399 if (eptr[0] != '\0' || isnan(spec->max)) return REDIS_ERR;
400 }
401 }
402
403 return REDIS_OK;
404 }
405
406 /*-----------------------------------------------------------------------------
407 * Ziplist-backed sorted set API
408 *----------------------------------------------------------------------------*/
409
410 double zzlGetScore(unsigned char *sptr) {
411 unsigned char *vstr;
412 unsigned int vlen;
413 long long vlong;
414 char buf[128];
415 double score;
416
417 redisAssert(sptr != NULL);
418 redisAssert(ziplistGet(sptr,&vstr,&vlen,&vlong));
419
420 if (vstr) {
421 memcpy(buf,vstr,vlen);
422 buf[vlen] = '\0';
423 score = strtod(buf,NULL);
424 } else {
425 score = vlong;
426 }
427
428 return score;
429 }
430
431 /* Compare element in sorted set with given element. */
432 int zzlCompareElements(unsigned char *eptr, unsigned char *cstr, unsigned int clen) {
433 unsigned char *vstr;
434 unsigned int vlen;
435 long long vlong;
436 unsigned char vbuf[32];
437 int minlen, cmp;
438
439 redisAssert(ziplistGet(eptr,&vstr,&vlen,&vlong));
440 if (vstr == NULL) {
441 /* Store string representation of long long in buf. */
442 vlen = ll2string((char*)vbuf,sizeof(vbuf),vlong);
443 vstr = vbuf;
444 }
445
446 minlen = (vlen < clen) ? vlen : clen;
447 cmp = memcmp(vstr,cstr,minlen);
448 if (cmp == 0) return vlen-clen;
449 return cmp;
450 }
451
452 unsigned int *zzlLength(robj *zobj) {
453 unsigned char *zl = zobj->ptr;
454 return ziplistLen(zl)/2;
455 }
456
457 unsigned char *zzlFind(robj *zobj, robj *ele, double *score) {
458 unsigned char *zl = zobj->ptr;
459 unsigned char *eptr = ziplistIndex(zl,0), *sptr;
460
461 ele = getDecodedObject(ele);
462 while (eptr != NULL) {
463 sptr = ziplistNext(zl,eptr);
464 redisAssert(sptr != NULL);
465
466 if (ziplistCompare(eptr,ele->ptr,sdslen(ele->ptr))) {
467 /* Matching element, pull out score. */
468 if (score != NULL) *score = zzlGetScore(sptr);
469 decrRefCount(ele);
470 return eptr;
471 }
472
473 /* Move to next element. */
474 eptr = ziplistNext(zl,sptr);
475 }
476
477 decrRefCount(ele);
478 return NULL;
479 }
480
481 /* Delete (element,score) pair from ziplist. Use local copy of eptr because we
482 * don't want to modify the one given as argument. */
483 int zzlDelete(robj *zobj, unsigned char *eptr) {
484 unsigned char *zl = zobj->ptr;
485 unsigned char *p = eptr;
486
487 /* TODO: add function to ziplist API to delete N elements from offset. */
488 zl = ziplistDelete(zl,&p);
489 zl = ziplistDelete(zl,&p);
490 zobj->ptr = zl;
491 return REDIS_OK;
492 }
493
494 int zzlInsertAt(robj *zobj, robj *ele, double score, unsigned char *eptr) {
495 unsigned char *zl = zobj->ptr;
496 unsigned char *sptr;
497 char scorebuf[128];
498 int scorelen;
499 int offset;
500
501 redisAssert(ele->encoding == REDIS_ENCODING_RAW);
502 scorelen = d2string(scorebuf,sizeof(scorebuf),score);
503 if (eptr == NULL) {
504 zl = ziplistPush(zl,ele->ptr,sdslen(ele->ptr),ZIPLIST_TAIL);
505 zl = ziplistPush(zl,(unsigned char*)scorebuf,scorelen,ZIPLIST_TAIL);
506 } else {
507 /* Keep offset relative to zl, as it might be re-allocated. */
508 offset = eptr-zl;
509 zl = ziplistInsert(zl,eptr,ele->ptr,sdslen(ele->ptr));
510 eptr = zl+offset;
511
512 /* Insert score after the element. */
513 redisAssert((sptr = ziplistNext(zl,eptr)) != NULL);
514 zl = ziplistInsert(zl,sptr,(unsigned char*)scorebuf,scorelen);
515 }
516
517 zobj->ptr = zl;
518 return REDIS_OK;
519 }
520
521 /* Insert (element,score) pair in ziplist. This function assumes the element is
522 * not yet present in the list. */
523 int zzlInsert(robj *zobj, robj *ele, double score) {
524 unsigned char *zl = zobj->ptr;
525 unsigned char *eptr = ziplistIndex(zl,0), *sptr;
526 double s;
527 int insert = 0;
528
529 ele = getDecodedObject(ele);
530 while (eptr != NULL) {
531 sptr = ziplistNext(zl,eptr);
532 redisAssert(sptr != NULL);
533 s = zzlGetScore(sptr);
534
535 if (s > score) {
536 /* First element with score larger than score for element to be
537 * inserted. This means we should take its spot in the list to
538 * maintain ordering. */
539 insert = 1;
540 } else if (s == score) {
541 /* Ensure lexicographical ordering for elements. */
542 if (zzlCompareElements(eptr,ele->ptr,sdslen(ele->ptr)) < 0)
543 insert = 1;
544 }
545
546 if (insert) {
547 zzlInsertAt(zobj,ele,score,eptr);
548 break;
549 }
550
551 /* Move to next element. */
552 eptr = ziplistNext(zl,sptr);
553 }
554
555 /* Push on tail of list when it was not yet inserted. */
556 if (!insert)
557 zzlInsertAt(zobj,ele,score,eptr);
558
559 decrRefCount(ele);
560 return REDIS_OK;
561 }
562
563 /*-----------------------------------------------------------------------------
564 * Sorted set commands
565 *----------------------------------------------------------------------------*/
566
567 /* This generic command implements both ZADD and ZINCRBY. */
568 void zaddGenericCommand(redisClient *c, int incr) {
569 static char *nanerr = "resulting score is not a number (NaN)";
570 robj *key = c->argv[1];
571 robj *ele;
572 robj *zobj;
573 robj *curobj;
574 double score, curscore = 0.0;
575
576 if (getDoubleFromObjectOrReply(c,c->argv[2],&score,NULL) != REDIS_OK)
577 return;
578
579 zobj = lookupKeyWrite(c->db,key);
580 if (zobj == NULL) {
581 zobj = createZsetZiplistObject();
582 dbAdd(c->db,key,zobj);
583 } else {
584 if (zobj->type != REDIS_ZSET) {
585 addReply(c,shared.wrongtypeerr);
586 return;
587 }
588 }
589
590 if (zobj->encoding == REDIS_ENCODING_ZIPLIST) {
591 unsigned char *eptr;
592
593 /* Prefer non-encoded element when dealing with ziplists. */
594 ele = c->argv[3];
595 if ((eptr = zzlFind(zobj,ele,&curscore)) != NULL) {
596 if (incr) {
597 score += curscore;
598 if (isnan(score)) {
599 addReplyError(c,nanerr);
600 /* Don't need to check if the sorted set is empty, because
601 * we know it has at least one element. */
602 return;
603 }
604 }
605
606 /* Remove and re-insert when score changed. */
607 if (score != curscore) {
608 redisAssert(zzlDelete(zobj,eptr) == REDIS_OK);
609 redisAssert(zzlInsert(zobj,ele,score) == REDIS_OK);
610
611 signalModifiedKey(c->db,key);
612 server.dirty++;
613 }
614
615 if (incr) /* ZINCRBY */
616 addReplyDouble(c,score);
617 else /* ZADD */
618 addReply(c,shared.czero);
619 } else {
620 redisAssert(zzlInsert(zobj,ele,score) == REDIS_OK);
621
622 signalModifiedKey(c->db,key);
623 server.dirty++;
624
625 if (incr) /* ZINCRBY */
626 addReplyDouble(c,score);
627 else /* ZADD */
628 addReply(c,shared.cone);
629 }
630 } else if (zobj->encoding == REDIS_ENCODING_RAW) {
631 zset *zs = zobj->ptr;
632 zskiplistNode *znode;
633 dictEntry *de;
634
635 ele = c->argv[3] = tryObjectEncoding(c->argv[3]);
636 de = dictFind(zs->dict,ele);
637 if (de != NULL) {
638 curobj = dictGetEntryKey(de);
639 curscore = *(double*)dictGetEntryVal(de);
640
641 if (incr) {
642 score += curscore;
643 if (isnan(score)) {
644 addReplyError(c,nanerr);
645 /* Don't need to check if the sorted set is empty, because
646 * we know it has at least one element. */
647 return;
648 }
649 }
650
651 /* Remove and re-insert when score changed. We can safely delete
652 * the key object from the skiplist, since the dictionary still has
653 * a reference to it. */
654 if (score != curscore) {
655 redisAssert(zslDelete(zs->zsl,curscore,curobj));
656 znode = zslInsert(zs->zsl,score,curobj);
657 incrRefCount(curobj); /* Re-inserted in skiplist. */
658 dictGetEntryVal(de) = &znode->score; /* Update score ptr. */
659
660 signalModifiedKey(c->db,key);
661 server.dirty++;
662 }
663
664 if (incr) /* ZINCRBY */
665 addReplyDouble(c,score);
666 else /* ZADD */
667 addReply(c,shared.czero);
668 } else {
669 znode = zslInsert(zs->zsl,score,ele);
670 incrRefCount(ele); /* Inserted in skiplist. */
671 redisAssert(dictAdd(zs->dict,ele,&znode->score) == DICT_OK);
672 incrRefCount(ele); /* Added to dictionary. */
673
674 signalModifiedKey(c->db,key);
675 server.dirty++;
676
677 if (incr) /* ZINCRBY */
678 addReplyDouble(c,score);
679 else /* ZADD */
680 addReply(c,shared.cone);
681 }
682 } else {
683 redisPanic("Unknown sorted set encoding");
684 }
685 }
686
687 void zaddCommand(redisClient *c) {
688 zaddGenericCommand(c,0);
689 }
690
691 void zincrbyCommand(redisClient *c) {
692 zaddGenericCommand(c,1);
693 }
694
695 void zremCommand(redisClient *c) {
696 robj *key = c->argv[1];
697 robj *ele = c->argv[2];
698 robj *zobj;
699
700 if ((zobj = lookupKeyWriteOrReply(c,key,shared.czero)) == NULL ||
701 checkType(c,zobj,REDIS_ZSET)) return;
702
703 if (zobj->encoding == REDIS_ENCODING_ZIPLIST) {
704 unsigned char *eptr;
705
706 if ((eptr = zzlFind(zobj,ele,NULL)) != NULL) {
707 redisAssert(zzlDelete(zobj,eptr) == REDIS_OK);
708 if (zzlLength(zobj) == 0) dbDelete(c->db,key);
709 } else {
710 addReply(c,shared.czero);
711 return;
712 }
713 } else if (zobj->encoding == REDIS_ENCODING_RAW) {
714 zset *zs = zobj->ptr;
715 dictEntry *de;
716 double score;
717
718 de = dictFind(zs->dict,ele);
719 if (de != NULL) {
720 /* Delete from the skiplist */
721 score = *(double*)dictGetEntryVal(de);
722 redisAssert(zslDelete(zs->zsl,score,ele));
723
724 /* Delete from the hash table */
725 dictDelete(zs->dict,ele);
726 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
727 if (dictSize(zs->dict) == 0) dbDelete(c->db,key);
728 } else {
729 addReply(c,shared.czero);
730 return;
731 }
732 } else {
733 redisPanic("Unknown sorted set encoding");
734 }
735
736 signalModifiedKey(c->db,key);
737 server.dirty++;
738 addReply(c,shared.cone);
739 }
740
741 void zremrangebyscoreCommand(redisClient *c) {
742 zrangespec range;
743 long deleted;
744 robj *o;
745 zset *zs;
746
747 /* Parse the range arguments. */
748 if (zslParseRange(c->argv[2],c->argv[3],&range) != REDIS_OK) {
749 addReplyError(c,"min or max is not a double");
750 return;
751 }
752
753 if ((o = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
754 checkType(c,o,REDIS_ZSET)) return;
755
756 zs = o->ptr;
757 deleted = zslDeleteRangeByScore(zs->zsl,range,zs->dict);
758 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
759 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
760 if (deleted) signalModifiedKey(c->db,c->argv[1]);
761 server.dirty += deleted;
762 addReplyLongLong(c,deleted);
763 }
764
765 void zremrangebyrankCommand(redisClient *c) {
766 long start;
767 long end;
768 int llen;
769 long deleted;
770 robj *zsetobj;
771 zset *zs;
772
773 if ((getLongFromObjectOrReply(c, c->argv[2], &start, NULL) != REDIS_OK) ||
774 (getLongFromObjectOrReply(c, c->argv[3], &end, NULL) != REDIS_OK)) return;
775
776 if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
777 checkType(c,zsetobj,REDIS_ZSET)) return;
778 zs = zsetobj->ptr;
779 llen = zs->zsl->length;
780
781 /* convert negative indexes */
782 if (start < 0) start = llen+start;
783 if (end < 0) end = llen+end;
784 if (start < 0) start = 0;
785
786 /* Invariant: start >= 0, so this test will be true when end < 0.
787 * The range is empty when start > end or start >= length. */
788 if (start > end || start >= llen) {
789 addReply(c,shared.czero);
790 return;
791 }
792 if (end >= llen) end = llen-1;
793
794 /* increment start and end because zsl*Rank functions
795 * use 1-based rank */
796 deleted = zslDeleteRangeByRank(zs->zsl,start+1,end+1,zs->dict);
797 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
798 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
799 if (deleted) signalModifiedKey(c->db,c->argv[1]);
800 server.dirty += deleted;
801 addReplyLongLong(c, deleted);
802 }
803
804 typedef struct {
805 dict *dict;
806 double weight;
807 } zsetopsrc;
808
809 int qsortCompareZsetopsrcByCardinality(const void *s1, const void *s2) {
810 zsetopsrc *d1 = (void*) s1, *d2 = (void*) s2;
811 unsigned long size1, size2;
812 size1 = d1->dict ? dictSize(d1->dict) : 0;
813 size2 = d2->dict ? dictSize(d2->dict) : 0;
814 return size1 - size2;
815 }
816
817 #define REDIS_AGGR_SUM 1
818 #define REDIS_AGGR_MIN 2
819 #define REDIS_AGGR_MAX 3
820 #define zunionInterDictValue(_e) (dictGetEntryVal(_e) == NULL ? 1.0 : *(double*)dictGetEntryVal(_e))
821
822 inline static void zunionInterAggregate(double *target, double val, int aggregate) {
823 if (aggregate == REDIS_AGGR_SUM) {
824 *target = *target + val;
825 /* The result of adding two doubles is NaN when one variable
826 * is +inf and the other is -inf. When these numbers are added,
827 * we maintain the convention of the result being 0.0. */
828 if (isnan(*target)) *target = 0.0;
829 } else if (aggregate == REDIS_AGGR_MIN) {
830 *target = val < *target ? val : *target;
831 } else if (aggregate == REDIS_AGGR_MAX) {
832 *target = val > *target ? val : *target;
833 } else {
834 /* safety net */
835 redisPanic("Unknown ZUNION/INTER aggregate type");
836 }
837 }
838
839 void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
840 int i, j, setnum;
841 int aggregate = REDIS_AGGR_SUM;
842 zsetopsrc *src;
843 robj *dstobj;
844 zset *dstzset;
845 zskiplistNode *znode;
846 dictIterator *di;
847 dictEntry *de;
848 int touched = 0;
849
850 /* expect setnum input keys to be given */
851 setnum = atoi(c->argv[2]->ptr);
852 if (setnum < 1) {
853 addReplyError(c,
854 "at least 1 input key is needed for ZUNIONSTORE/ZINTERSTORE");
855 return;
856 }
857
858 /* test if the expected number of keys would overflow */
859 if (3+setnum > c->argc) {
860 addReply(c,shared.syntaxerr);
861 return;
862 }
863
864 /* read keys to be used for input */
865 src = zmalloc(sizeof(zsetopsrc) * setnum);
866 for (i = 0, j = 3; i < setnum; i++, j++) {
867 robj *obj = lookupKeyWrite(c->db,c->argv[j]);
868 if (!obj) {
869 src[i].dict = NULL;
870 } else {
871 if (obj->type == REDIS_ZSET) {
872 src[i].dict = ((zset*)obj->ptr)->dict;
873 } else if (obj->type == REDIS_SET) {
874 src[i].dict = (obj->ptr);
875 } else {
876 zfree(src);
877 addReply(c,shared.wrongtypeerr);
878 return;
879 }
880 }
881
882 /* default all weights to 1 */
883 src[i].weight = 1.0;
884 }
885
886 /* parse optional extra arguments */
887 if (j < c->argc) {
888 int remaining = c->argc - j;
889
890 while (remaining) {
891 if (remaining >= (setnum + 1) && !strcasecmp(c->argv[j]->ptr,"weights")) {
892 j++; remaining--;
893 for (i = 0; i < setnum; i++, j++, remaining--) {
894 if (getDoubleFromObjectOrReply(c,c->argv[j],&src[i].weight,
895 "weight value is not a double") != REDIS_OK)
896 {
897 zfree(src);
898 return;
899 }
900 }
901 } else if (remaining >= 2 && !strcasecmp(c->argv[j]->ptr,"aggregate")) {
902 j++; remaining--;
903 if (!strcasecmp(c->argv[j]->ptr,"sum")) {
904 aggregate = REDIS_AGGR_SUM;
905 } else if (!strcasecmp(c->argv[j]->ptr,"min")) {
906 aggregate = REDIS_AGGR_MIN;
907 } else if (!strcasecmp(c->argv[j]->ptr,"max")) {
908 aggregate = REDIS_AGGR_MAX;
909 } else {
910 zfree(src);
911 addReply(c,shared.syntaxerr);
912 return;
913 }
914 j++; remaining--;
915 } else {
916 zfree(src);
917 addReply(c,shared.syntaxerr);
918 return;
919 }
920 }
921 }
922
923 /* sort sets from the smallest to largest, this will improve our
924 * algorithm's performance */
925 qsort(src,setnum,sizeof(zsetopsrc),qsortCompareZsetopsrcByCardinality);
926
927 dstobj = createZsetObject();
928 dstzset = dstobj->ptr;
929
930 if (op == REDIS_OP_INTER) {
931 /* skip going over all entries if the smallest zset is NULL or empty */
932 if (src[0].dict && dictSize(src[0].dict) > 0) {
933 /* precondition: as src[0].dict is non-empty and the zsets are ordered
934 * from small to large, all src[i > 0].dict are non-empty too */
935 di = dictGetIterator(src[0].dict);
936 while((de = dictNext(di)) != NULL) {
937 double score, value;
938
939 score = src[0].weight * zunionInterDictValue(de);
940 for (j = 1; j < setnum; j++) {
941 dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
942 if (other) {
943 value = src[j].weight * zunionInterDictValue(other);
944 zunionInterAggregate(&score,value,aggregate);
945 } else {
946 break;
947 }
948 }
949
950 /* Only continue when present in every source dict. */
951 if (j == setnum) {
952 robj *o = dictGetEntryKey(de);
953 znode = zslInsert(dstzset->zsl,score,o);
954 incrRefCount(o); /* added to skiplist */
955 dictAdd(dstzset->dict,o,&znode->score);
956 incrRefCount(o); /* added to dictionary */
957 }
958 }
959 dictReleaseIterator(di);
960 }
961 } else if (op == REDIS_OP_UNION) {
962 for (i = 0; i < setnum; i++) {
963 if (!src[i].dict) continue;
964
965 di = dictGetIterator(src[i].dict);
966 while((de = dictNext(di)) != NULL) {
967 double score, value;
968
969 /* skip key when already processed */
970 if (dictFind(dstzset->dict,dictGetEntryKey(de)) != NULL)
971 continue;
972
973 /* initialize score */
974 score = src[i].weight * zunionInterDictValue(de);
975
976 /* because the zsets are sorted by size, its only possible
977 * for sets at larger indices to hold this entry */
978 for (j = (i+1); j < setnum; j++) {
979 dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
980 if (other) {
981 value = src[j].weight * zunionInterDictValue(other);
982 zunionInterAggregate(&score,value,aggregate);
983 }
984 }
985
986 robj *o = dictGetEntryKey(de);
987 znode = zslInsert(dstzset->zsl,score,o);
988 incrRefCount(o); /* added to skiplist */
989 dictAdd(dstzset->dict,o,&znode->score);
990 incrRefCount(o); /* added to dictionary */
991 }
992 dictReleaseIterator(di);
993 }
994 } else {
995 /* unknown operator */
996 redisAssert(op == REDIS_OP_INTER || op == REDIS_OP_UNION);
997 }
998
999 if (dbDelete(c->db,dstkey)) {
1000 signalModifiedKey(c->db,dstkey);
1001 touched = 1;
1002 server.dirty++;
1003 }
1004 if (dstzset->zsl->length) {
1005 dbAdd(c->db,dstkey,dstobj);
1006 addReplyLongLong(c, dstzset->zsl->length);
1007 if (!touched) signalModifiedKey(c->db,dstkey);
1008 server.dirty++;
1009 } else {
1010 decrRefCount(dstobj);
1011 addReply(c, shared.czero);
1012 }
1013 zfree(src);
1014 }
1015
1016 void zunionstoreCommand(redisClient *c) {
1017 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_UNION);
1018 }
1019
1020 void zinterstoreCommand(redisClient *c) {
1021 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_INTER);
1022 }
1023
1024 void zrangeGenericCommand(redisClient *c, int reverse) {
1025 robj *o;
1026 long start;
1027 long end;
1028 int withscores = 0;
1029 int llen;
1030 int rangelen, j;
1031 zset *zsetobj;
1032 zskiplist *zsl;
1033 zskiplistNode *ln;
1034 robj *ele;
1035
1036 if ((getLongFromObjectOrReply(c, c->argv[2], &start, NULL) != REDIS_OK) ||
1037 (getLongFromObjectOrReply(c, c->argv[3], &end, NULL) != REDIS_OK)) return;
1038
1039 if (c->argc == 5 && !strcasecmp(c->argv[4]->ptr,"withscores")) {
1040 withscores = 1;
1041 } else if (c->argc >= 5) {
1042 addReply(c,shared.syntaxerr);
1043 return;
1044 }
1045
1046 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.emptymultibulk)) == NULL
1047 || checkType(c,o,REDIS_ZSET)) return;
1048 zsetobj = o->ptr;
1049 zsl = zsetobj->zsl;
1050 llen = zsl->length;
1051
1052 /* convert negative indexes */
1053 if (start < 0) start = llen+start;
1054 if (end < 0) end = llen+end;
1055 if (start < 0) start = 0;
1056
1057 /* Invariant: start >= 0, so this test will be true when end < 0.
1058 * The range is empty when start > end or start >= length. */
1059 if (start > end || start >= llen) {
1060 addReply(c,shared.emptymultibulk);
1061 return;
1062 }
1063 if (end >= llen) end = llen-1;
1064 rangelen = (end-start)+1;
1065
1066 /* check if starting point is trivial, before searching
1067 * the element in log(N) time */
1068 if (reverse) {
1069 ln = start == 0 ? zsl->tail : zslGetElementByRank(zsl, llen-start);
1070 } else {
1071 ln = start == 0 ?
1072 zsl->header->level[0].forward : zslGetElementByRank(zsl, start+1);
1073 }
1074
1075 /* Return the result in form of a multi-bulk reply */
1076 addReplyMultiBulkLen(c,withscores ? (rangelen*2) : rangelen);
1077 for (j = 0; j < rangelen; j++) {
1078 ele = ln->obj;
1079 addReplyBulk(c,ele);
1080 if (withscores)
1081 addReplyDouble(c,ln->score);
1082 ln = reverse ? ln->backward : ln->level[0].forward;
1083 }
1084 }
1085
1086 void zrangeCommand(redisClient *c) {
1087 zrangeGenericCommand(c,0);
1088 }
1089
1090 void zrevrangeCommand(redisClient *c) {
1091 zrangeGenericCommand(c,1);
1092 }
1093
1094 /* This command implements ZRANGEBYSCORE, ZREVRANGEBYSCORE and ZCOUNT.
1095 * If "justcount", only the number of elements in the range is returned. */
1096 void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) {
1097 zrangespec range;
1098 robj *o, *emptyreply;
1099 zset *zsetobj;
1100 zskiplist *zsl;
1101 zskiplistNode *ln;
1102 int offset = 0, limit = -1;
1103 int withscores = 0;
1104 unsigned long rangelen = 0;
1105 void *replylen = NULL;
1106 int minidx, maxidx;
1107
1108 /* Parse the range arguments. */
1109 if (reverse) {
1110 /* Range is given as [max,min] */
1111 maxidx = 2; minidx = 3;
1112 } else {
1113 /* Range is given as [min,max] */
1114 minidx = 2; maxidx = 3;
1115 }
1116
1117 if (zslParseRange(c->argv[minidx],c->argv[maxidx],&range) != REDIS_OK) {
1118 addReplyError(c,"min or max is not a double");
1119 return;
1120 }
1121
1122 /* Parse optional extra arguments. Note that ZCOUNT will exactly have
1123 * 4 arguments, so we'll never enter the following code path. */
1124 if (c->argc > 4) {
1125 int remaining = c->argc - 4;
1126 int pos = 4;
1127
1128 while (remaining) {
1129 if (remaining >= 1 && !strcasecmp(c->argv[pos]->ptr,"withscores")) {
1130 pos++; remaining--;
1131 withscores = 1;
1132 } else if (remaining >= 3 && !strcasecmp(c->argv[pos]->ptr,"limit")) {
1133 offset = atoi(c->argv[pos+1]->ptr);
1134 limit = atoi(c->argv[pos+2]->ptr);
1135 pos += 3; remaining -= 3;
1136 } else {
1137 addReply(c,shared.syntaxerr);
1138 return;
1139 }
1140 }
1141 }
1142
1143 /* Ok, lookup the key and get the range */
1144 emptyreply = justcount ? shared.czero : shared.emptymultibulk;
1145 if ((o = lookupKeyReadOrReply(c,c->argv[1],emptyreply)) == NULL ||
1146 checkType(c,o,REDIS_ZSET)) return;
1147 zsetobj = o->ptr;
1148 zsl = zsetobj->zsl;
1149
1150 /* If reversed, get the last node in range as starting point. */
1151 if (reverse) {
1152 ln = zslLastInRange(zsl,range);
1153 } else {
1154 ln = zslFirstInRange(zsl,range);
1155 }
1156
1157 /* No "first" element in the specified interval. */
1158 if (ln == NULL) {
1159 addReply(c,emptyreply);
1160 return;
1161 }
1162
1163 /* We don't know in advance how many matching elements there are in the
1164 * list, so we push this object that will represent the multi-bulk length
1165 * in the output buffer, and will "fix" it later */
1166 if (!justcount)
1167 replylen = addDeferredMultiBulkLength(c);
1168
1169 /* If there is an offset, just traverse the number of elements without
1170 * checking the score because that is done in the next loop. */
1171 while(ln && offset--) {
1172 ln = reverse ? ln->backward : ln->level[0].forward;
1173 }
1174
1175 while (ln && limit--) {
1176 /* Abort when the node is no longer in range. */
1177 if (reverse) {
1178 if (!zslValueGteMin(ln->score,&range)) break;
1179 } else {
1180 if (!zslValueLteMax(ln->score,&range)) break;
1181 }
1182
1183 /* Do our magic */
1184 rangelen++;
1185 if (!justcount) {
1186 addReplyBulk(c,ln->obj);
1187 if (withscores)
1188 addReplyDouble(c,ln->score);
1189 }
1190
1191 /* Move to next node */
1192 ln = reverse ? ln->backward : ln->level[0].forward;
1193 }
1194
1195 if (justcount) {
1196 addReplyLongLong(c,(long)rangelen);
1197 } else {
1198 setDeferredMultiBulkLength(c,replylen,
1199 withscores ? (rangelen*2) : rangelen);
1200 }
1201 }
1202
1203 void zrangebyscoreCommand(redisClient *c) {
1204 genericZrangebyscoreCommand(c,0,0);
1205 }
1206
1207 void zrevrangebyscoreCommand(redisClient *c) {
1208 genericZrangebyscoreCommand(c,1,0);
1209 }
1210
1211 void zcountCommand(redisClient *c) {
1212 genericZrangebyscoreCommand(c,0,1);
1213 }
1214
1215 void zcardCommand(redisClient *c) {
1216 robj *o;
1217 zset *zs;
1218
1219 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
1220 checkType(c,o,REDIS_ZSET)) return;
1221
1222 zs = o->ptr;
1223 addReplyLongLong(c,zs->zsl->length);
1224 }
1225
1226 void zscoreCommand(redisClient *c) {
1227 robj *o;
1228 zset *zs;
1229 dictEntry *de;
1230
1231 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
1232 checkType(c,o,REDIS_ZSET)) return;
1233
1234 zs = o->ptr;
1235 c->argv[2] = tryObjectEncoding(c->argv[2]);
1236 de = dictFind(zs->dict,c->argv[2]);
1237 if (!de) {
1238 addReply(c,shared.nullbulk);
1239 } else {
1240 double *score = dictGetEntryVal(de);
1241
1242 addReplyDouble(c,*score);
1243 }
1244 }
1245
1246 void zrankGenericCommand(redisClient *c, int reverse) {
1247 robj *o;
1248 zset *zs;
1249 zskiplist *zsl;
1250 dictEntry *de;
1251 unsigned long rank;
1252 double *score;
1253
1254 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
1255 checkType(c,o,REDIS_ZSET)) return;
1256
1257 zs = o->ptr;
1258 zsl = zs->zsl;
1259 c->argv[2] = tryObjectEncoding(c->argv[2]);
1260 de = dictFind(zs->dict,c->argv[2]);
1261 if (!de) {
1262 addReply(c,shared.nullbulk);
1263 return;
1264 }
1265
1266 score = dictGetEntryVal(de);
1267 rank = zslGetRank(zsl, *score, c->argv[2]);
1268 if (rank) {
1269 if (reverse) {
1270 addReplyLongLong(c, zsl->length - rank);
1271 } else {
1272 addReplyLongLong(c, rank-1);
1273 }
1274 } else {
1275 addReply(c,shared.nullbulk);
1276 }
1277 }
1278
1279 void zrankCommand(redisClient *c) {
1280 zrankGenericCommand(c, 0);
1281 }
1282
1283 void zrevrankCommand(redisClient *c) {
1284 zrankGenericCommand(c, 1);
1285 }