漫话Redis源码之六十一

这个文件不难，主要是实现：geoadd, georadius等命令功能

#include "geo.h"
#include "geohash_helper.h"
#include "debugmacro.h"
#include "pqsort.h"
 
/* Things exported from t_zset.c only for geo.c, since it is the only other
 * part of Redis that requires close zset introspection. */
unsigned char *zzlFirstInRange(unsigned char *zl, zrangespec *range);
int zslValueLteMax(double value, zrangespec *spec);
 
/* ====================================================================
 * This file implements the following commands:
 *
 *   - geoadd - add coordinates for value to geoset
 *   - georadius - search radius by coordinates in geoset
 *   - georadiusbymember - search radius based on geoset member position
 * ==================================================================== */
 
/* ====================================================================
 * geoArray implementation
 * ==================================================================== */
 
/* Create a new array of geoPoints. */
geoArray *geoArrayCreate(void) {
    geoArray *ga = zmalloc(sizeof(*ga));
    /* It gets allocated on first geoArrayAppend() call. */
    ga->array = NULL;
    ga->buckets = 0;
    ga->used = 0;
    return ga;
}
 
/* Add a new entry and return its pointer so that the caller can populate
 * it with data. */
geoPoint *geoArrayAppend(geoArray *ga) {
    if (ga->used == ga->buckets) {
        ga->buckets = (ga->buckets == 0) ? 8 : ga->buckets*2;
        ga->array = zrealloc(ga->array,sizeof(geoPoint)*ga->buckets);
    }
    geoPoint *gp = ga->array+ga->used;
    ga->used++;
    return gp;
}
 
/* Destroy a geoArray created with geoArrayCreate(). */
void geoArrayFree(geoArray *ga) {
    size_t i;
    for (i = 0; i < ga->used; i++) sdsfree(ga->array[i].member);
    zfree(ga->array);
    zfree(ga);
}
 
/* ====================================================================
 * Helpers
 * ==================================================================== */
int decodeGeohash(double bits, double *xy) {
    GeoHashBits hash = { .bits = (uint64_t)bits, .step = GEO_STEP_MAX };
    return geohashDecodeToLongLatWGS84(hash, xy);
}
 
/* Input Argument Helper */
/* Take a pointer to the latitude arg then use the next arg for longitude.
 * On parse error C_ERR is returned, otherwise C_OK. */
int extractLongLatOrReply(client *c, robj **argv, double *xy) {
    int i;
    for (i = 0; i < 2; i++) {
        if (getDoubleFromObjectOrReply(c, argv[i], xy + i, NULL) !=
            C_OK) {
            return C_ERR;
        }
    }
    if (xy[0] < GEO_LONG_MIN || xy[0] > GEO_LONG_MAX ||
        xy[1] < GEO_LAT_MIN  || xy[1] > GEO_LAT_MAX) {
        addReplyErrorFormat(c,
            "-ERR invalid longitude,latitude pair %f,%f\r\n",xy[0],xy[1]);
        return C_ERR;
    }
    return C_OK;
}
 
/* Input Argument Helper */
/* Decode lat/long from a zset member's score.
 * Returns C_OK on successful decoding, otherwise C_ERR is returned. */
int longLatFromMember(robj *zobj, robj *member, double *xy) {
    double score = 0;
 
    if (zsetScore(zobj, member->ptr, &score) == C_ERR) return C_ERR;
    if (!decodeGeohash(score, xy)) return C_ERR;
    return C_OK;
}
 
/* Check that the unit argument matches one of the known units, and returns
 * the conversion factor to meters (you need to divide meters by the conversion
 * factor to convert to the right unit).
 *
 * If the unit is not valid, an error is reported to the client, and a value
 * less than zero is returned. */
double extractUnitOrReply(client *c, robj *unit) {
    char *u = unit->ptr;
 
    if (!strcmp(u, "m")) {
        return 1;
    } else if (!strcmp(u, "km")) {
        return 1000;
    } else if (!strcmp(u, "ft")) {
        return 0.3048;
    } else if (!strcmp(u, "mi")) {
        return 1609.34;
    } else {
        addReplyError(c,
            "unsupported unit provided. please use m, km, ft, mi");
        return -1;
    }
}
 
/* Input Argument Helper.
 * Extract the distance from the specified two arguments starting at 'argv'
 * that should be in the form: <number> <unit>, and return C_OK or C_ERR means success or failure
 * *conversions is populated with the coefficient to use in order to convert meters to the unit.*/
int extractDistanceOrReply(client *c, robj **argv,
                              double *conversion, double *radius) {
    double distance;
    if (getDoubleFromObjectOrReply(c, argv[0], &distance,
                                   "need numeric radius") != C_OK) {
        return C_ERR;
    }
 
    if (distance < 0) {
        addReplyError(c,"radius cannot be negative");
        return C_ERR;
    }
    if (radius) *radius = distance;
 
    double to_meters = extractUnitOrReply(c,argv[1]);
    if (to_meters < 0) {
        return C_ERR;
    }
 
    if (conversion) *conversion = to_meters;
    return C_OK;
}
 
/* Input Argument Helper.
 * Extract height and width from the specified three arguments starting at 'argv'
 * that should be in the form: <number> <number> <unit>, and return C_OK or C_ERR means success or failure
 * *conversions is populated with the coefficient to use in order to convert meters to the unit.*/
int extractBoxOrReply(client *c, robj **argv, double *conversion,
                         double *width, double *height) {
    double h, w;
    if ((getDoubleFromObjectOrReply(c, argv[0], &w, "need numeric width") != C_OK) ||
        (getDoubleFromObjectOrReply(c, argv[1], &h, "need numeric height") != C_OK)) {
        return C_ERR;
    }
 
    if (h < 0 || w < 0) {
        addReplyError(c, "height or width cannot be negative");
        return C_ERR;
    }
    if (height) *height = h;
    if (width) *width = w;
 
    double to_meters = extractUnitOrReply(c,argv[2]);
    if (to_meters < 0) {
        return C_ERR;
    }
 
    if (conversion) *conversion = to_meters;
    return C_OK;
}
 
/* The default addReplyDouble has too much accuracy.  We use this
 * for returning location distances. "5.2145 meters away" is nicer
 * than "5.2144992818115 meters away." We provide 4 digits after the dot
 * so that the returned value is decently accurate even when the unit is
 * the kilometer. */
void addReplyDoubleDistance(client *c, double d) {
    char dbuf[128];
    int dlen = snprintf(dbuf, sizeof(dbuf), "%.4f", d);
    addReplyBulkCBuffer(c, dbuf, dlen);
}
 
/* Helper function for geoGetPointsInRange(): given a sorted set score
 * representing a point, and a GeoShape, appends this entry as a geoPoint
 * into the specified geoArray only if the point is within the search area.
 *
 * returns C_OK if the point is included, or REIDS_ERR if it is outside. */
int geoAppendIfWithinShape(geoArray *ga, GeoShape *shape, double score, sds member) {
    double distance = 0, xy[2];
 
    if (!decodeGeohash(score,xy)) return C_ERR; /* Can't decode. */
    /* Note that geohashGetDistanceIfInRadiusWGS84() takes arguments in
     * reverse order: longitude first, latitude later. */
    if (shape->type == CIRCULAR_TYPE) {
        if (!geohashGetDistanceIfInRadiusWGS84(shape->xy[0], shape->xy[1], xy[0], xy[1],
                                               shape->t.radius*shape->conversion, &distance)) return C_ERR;
    } else if (shape->type == RECTANGLE_TYPE) {
        if (!geohashGetDistanceIfInRectangle(shape->t.r.width * shape->conversion,
                                             shape->t.r.height * shape->conversion,
                                             shape->xy[0], shape->xy[1], xy[0], xy[1], &distance))
            return C_ERR;
    }
 
    /* Append the new element. */
    geoPoint *gp = geoArrayAppend(ga);
    gp->longitude = xy[0];
    gp->latitude = xy[1];
    gp->dist = distance;
    gp->member = member;
    gp->score = score;
    return C_OK;
}