]> git.saurik.com Git - redis.git/blob - src/t_zset.c
Don't encode element argument when dealing with ziplist
[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 char *zzlFind(robj *zobj, robj *ele, double *score) {
453 unsigned char *zl = zobj->ptr;
454 unsigned char *eptr = ziplistIndex(zl,0), *sptr;
455
456 ele = getDecodedObject(ele);
457 while (eptr != NULL) {
458 sptr = ziplistNext(zl,eptr);
459 redisAssert(sptr != NULL);
460
461 if (ziplistCompare(eptr,ele->ptr,sdslen(ele->ptr))) {
462 /* Matching element, pull out score. */
463 *score = zzlGetScore(sptr);
464 decrRefCount(ele);
465 return eptr;
466 }
467
468 /* Move to next element. */
469 eptr = ziplistNext(zl,sptr);
470 }
471
472 decrRefCount(ele);
473 return NULL;
474 }
475
476 /* Delete (element,score) pair from ziplist. Use local copy of eptr because we
477 * don't want to modify the one given as argument. */
478 int zzlDelete(robj *zobj, unsigned char *eptr) {
479 unsigned char *zl = zobj->ptr;
480 unsigned char *p = eptr;
481
482 /* TODO: add function to ziplist API to delete N elements from offset. */
483 zl = ziplistDelete(zl,&p);
484 zl = ziplistDelete(zl,&p);
485 zobj->ptr = zl;
486 return REDIS_OK;
487 }
488
489 int zzlInsertAt(robj *zobj, robj *ele, double score, unsigned char *eptr) {
490 unsigned char *zl = zobj->ptr;
491 unsigned char *sptr;
492 char scorebuf[128];
493 int scorelen;
494 int offset;
495
496 redisAssert(ele->encoding == REDIS_ENCODING_RAW);
497 scorelen = d2string(scorebuf,sizeof(scorebuf),score);
498 if (eptr == NULL) {
499 zl = ziplistPush(zl,ele->ptr,sdslen(ele->ptr),ZIPLIST_TAIL);
500 zl = ziplistPush(zl,(unsigned char*)scorebuf,scorelen,ZIPLIST_TAIL);
501 } else {
502 /* Keep offset relative to zl, as it might be re-allocated. */
503 offset = eptr-zl;
504 zl = ziplistInsert(zl,eptr,ele->ptr,sdslen(ele->ptr));
505 eptr = zl+offset;
506
507 /* Insert score after the element. */
508 redisAssert((sptr = ziplistNext(zl,eptr)) != NULL);
509 zl = ziplistInsert(zl,sptr,(unsigned char*)scorebuf,scorelen);
510 }
511
512 zobj->ptr = zl;
513 return REDIS_OK;
514 }
515
516 /* Insert (element,score) pair in ziplist. This function assumes the element is
517 * not yet present in the list. */
518 int zzlInsert(robj *zobj, robj *ele, double score) {
519 unsigned char *zl = zobj->ptr;
520 unsigned char *eptr = ziplistIndex(zl,0), *sptr;
521 double s;
522 int insert = 0;
523
524 ele = getDecodedObject(ele);
525 while (eptr != NULL) {
526 sptr = ziplistNext(zl,eptr);
527 redisAssert(sptr != NULL);
528 s = zzlGetScore(sptr);
529
530 if (s > score) {
531 /* First element with score larger than score for element to be
532 * inserted. This means we should take its spot in the list to
533 * maintain ordering. */
534 insert = 1;
535 } else if (s == score) {
536 /* Ensure lexicographical ordering for elements. */
537 if (zzlCompareElements(eptr,ele->ptr,sdslen(ele->ptr)) < 0)
538 insert = 1;
539 }
540
541 if (insert) {
542 zzlInsertAt(zobj,ele,score,eptr);
543 break;
544 }
545
546 /* Move to next element. */
547 eptr = ziplistNext(zl,sptr);
548 }
549
550 /* Push on tail of list when it was not yet inserted. */
551 if (!insert)
552 zzlInsertAt(zobj,ele,score,eptr);
553
554 decrRefCount(ele);
555 return REDIS_OK;
556 }
557
558 /*-----------------------------------------------------------------------------
559 * Sorted set commands
560 *----------------------------------------------------------------------------*/
561
562 /* This generic command implements both ZADD and ZINCRBY. */
563 void zaddGenericCommand(redisClient *c, int incr) {
564 static char *nanerr = "resulting score is not a number (NaN)";
565 robj *key = c->argv[1];
566 robj *ele;
567 robj *zobj;
568 robj *curobj;
569 double score, curscore = 0.0;
570
571 if (getDoubleFromObjectOrReply(c,c->argv[2],&score,NULL) != REDIS_OK)
572 return;
573
574 zobj = lookupKeyWrite(c->db,key);
575 if (zobj == NULL) {
576 zobj = createZsetZiplistObject();
577 dbAdd(c->db,key,zobj);
578 } else {
579 if (zobj->type != REDIS_ZSET) {
580 addReply(c,shared.wrongtypeerr);
581 return;
582 }
583 }
584
585 if (zobj->encoding == REDIS_ENCODING_ZIPLIST) {
586 unsigned char *eptr;
587
588 /* Prefer non-encoded element when dealing with ziplists. */
589 ele = c->argv[3];
590 if ((eptr = zzlFind(zobj,ele,&curscore)) != NULL) {
591 if (incr) {
592 score += curscore;
593 if (isnan(score)) {
594 addReplyError(c,nanerr);
595 /* Don't need to check if the sorted set is empty, because
596 * we know it has at least one element. */
597 return;
598 }
599 }
600
601 /* Remove and re-insert when score changed. */
602 if (score != curscore) {
603 redisAssert(zzlDelete(zobj,eptr) == REDIS_OK);
604 redisAssert(zzlInsert(zobj,ele,score) == REDIS_OK);
605
606 signalModifiedKey(c->db,key);
607 server.dirty++;
608 }
609
610 if (incr) /* ZINCRBY */
611 addReplyDouble(c,score);
612 else /* ZADD */
613 addReply(c,shared.czero);
614 } else {
615 redisAssert(zzlInsert(zobj,ele,score) == REDIS_OK);
616
617 signalModifiedKey(c->db,key);
618 server.dirty++;
619
620 if (incr) /* ZINCRBY */
621 addReplyDouble(c,score);
622 else /* ZADD */
623 addReply(c,shared.cone);
624 }
625 } else if (zobj->encoding == REDIS_ENCODING_RAW) {
626 zset *zs = zobj->ptr;
627 zskiplistNode *znode;
628 dictEntry *de;
629
630 ele = c->argv[3] = tryObjectEncoding(c->argv[3]);
631 de = dictFind(zs->dict,ele);
632 if (de != NULL) {
633 curobj = dictGetEntryKey(de);
634 curscore = *(double*)dictGetEntryVal(de);
635
636 if (incr) {
637 score += curscore;
638 if (isnan(score)) {
639 addReplyError(c,nanerr);
640 /* Don't need to check if the sorted set is empty, because
641 * we know it has at least one element. */
642 return;
643 }
644 }
645
646 /* Remove and re-insert when score changed. We can safely delete
647 * the key object from the skiplist, since the dictionary still has
648 * a reference to it. */
649 if (score != curscore) {
650 redisAssert(zslDelete(zs->zsl,curscore,curobj));
651 znode = zslInsert(zs->zsl,score,curobj);
652 incrRefCount(curobj); /* Re-inserted in skiplist. */
653 dictGetEntryVal(de) = &znode->score; /* Update score ptr. */
654
655 signalModifiedKey(c->db,key);
656 server.dirty++;
657 }
658
659 if (incr) /* ZINCRBY */
660 addReplyDouble(c,score);
661 else /* ZADD */
662 addReply(c,shared.czero);
663 } else {
664 znode = zslInsert(zs->zsl,score,ele);
665 incrRefCount(ele); /* Inserted in skiplist. */
666 redisAssert(dictAdd(zs->dict,ele,&znode->score) == DICT_OK);
667 incrRefCount(ele); /* Added to dictionary. */
668
669 signalModifiedKey(c->db,key);
670 server.dirty++;
671
672 if (incr) /* ZINCRBY */
673 addReplyDouble(c,score);
674 else /* ZADD */
675 addReply(c,shared.cone);
676 }
677 } else {
678 redisPanic("Unknown sorted set encoding");
679 }
680 }
681
682 void zaddCommand(redisClient *c) {
683 zaddGenericCommand(c,0);
684 }
685
686 void zincrbyCommand(redisClient *c) {
687 zaddGenericCommand(c,1);
688 }
689
690 void zremCommand(redisClient *c) {
691 robj *zsetobj;
692 zset *zs;
693 dictEntry *de;
694 double curscore;
695 int deleted;
696
697 if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
698 checkType(c,zsetobj,REDIS_ZSET)) return;
699
700 zs = zsetobj->ptr;
701 c->argv[2] = tryObjectEncoding(c->argv[2]);
702 de = dictFind(zs->dict,c->argv[2]);
703 if (de == NULL) {
704 addReply(c,shared.czero);
705 return;
706 }
707 /* Delete from the skiplist */
708 curscore = *(double*)dictGetEntryVal(de);
709 deleted = zslDelete(zs->zsl,curscore,c->argv[2]);
710 redisAssert(deleted != 0);
711
712 /* Delete from the hash table */
713 dictDelete(zs->dict,c->argv[2]);
714 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
715 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
716 signalModifiedKey(c->db,c->argv[1]);
717 server.dirty++;
718 addReply(c,shared.cone);
719 }
720
721 void zremrangebyscoreCommand(redisClient *c) {
722 zrangespec range;
723 long deleted;
724 robj *o;
725 zset *zs;
726
727 /* Parse the range arguments. */
728 if (zslParseRange(c->argv[2],c->argv[3],&range) != REDIS_OK) {
729 addReplyError(c,"min or max is not a double");
730 return;
731 }
732
733 if ((o = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
734 checkType(c,o,REDIS_ZSET)) return;
735
736 zs = o->ptr;
737 deleted = zslDeleteRangeByScore(zs->zsl,range,zs->dict);
738 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
739 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
740 if (deleted) signalModifiedKey(c->db,c->argv[1]);
741 server.dirty += deleted;
742 addReplyLongLong(c,deleted);
743 }
744
745 void zremrangebyrankCommand(redisClient *c) {
746 long start;
747 long end;
748 int llen;
749 long deleted;
750 robj *zsetobj;
751 zset *zs;
752
753 if ((getLongFromObjectOrReply(c, c->argv[2], &start, NULL) != REDIS_OK) ||
754 (getLongFromObjectOrReply(c, c->argv[3], &end, NULL) != REDIS_OK)) return;
755
756 if ((zsetobj = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL ||
757 checkType(c,zsetobj,REDIS_ZSET)) return;
758 zs = zsetobj->ptr;
759 llen = zs->zsl->length;
760
761 /* convert negative indexes */
762 if (start < 0) start = llen+start;
763 if (end < 0) end = llen+end;
764 if (start < 0) start = 0;
765
766 /* Invariant: start >= 0, so this test will be true when end < 0.
767 * The range is empty when start > end or start >= length. */
768 if (start > end || start >= llen) {
769 addReply(c,shared.czero);
770 return;
771 }
772 if (end >= llen) end = llen-1;
773
774 /* increment start and end because zsl*Rank functions
775 * use 1-based rank */
776 deleted = zslDeleteRangeByRank(zs->zsl,start+1,end+1,zs->dict);
777 if (htNeedsResize(zs->dict)) dictResize(zs->dict);
778 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]);
779 if (deleted) signalModifiedKey(c->db,c->argv[1]);
780 server.dirty += deleted;
781 addReplyLongLong(c, deleted);
782 }
783
784 typedef struct {
785 dict *dict;
786 double weight;
787 } zsetopsrc;
788
789 int qsortCompareZsetopsrcByCardinality(const void *s1, const void *s2) {
790 zsetopsrc *d1 = (void*) s1, *d2 = (void*) s2;
791 unsigned long size1, size2;
792 size1 = d1->dict ? dictSize(d1->dict) : 0;
793 size2 = d2->dict ? dictSize(d2->dict) : 0;
794 return size1 - size2;
795 }
796
797 #define REDIS_AGGR_SUM 1
798 #define REDIS_AGGR_MIN 2
799 #define REDIS_AGGR_MAX 3
800 #define zunionInterDictValue(_e) (dictGetEntryVal(_e) == NULL ? 1.0 : *(double*)dictGetEntryVal(_e))
801
802 inline static void zunionInterAggregate(double *target, double val, int aggregate) {
803 if (aggregate == REDIS_AGGR_SUM) {
804 *target = *target + val;
805 /* The result of adding two doubles is NaN when one variable
806 * is +inf and the other is -inf. When these numbers are added,
807 * we maintain the convention of the result being 0.0. */
808 if (isnan(*target)) *target = 0.0;
809 } else if (aggregate == REDIS_AGGR_MIN) {
810 *target = val < *target ? val : *target;
811 } else if (aggregate == REDIS_AGGR_MAX) {
812 *target = val > *target ? val : *target;
813 } else {
814 /* safety net */
815 redisPanic("Unknown ZUNION/INTER aggregate type");
816 }
817 }
818
819 void zunionInterGenericCommand(redisClient *c, robj *dstkey, int op) {
820 int i, j, setnum;
821 int aggregate = REDIS_AGGR_SUM;
822 zsetopsrc *src;
823 robj *dstobj;
824 zset *dstzset;
825 zskiplistNode *znode;
826 dictIterator *di;
827 dictEntry *de;
828 int touched = 0;
829
830 /* expect setnum input keys to be given */
831 setnum = atoi(c->argv[2]->ptr);
832 if (setnum < 1) {
833 addReplyError(c,
834 "at least 1 input key is needed for ZUNIONSTORE/ZINTERSTORE");
835 return;
836 }
837
838 /* test if the expected number of keys would overflow */
839 if (3+setnum > c->argc) {
840 addReply(c,shared.syntaxerr);
841 return;
842 }
843
844 /* read keys to be used for input */
845 src = zmalloc(sizeof(zsetopsrc) * setnum);
846 for (i = 0, j = 3; i < setnum; i++, j++) {
847 robj *obj = lookupKeyWrite(c->db,c->argv[j]);
848 if (!obj) {
849 src[i].dict = NULL;
850 } else {
851 if (obj->type == REDIS_ZSET) {
852 src[i].dict = ((zset*)obj->ptr)->dict;
853 } else if (obj->type == REDIS_SET) {
854 src[i].dict = (obj->ptr);
855 } else {
856 zfree(src);
857 addReply(c,shared.wrongtypeerr);
858 return;
859 }
860 }
861
862 /* default all weights to 1 */
863 src[i].weight = 1.0;
864 }
865
866 /* parse optional extra arguments */
867 if (j < c->argc) {
868 int remaining = c->argc - j;
869
870 while (remaining) {
871 if (remaining >= (setnum + 1) && !strcasecmp(c->argv[j]->ptr,"weights")) {
872 j++; remaining--;
873 for (i = 0; i < setnum; i++, j++, remaining--) {
874 if (getDoubleFromObjectOrReply(c,c->argv[j],&src[i].weight,
875 "weight value is not a double") != REDIS_OK)
876 {
877 zfree(src);
878 return;
879 }
880 }
881 } else if (remaining >= 2 && !strcasecmp(c->argv[j]->ptr,"aggregate")) {
882 j++; remaining--;
883 if (!strcasecmp(c->argv[j]->ptr,"sum")) {
884 aggregate = REDIS_AGGR_SUM;
885 } else if (!strcasecmp(c->argv[j]->ptr,"min")) {
886 aggregate = REDIS_AGGR_MIN;
887 } else if (!strcasecmp(c->argv[j]->ptr,"max")) {
888 aggregate = REDIS_AGGR_MAX;
889 } else {
890 zfree(src);
891 addReply(c,shared.syntaxerr);
892 return;
893 }
894 j++; remaining--;
895 } else {
896 zfree(src);
897 addReply(c,shared.syntaxerr);
898 return;
899 }
900 }
901 }
902
903 /* sort sets from the smallest to largest, this will improve our
904 * algorithm's performance */
905 qsort(src,setnum,sizeof(zsetopsrc),qsortCompareZsetopsrcByCardinality);
906
907 dstobj = createZsetObject();
908 dstzset = dstobj->ptr;
909
910 if (op == REDIS_OP_INTER) {
911 /* skip going over all entries if the smallest zset is NULL or empty */
912 if (src[0].dict && dictSize(src[0].dict) > 0) {
913 /* precondition: as src[0].dict is non-empty and the zsets are ordered
914 * from small to large, all src[i > 0].dict are non-empty too */
915 di = dictGetIterator(src[0].dict);
916 while((de = dictNext(di)) != NULL) {
917 double score, value;
918
919 score = src[0].weight * zunionInterDictValue(de);
920 for (j = 1; j < setnum; j++) {
921 dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
922 if (other) {
923 value = src[j].weight * zunionInterDictValue(other);
924 zunionInterAggregate(&score,value,aggregate);
925 } else {
926 break;
927 }
928 }
929
930 /* Only continue when present in every source dict. */
931 if (j == setnum) {
932 robj *o = dictGetEntryKey(de);
933 znode = zslInsert(dstzset->zsl,score,o);
934 incrRefCount(o); /* added to skiplist */
935 dictAdd(dstzset->dict,o,&znode->score);
936 incrRefCount(o); /* added to dictionary */
937 }
938 }
939 dictReleaseIterator(di);
940 }
941 } else if (op == REDIS_OP_UNION) {
942 for (i = 0; i < setnum; i++) {
943 if (!src[i].dict) continue;
944
945 di = dictGetIterator(src[i].dict);
946 while((de = dictNext(di)) != NULL) {
947 double score, value;
948
949 /* skip key when already processed */
950 if (dictFind(dstzset->dict,dictGetEntryKey(de)) != NULL)
951 continue;
952
953 /* initialize score */
954 score = src[i].weight * zunionInterDictValue(de);
955
956 /* because the zsets are sorted by size, its only possible
957 * for sets at larger indices to hold this entry */
958 for (j = (i+1); j < setnum; j++) {
959 dictEntry *other = dictFind(src[j].dict,dictGetEntryKey(de));
960 if (other) {
961 value = src[j].weight * zunionInterDictValue(other);
962 zunionInterAggregate(&score,value,aggregate);
963 }
964 }
965
966 robj *o = dictGetEntryKey(de);
967 znode = zslInsert(dstzset->zsl,score,o);
968 incrRefCount(o); /* added to skiplist */
969 dictAdd(dstzset->dict,o,&znode->score);
970 incrRefCount(o); /* added to dictionary */
971 }
972 dictReleaseIterator(di);
973 }
974 } else {
975 /* unknown operator */
976 redisAssert(op == REDIS_OP_INTER || op == REDIS_OP_UNION);
977 }
978
979 if (dbDelete(c->db,dstkey)) {
980 signalModifiedKey(c->db,dstkey);
981 touched = 1;
982 server.dirty++;
983 }
984 if (dstzset->zsl->length) {
985 dbAdd(c->db,dstkey,dstobj);
986 addReplyLongLong(c, dstzset->zsl->length);
987 if (!touched) signalModifiedKey(c->db,dstkey);
988 server.dirty++;
989 } else {
990 decrRefCount(dstobj);
991 addReply(c, shared.czero);
992 }
993 zfree(src);
994 }
995
996 void zunionstoreCommand(redisClient *c) {
997 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_UNION);
998 }
999
1000 void zinterstoreCommand(redisClient *c) {
1001 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_INTER);
1002 }
1003
1004 void zrangeGenericCommand(redisClient *c, int reverse) {
1005 robj *o;
1006 long start;
1007 long end;
1008 int withscores = 0;
1009 int llen;
1010 int rangelen, j;
1011 zset *zsetobj;
1012 zskiplist *zsl;
1013 zskiplistNode *ln;
1014 robj *ele;
1015
1016 if ((getLongFromObjectOrReply(c, c->argv[2], &start, NULL) != REDIS_OK) ||
1017 (getLongFromObjectOrReply(c, c->argv[3], &end, NULL) != REDIS_OK)) return;
1018
1019 if (c->argc == 5 && !strcasecmp(c->argv[4]->ptr,"withscores")) {
1020 withscores = 1;
1021 } else if (c->argc >= 5) {
1022 addReply(c,shared.syntaxerr);
1023 return;
1024 }
1025
1026 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.emptymultibulk)) == NULL
1027 || checkType(c,o,REDIS_ZSET)) return;
1028 zsetobj = o->ptr;
1029 zsl = zsetobj->zsl;
1030 llen = zsl->length;
1031
1032 /* convert negative indexes */
1033 if (start < 0) start = llen+start;
1034 if (end < 0) end = llen+end;
1035 if (start < 0) start = 0;
1036
1037 /* Invariant: start >= 0, so this test will be true when end < 0.
1038 * The range is empty when start > end or start >= length. */
1039 if (start > end || start >= llen) {
1040 addReply(c,shared.emptymultibulk);
1041 return;
1042 }
1043 if (end >= llen) end = llen-1;
1044 rangelen = (end-start)+1;
1045
1046 /* check if starting point is trivial, before searching
1047 * the element in log(N) time */
1048 if (reverse) {
1049 ln = start == 0 ? zsl->tail : zslGetElementByRank(zsl, llen-start);
1050 } else {
1051 ln = start == 0 ?
1052 zsl->header->level[0].forward : zslGetElementByRank(zsl, start+1);
1053 }
1054
1055 /* Return the result in form of a multi-bulk reply */
1056 addReplyMultiBulkLen(c,withscores ? (rangelen*2) : rangelen);
1057 for (j = 0; j < rangelen; j++) {
1058 ele = ln->obj;
1059 addReplyBulk(c,ele);
1060 if (withscores)
1061 addReplyDouble(c,ln->score);
1062 ln = reverse ? ln->backward : ln->level[0].forward;
1063 }
1064 }
1065
1066 void zrangeCommand(redisClient *c) {
1067 zrangeGenericCommand(c,0);
1068 }
1069
1070 void zrevrangeCommand(redisClient *c) {
1071 zrangeGenericCommand(c,1);
1072 }
1073
1074 /* This command implements ZRANGEBYSCORE, ZREVRANGEBYSCORE and ZCOUNT.
1075 * If "justcount", only the number of elements in the range is returned. */
1076 void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) {
1077 zrangespec range;
1078 robj *o, *emptyreply;
1079 zset *zsetobj;
1080 zskiplist *zsl;
1081 zskiplistNode *ln;
1082 int offset = 0, limit = -1;
1083 int withscores = 0;
1084 unsigned long rangelen = 0;
1085 void *replylen = NULL;
1086 int minidx, maxidx;
1087
1088 /* Parse the range arguments. */
1089 if (reverse) {
1090 /* Range is given as [max,min] */
1091 maxidx = 2; minidx = 3;
1092 } else {
1093 /* Range is given as [min,max] */
1094 minidx = 2; maxidx = 3;
1095 }
1096
1097 if (zslParseRange(c->argv[minidx],c->argv[maxidx],&range) != REDIS_OK) {
1098 addReplyError(c,"min or max is not a double");
1099 return;
1100 }
1101
1102 /* Parse optional extra arguments. Note that ZCOUNT will exactly have
1103 * 4 arguments, so we'll never enter the following code path. */
1104 if (c->argc > 4) {
1105 int remaining = c->argc - 4;
1106 int pos = 4;
1107
1108 while (remaining) {
1109 if (remaining >= 1 && !strcasecmp(c->argv[pos]->ptr,"withscores")) {
1110 pos++; remaining--;
1111 withscores = 1;
1112 } else if (remaining >= 3 && !strcasecmp(c->argv[pos]->ptr,"limit")) {
1113 offset = atoi(c->argv[pos+1]->ptr);
1114 limit = atoi(c->argv[pos+2]->ptr);
1115 pos += 3; remaining -= 3;
1116 } else {
1117 addReply(c,shared.syntaxerr);
1118 return;
1119 }
1120 }
1121 }
1122
1123 /* Ok, lookup the key and get the range */
1124 emptyreply = justcount ? shared.czero : shared.emptymultibulk;
1125 if ((o = lookupKeyReadOrReply(c,c->argv[1],emptyreply)) == NULL ||
1126 checkType(c,o,REDIS_ZSET)) return;
1127 zsetobj = o->ptr;
1128 zsl = zsetobj->zsl;
1129
1130 /* If reversed, get the last node in range as starting point. */
1131 if (reverse) {
1132 ln = zslLastInRange(zsl,range);
1133 } else {
1134 ln = zslFirstInRange(zsl,range);
1135 }
1136
1137 /* No "first" element in the specified interval. */
1138 if (ln == NULL) {
1139 addReply(c,emptyreply);
1140 return;
1141 }
1142
1143 /* We don't know in advance how many matching elements there are in the
1144 * list, so we push this object that will represent the multi-bulk length
1145 * in the output buffer, and will "fix" it later */
1146 if (!justcount)
1147 replylen = addDeferredMultiBulkLength(c);
1148
1149 /* If there is an offset, just traverse the number of elements without
1150 * checking the score because that is done in the next loop. */
1151 while(ln && offset--) {
1152 ln = reverse ? ln->backward : ln->level[0].forward;
1153 }
1154
1155 while (ln && limit--) {
1156 /* Abort when the node is no longer in range. */
1157 if (reverse) {
1158 if (!zslValueGteMin(ln->score,&range)) break;
1159 } else {
1160 if (!zslValueLteMax(ln->score,&range)) break;
1161 }
1162
1163 /* Do our magic */
1164 rangelen++;
1165 if (!justcount) {
1166 addReplyBulk(c,ln->obj);
1167 if (withscores)
1168 addReplyDouble(c,ln->score);
1169 }
1170
1171 /* Move to next node */
1172 ln = reverse ? ln->backward : ln->level[0].forward;
1173 }
1174
1175 if (justcount) {
1176 addReplyLongLong(c,(long)rangelen);
1177 } else {
1178 setDeferredMultiBulkLength(c,replylen,
1179 withscores ? (rangelen*2) : rangelen);
1180 }
1181 }
1182
1183 void zrangebyscoreCommand(redisClient *c) {
1184 genericZrangebyscoreCommand(c,0,0);
1185 }
1186
1187 void zrevrangebyscoreCommand(redisClient *c) {
1188 genericZrangebyscoreCommand(c,1,0);
1189 }
1190
1191 void zcountCommand(redisClient *c) {
1192 genericZrangebyscoreCommand(c,0,1);
1193 }
1194
1195 void zcardCommand(redisClient *c) {
1196 robj *o;
1197 zset *zs;
1198
1199 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
1200 checkType(c,o,REDIS_ZSET)) return;
1201
1202 zs = o->ptr;
1203 addReplyLongLong(c,zs->zsl->length);
1204 }
1205
1206 void zscoreCommand(redisClient *c) {
1207 robj *o;
1208 zset *zs;
1209 dictEntry *de;
1210
1211 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
1212 checkType(c,o,REDIS_ZSET)) return;
1213
1214 zs = o->ptr;
1215 c->argv[2] = tryObjectEncoding(c->argv[2]);
1216 de = dictFind(zs->dict,c->argv[2]);
1217 if (!de) {
1218 addReply(c,shared.nullbulk);
1219 } else {
1220 double *score = dictGetEntryVal(de);
1221
1222 addReplyDouble(c,*score);
1223 }
1224 }
1225
1226 void zrankGenericCommand(redisClient *c, int reverse) {
1227 robj *o;
1228 zset *zs;
1229 zskiplist *zsl;
1230 dictEntry *de;
1231 unsigned long rank;
1232 double *score;
1233
1234 if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.nullbulk)) == NULL ||
1235 checkType(c,o,REDIS_ZSET)) return;
1236
1237 zs = o->ptr;
1238 zsl = zs->zsl;
1239 c->argv[2] = tryObjectEncoding(c->argv[2]);
1240 de = dictFind(zs->dict,c->argv[2]);
1241 if (!de) {
1242 addReply(c,shared.nullbulk);
1243 return;
1244 }
1245
1246 score = dictGetEntryVal(de);
1247 rank = zslGetRank(zsl, *score, c->argv[2]);
1248 if (rank) {
1249 if (reverse) {
1250 addReplyLongLong(c, zsl->length - rank);
1251 } else {
1252 addReplyLongLong(c, rank-1);
1253 }
1254 } else {
1255 addReply(c,shared.nullbulk);
1256 }
1257 }
1258
1259 void zrankCommand(redisClient *c) {
1260 zrankGenericCommand(c, 0);
1261 }
1262
1263 void zrevrankCommand(redisClient *c) {
1264 zrankGenericCommand(c, 1);
1265 }