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