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