1 #include "redis.h" 2 3 #include <math.h> 4 5 /*----------------------------------------------------------------------------- 6 * Sorted set API 7 *----------------------------------------------------------------------------*/ 8 9 /* ZSETs are ordered sets using two data structures to hold the same elements 10 * in order to get O(log(N)) INSERT and REMOVE operations into a sorted 11 * data structure. 12 * 13 * The elements are added to an hash table mapping Redis objects to scores. 14 * At the same time the elements are added to a skip list mapping scores 15 * to Redis objects (so objects are sorted by scores in this "view"). */ 16 17 /* This skiplist implementation is almost a C translation of the original 18 * algorithm described by William Pugh in "Skip Lists: A Probabilistic 19 * Alternative to Balanced Trees", modified in three ways: 20 * a) this implementation allows for repeated values. 21 * b) the comparison is not just by key (our 'score') but by satellite data. 22 * c) there is a back pointer, so it's a doubly linked list with the back 23 * pointers being only at "level 1". This allows to traverse the list 24 * from tail to head, useful for ZREVRANGE. */ 25 26 zskiplistNode *zslCreateNode(int level, double score, robj *obj) { 27 zskiplistNode *zn = zmalloc(sizeof(*zn)+level*sizeof(struct zskiplistLevel)); 28 zn->score = score; 29 zn->obj = obj; 30 return zn; 31 } 32 33 zskiplist *zslCreate(void) { 34 int j; 35 zskiplist *zsl; 36 37 zsl = zmalloc(sizeof(*zsl)); 38 zsl->level = 1; 39 zsl->length = 0; 40 zsl->header = zslCreateNode(ZSKIPLIST_MAXLEVEL,0,NULL); 41 for (j = 0; j < ZSKIPLIST_MAXLEVEL; j++) { 42 zsl->header->level[j].forward = NULL; 43 zsl->header->level[j].span = 0; 44 } 45 zsl->header->backward = NULL; 46 zsl->tail = NULL; 47 return zsl; 48 } 49 50 void zslFreeNode(zskiplistNode *node) { 51 decrRefCount(node->obj); 52 zfree(node); 53 } 54 55 void zslFree(zskiplist *zsl) { 56 zskiplistNode *node = zsl->header->level[0].forward, *next; 57 58 zfree(zsl->header); 59 while(node) { 60 next = node->level[0].forward; 61 zslFreeNode(node); 62 node = next; 63 } 64 zfree(zsl); 65 } 66 67 int zslRandomLevel(void) { 68 int level = 1; 69 while ((random()&0xFFFF) < (ZSKIPLIST_P * 0xFFFF)) 70 level += 1; 71 return (level<ZSKIPLIST_MAXLEVEL) ? level : ZSKIPLIST_MAXLEVEL; 72 } 73 74 zskiplistNode *zslInsert(zskiplist *zsl, double score, robj *obj) { 75 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x; 76 unsigned int rank[ZSKIPLIST_MAXLEVEL]; 77 int i, level; 78 79 x = zsl->header; 80 for (i = zsl->level-1; i >= 0; i--) { 81 /* store rank that is crossed to reach the insert position */ 82 rank[i] = i == (zsl->level-1) ? 0 : rank[i+1]; 83 while (x->level[i].forward && 84 (x->level[i].forward->score < score || 85 (x->level[i].forward->score == score && 86 compareStringObjects(x->level[i].forward->obj,obj) < 0))) { 87 rank[i] += x->level[i].span; 88 x = x->level[i].forward; 89 } 90 update[i] = x; 91 } 92 /* we assume the key is not already inside, since we allow duplicated 93 * scores, and the re-insertion of score and redis object should never 94 * happpen since the caller of zslInsert() should test in the hash table 95 * if the element is already inside or not. */ 96 level = zslRandomLevel(); 97 if (level > zsl->level) { 98 for (i = zsl->level; i < level; i++) { 99 rank[i] = 0; 100 update[i] = zsl->header; 101 update[i]->level[i].span = zsl->length; 102 } 103 zsl->level = level; 104 } 105 x = zslCreateNode(level,score,obj); 106 for (i = 0; i < level; i++) { 107 x->level[i].forward = update[i]->level[i].forward; 108 update[i]->level[i].forward = x; 109 110 /* update span covered by update[i] as x is inserted here */ 111 x->level[i].span = update[i]->level[i].span - (rank[0] - rank[i]); 112 update[i]->level[i].span = (rank[0] - rank[i]) + 1; 113 } 114 115 /* increment span for untouched levels */ 116 for (i = level; i < zsl->level; i++) { 117 update[i]->level[i].span++; 118 } 119 120 x->backward = (update[0] == zsl->header) ? NULL : update[0]; 121 if (x->level[0].forward) 122 x->level[0].forward->backward = x; 123 else 124 zsl->tail = x; 125 zsl->length++; 126 return x; 127 } 128 129 /* Internal function used by zslDelete, zslDeleteByScore and zslDeleteByRank */ 130 void zslDeleteNode(zskiplist *zsl, zskiplistNode *x, zskiplistNode **update) { 131 int i; 132 for (i = 0; i < zsl->level; i++) { 133 if (update[i]->level[i].forward == x) { 134 update[i]->level[i].span += x->level[i].span - 1; 135 update[i]->level[i].forward = x->level[i].forward; 136 } else { 137 update[i]->level[i].span -= 1; 138 } 139 } 140 if (x->level[0].forward) { 141 x->level[0].forward->backward = x->backward; 142 } else { 143 zsl->tail = x->backward; 144 } 145 while(zsl->level > 1 && zsl->header->level[zsl->level-1].forward == NULL) 146 zsl->level--; 147 zsl->length--; 148 } 149 150 /* Delete an element with matching score/object from the skiplist. */ 151 int zslDelete(zskiplist *zsl, double score, robj *obj) { 152 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x; 153 int i; 154 155 x = zsl->header; 156 for (i = zsl->level-1; i >= 0; i--) { 157 while (x->level[i].forward && 158 (x->level[i].forward->score < score || 159 (x->level[i].forward->score == score && 160 compareStringObjects(x->level[i].forward->obj,obj) < 0))) 161 x = x->level[i].forward; 162 update[i] = x; 163 } 164 /* We may have multiple elements with the same score, what we need 165 * is to find the element with both the right score and object. */ 166 x = x->level[0].forward; 167 if (x && score == x->score && equalStringObjects(x->obj,obj)) { 168 zslDeleteNode(zsl, x, update); 169 zslFreeNode(x); 170 return 1; 171 } else { 172 return 0; /* not found */ 173 } 174 return 0; /* not found */ 175 } 176 177 /* Struct to hold a inclusive/exclusive range spec. */ 178 typedef struct { 179 double min, max; 180 int minex, maxex; /* are min or max exclusive? */ 181 } zrangespec; 182 183 /* Delete all the elements with score between min and max from the skiplist. 184 * Min and mx are inclusive, so a score >= min || score <= max is deleted. 185 * Note that this function takes the reference to the hash table view of the 186 * sorted set, in order to remove the elements from the hash table too. */ 187 unsigned long zslDeleteRangeByScore(zskiplist *zsl, zrangespec range, dict *dict) { 188 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x; 189 unsigned long removed = 0; 190 int i; 191 192 x = zsl->header; 193 for (i = zsl->level-1; i >= 0; i--) { 194 while (x->level[i].forward && (range.minex ? 195 x->level[i].forward->score <= range.min : 196 x->level[i].forward->score < range.min)) 197 x = x->level[i].forward; 198 update[i] = x; 199 } 200 201 /* Current node is the last with score < or <= min. */ 202 x = x->level[0].forward; 203 204 /* Delete nodes while in range. */ 205 while (x && (range.maxex ? x->score < range.max : x->score <= range.max)) { 206 zskiplistNode *next = x->level[0].forward; 207 zslDeleteNode(zsl,x,update); 208 dictDelete(dict,x->obj); 209 zslFreeNode(x); 210 removed++; 211 x = next; 212 } 213 return removed; 214 } 215 216 /* Delete all the elements with rank between start and end from the skiplist. 217 * Start and end are inclusive. Note that start and end need to be 1-based */ 218 unsigned long zslDeleteRangeByRank(zskiplist *zsl, unsigned int start, unsigned int end, dict *dict) { 219 zskiplistNode *update[ZSKIPLIST_MAXLEVEL], *x; 220 unsigned long traversed = 0, removed = 0; 221 int i; 222 223 x = zsl->header; 224 for (i = zsl->level-1; i >= 0; i--) { 225 while (x->level[i].forward && (traversed + x->level[i].span) < start) { 226 traversed += x->level[i].span; 227 x = x->level[i].forward; 228 } 229 update[i] = x; 230 } 231 232 traversed++; 233 x = x->level[0].forward; 234 while (x && traversed <= end) { 235 zskiplistNode *next = x->level[0].forward; 236 zslDeleteNode(zsl,x,update); 237 dictDelete(dict,x->obj); 238 zslFreeNode(x); 239 removed++; 240 traversed++; 241 x = next; 242 } 243 return removed; 244 } 245 246 /* Find the first node having a score equal or greater than the specified one. 247 * Returns NULL if there is no match. */ 248 zskiplistNode *zslFirstWithScore(zskiplist *zsl, double score) { 249 zskiplistNode *x; 250 int i; 251 252 x = zsl->header; 253 for (i = zsl->level-1; i >= 0; i--) { 254 while (x->level[i].forward && x->level[i].forward->score < score) 255 x = x->level[i].forward; 256 } 257 /* We may have multiple elements with the same score, what we need 258 * is to find the element with both the right score and object. */ 259 return x->level[0].forward; 260 } 261 262 /* Find the rank for an element by both score and key. 263 * Returns 0 when the element cannot be found, rank otherwise. 264 * Note that the rank is 1-based due to the span of zsl->header to the 265 * first element. */ 266 unsigned long zslistTypeGetRank(zskiplist *zsl, double score, robj *o) { 267 zskiplistNode *x; 268 unsigned long rank = 0; 269 int i; 270 271 x = zsl->header; 272 for (i = zsl->level-1; i >= 0; i--) { 273 while (x->level[i].forward && 274 (x->level[i].forward->score < score || 275 (x->level[i].forward->score == score && 276 compareStringObjects(x->level[i].forward->obj,o) <= 0))) { 277 rank += x->level[i].span; 278 x = x->level[i].forward; 279 } 280 281 /* x might be equal to zsl->header, so test if obj is non-NULL */ 282 if (x->obj && equalStringObjects(x->obj,o)) { 283 return rank; 284 } 285 } 286 return 0; 287 } 288 289 /* Finds an element by its rank. The rank argument needs to be 1-based. */ 290 zskiplistNode* zslistTypeGetElementByRank(zskiplist *zsl, unsigned long rank) { 291 zskiplistNode *x; 292 unsigned long traversed = 0; 293 int i; 294 295 x = zsl->header; 296 for (i = zsl->level-1; i >= 0; i--) { 297 while (x->level[i].forward && (traversed + x->level[i].span) <= rank) 298 { 299 traversed += x->level[i].span; 300 x = x->level[i].forward; 301 } 302 if (traversed == rank) { 303 return x; 304 } 305 } 306 return NULL; 307 } 308 309 /* Populate the rangespec according to the objects min and max. */ 310 static int zslParseRange(robj *min, robj *max, zrangespec *spec) { 311 char *eptr; 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] == '(') { 322 spec->min = strtod((char*)min->ptr+1,&eptr); 323 if (eptr[0] != '\0' || isnan(spec->min)) return REDIS_ERR; 324 spec->minex = 1; 325 } else { 326 spec->min = strtod((char*)min->ptr,&eptr); 327 if (eptr[0] != '\0' || isnan(spec->min)) return REDIS_ERR; 328 } 329 } 330 if (max->encoding == REDIS_ENCODING_INT) { 331 spec->max = (long)max->ptr; 332 } else { 333 if (((char*)max->ptr)[0] == '(') { 334 spec->max = strtod((char*)max->ptr+1,&eptr); 335 if (eptr[0] != '\0' || isnan(spec->max)) return REDIS_ERR; 336 spec->maxex = 1; 337 } else { 338 spec->max = strtod((char*)max->ptr,&eptr); 339 if (eptr[0] != '\0' || isnan(spec->max)) return REDIS_ERR; 340 } 341 } 342 343 return REDIS_OK; 344 } 345 346 347 /*----------------------------------------------------------------------------- 348 * Sorted set commands 349 *----------------------------------------------------------------------------*/ 350 351 /* This generic command implements both ZADD and ZINCRBY. */ 352 void zaddGenericCommand(redisClient *c, robj *key, robj *ele, double score, int incr) { 353 robj *zsetobj; 354 zset *zs; 355 zskiplistNode *znode; 356 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 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) { 372 /* Read the old score. If the element was not present starts from 0 */ 373 dictEntry *de = dictFind(zs->dict,ele); 374 if (de != NULL) 375 score += *(double*)dictGetEntryVal(de); 376 377 if (isnan(score)) { 378 addReplyError(c,"resulting score is not a number (NaN)"); 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 } 384 } 385 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 */ 394 incrRefCount(ele); /* added to hash */ 395 znode = zslInsert(zs->zsl,score,ele); 396 incrRefCount(ele); /* added to skiplist */ 397 398 /* Update the score in the dict entry */ 399 de = dictFind(zs->dict,ele); 400 redisAssert(de != NULL); 401 dictGetEntryVal(de) = &znode->score; 402 signalModifiedKey(c->db,c->argv[1]); 403 server.dirty++; 404 if (incr) 405 addReplyDouble(c,score); 406 else 407 addReply(c,shared.cone); 408 } else { 409 dictEntry *de; 410 robj *curobj; 411 double *curscore; 412 int deleted; 413 414 /* Update score */ 415 de = dictFind(zs->dict,ele); 416 redisAssert(de != NULL); 417 curobj = dictGetEntryKey(de); 418 curscore = dictGetEntryVal(de); 419 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); 424 redisAssert(deleted != 0); 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; 430 signalModifiedKey(c->db,c->argv[1]); 431 server.dirty++; 432 } 433 if (incr) 434 addReplyDouble(c,score); 435 else 436 addReply(c,shared.czero); 437 } 438 } 439 440 void zaddCommand(redisClient *c) { 441 double scoreval; 442 if (getDoubleFromObjectOrReply(c,c->argv[2],&scoreval,NULL) != REDIS_OK) return; 443 c->argv[3] = tryObjectEncoding(c->argv[3]); 444 zaddGenericCommand(c,c->argv[1],c->argv[3],scoreval,0); 445 } 446 447 void zincrbyCommand(redisClient *c) { 448 double scoreval; 449 if (getDoubleFromObjectOrReply(c,c->argv[2],&scoreval,NULL) != REDIS_OK) return; 450 c->argv[3] = tryObjectEncoding(c->argv[3]); 451 zaddGenericCommand(c,c->argv[1],c->argv[3],scoreval,1); 452 } 453 454 void zremCommand(redisClient *c) { 455 robj *zsetobj; 456 zset *zs; 457 dictEntry *de; 458 double curscore; 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; 465 c->argv[2] = tryObjectEncoding(c->argv[2]); 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 */ 472 curscore = *(double*)dictGetEntryVal(de); 473 deleted = zslDelete(zs->zsl,curscore,c->argv[2]); 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]); 480 signalModifiedKey(c->db,c->argv[1]); 481 server.dirty++; 482 addReply(c,shared.cone); 483 } 484 485 void zremrangebyscoreCommand(redisClient *c) { 486 zrangespec range; 487 long deleted; 488 robj *o; 489 zset *zs; 490 491 /* Parse the range arguments. */ 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 } 496 497 if ((o = lookupKeyWriteOrReply(c,c->argv[1],shared.czero)) == NULL || 498 checkType(c,o,REDIS_ZSET)) return; 499 500 zs = o->ptr; 501 deleted = zslDeleteRangeByScore(zs->zsl,range,zs->dict); 502 if (htNeedsResize(zs->dict)) dictResize(zs->dict); 503 if (dictSize(zs->dict) == 0) dbDelete(c->db,c->argv[1]); 504 if (deleted) signalModifiedKey(c->db,c->argv[1]); 505 server.dirty += deleted; 506 addReplyLongLong(c,deleted); 507 } 508 509 void 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; 529 530 /* Invariant: start >= 0, so this test will be true when end < 0. 531 * The range is empty when start > end or start >= length. */ 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]); 543 if (deleted) signalModifiedKey(c->db,c->argv[1]); 544 server.dirty += deleted; 545 addReplyLongLong(c, deleted); 546 } 547 548 typedef struct { 549 dict *dict; 550 double weight; 551 } zsetopsrc; 552 553 int 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 566 inline static void zunionInterAggregate(double *target, double val, int aggregate) { 567 if (aggregate == REDIS_AGGR_SUM) { 568 *target = *target + val; 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; 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 583 void 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; 589 zskiplistNode *znode; 590 dictIterator *di; 591 dictEntry *de; 592 int touched = 0; 593 594 /* expect setnum input keys to be given */ 595 setnum = atoi(c->argv[2]->ptr); 596 if (setnum < 1) { 597 addReplyError(c, 598 "at least 1 input key is needed for ZUNIONSTORE/ZINTERSTORE"); 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--) { 638 if (getDoubleFromObjectOrReply(c,c->argv[j],&src[i].weight, 639 "weight value is not a double") != REDIS_OK) 640 { 641 zfree(src); 642 return; 643 } 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) { 681 double score, value; 682 683 score = src[0].weight * zunionInterDictValue(de); 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); 688 zunionInterAggregate(&score,value,aggregate); 689 } else { 690 break; 691 } 692 } 693 694 /* Only continue when present in every source dict. */ 695 if (j == setnum) { 696 robj *o = dictGetEntryKey(de); 697 znode = zslInsert(dstzset->zsl,score,o); 698 incrRefCount(o); /* added to skiplist */ 699 dictAdd(dstzset->dict,o,&znode->score); 700 incrRefCount(o); /* added to dictionary */ 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) { 711 double score, value; 712 713 /* skip key when already processed */ 714 if (dictFind(dstzset->dict,dictGetEntryKey(de)) != NULL) 715 continue; 716 717 /* initialize score */ 718 score = src[i].weight * zunionInterDictValue(de); 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); 726 zunionInterAggregate(&score,value,aggregate); 727 } 728 } 729 730 robj *o = dictGetEntryKey(de); 731 znode = zslInsert(dstzset->zsl,score,o); 732 incrRefCount(o); /* added to skiplist */ 733 dictAdd(dstzset->dict,o,&znode->score); 734 incrRefCount(o); /* added to dictionary */ 735 } 736 dictReleaseIterator(di); 737 } 738 } else { 739 /* unknown operator */ 740 redisAssert(op == REDIS_OP_INTER || op == REDIS_OP_UNION); 741 } 742 743 if (dbDelete(c->db,dstkey)) { 744 signalModifiedKey(c->db,dstkey); 745 touched = 1; 746 server.dirty++; 747 } 748 if (dstzset->zsl->length) { 749 dbAdd(c->db,dstkey,dstobj); 750 addReplyLongLong(c, dstzset->zsl->length); 751 if (!touched) signalModifiedKey(c->db,dstkey); 752 server.dirty++; 753 } else { 754 decrRefCount(dstobj); 755 addReply(c, shared.czero); 756 } 757 zfree(src); 758 } 759 760 void zunionstoreCommand(redisClient *c) { 761 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_UNION); 762 } 763 764 void zinterstoreCommand(redisClient *c) { 765 zunionInterGenericCommand(c,c->argv[1], REDIS_OP_INTER); 766 } 767 768 void 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; 800 801 /* Invariant: start >= 0, so this test will be true when end < 0. 802 * The range is empty when start > end or start >= length. */ 803 if (start > end || start >= llen) { 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 ? 816 zsl->header->level[0].forward : zslistTypeGetElementByRank(zsl, start+1); 817 } 818 819 /* Return the result in form of a multi-bulk reply */ 820 addReplyMultiBulkLen(c,withscores ? (rangelen*2) : rangelen); 821 for (j = 0; j < rangelen; j++) { 822 ele = ln->obj; 823 addReplyBulk(c,ele); 824 if (withscores) 825 addReplyDouble(c,ln->score); 826 ln = reverse ? ln->backward : ln->level[0].forward; 827 } 828 } 829 830 void zrangeCommand(redisClient *c) { 831 zrangeGenericCommand(c,0); 832 } 833 834 void zrevrangeCommand(redisClient *c) { 835 zrangeGenericCommand(c,1); 836 } 837 838 /* This command implements ZRANGEBYSCORE, ZREVRANGEBYSCORE and ZCOUNT. 839 * If "justcount", only the number of elements in the range is returned. */ 840 void genericZrangebyscoreCommand(redisClient *c, int reverse, int justcount) { 841 zrangespec range; 842 robj *o, *emptyreply; 843 zset *zsetobj; 844 zskiplist *zsl; 845 zskiplistNode *ln; 846 int offset = 0, limit = -1; 847 int withscores = 0; 848 unsigned long rangelen = 0; 849 void *replylen = NULL; 850 851 /* Parse the range arguments. */ 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 } 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 } 876 } 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 } 907 } else { 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 } 912 } 913 914 /* No "first" element in the specified interval. */ 915 if (ln == NULL) { 916 addReply(c,emptyreply); 917 return; 918 } 919 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); 926 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 } 935 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; 945 } 946 } else { 947 if (range.maxex) { 948 /* Element should have score < range.max */ 949 if (ln->score >= range.max) break; 950 } else { 951 /* Element should have score <= range.max */ 952 if (ln->score > range.max) break; 953 } 954 } 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); 975 } 976 } 977 978 void zrangebyscoreCommand(redisClient *c) { 979 genericZrangebyscoreCommand(c,0,0); 980 } 981 982 void zrevrangebyscoreCommand(redisClient *c) { 983 genericZrangebyscoreCommand(c,1,0); 984 } 985 986 void zcountCommand(redisClient *c) { 987 genericZrangebyscoreCommand(c,0,1); 988 } 989 990 void 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; 998 addReplyLongLong(c,zs->zsl->length); 999 } 1000 1001 void 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; 1010 c->argv[2] = tryObjectEncoding(c->argv[2]); 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 1021 void 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; 1034 c->argv[2] = tryObjectEncoding(c->argv[2]); 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 1054 void zrankCommand(redisClient *c) { 1055 zrankGenericCommand(c, 0); 1056 } 1057 1058 void zrevrankCommand(redisClient *c) { 1059 zrankGenericCommand(c, 1); 1060 } 1061