Add generic function to grow an sds value

Move logic concerned with setting a bit in an sds to the SETBIT command
instead of keeping it in sds.c. The function to grow an sds can and will
be reused for a command to set a range within a string value.

src/sds.c

@@ -116,6 +116,25 @@ static sds sdsMakeRoomFor(sds s, size_t addlen) {
     return newsh->buf;
 }
 
+/* Grow the sds to have the specified length. Bytes that were not part of
+ * the original length of the sds will be set to NULL. */
+sds sdsgrowsafe(sds s, size_t len) {
+    struct sdshdr *sh = (void*)(s-(sizeof(struct sdshdr)));
+    size_t totlen, curlen = sh->len;
+
+    if (len <= curlen) return s;
+    s = sdsMakeRoomFor(s,len-curlen);
+    if (s == NULL) return NULL;
+
+    /* Make sure added region doesn't contain garbage */
+    sh = (void*)(s-(sizeof(struct sdshdr)));
+    memset(s+curlen,0,(len-curlen+1)); /* also set trailing NULL byte */
+    totlen = sh->len+sh->free;
+    sh->len = len;
+    sh->free = totlen-sh->len;
+    return s;
+}
+
 sds sdscatlen(sds s, void *t, size_t len) {
     struct sdshdr *sh;
     size_t curlen = sdslen(s);
@@ -155,32 +174,6 @@ sds sdscpy(sds s, char *t) {
     return sdscpylen(s, t, strlen(t));
 }
 
-sds sdssetbit(sds s, size_t bit, int on) {
-    struct sdshdr *sh = (void*)(s-(sizeof(struct sdshdr)));
-    int byte = bit >> 3;
-    int reqlen = byte+1;
-
-    if (reqlen > sh->len) {
-        size_t totlen;
-
-        s = sdsMakeRoomFor(s,reqlen-sh->len);
-        if (s == NULL) return NULL;
-        sh = (void*)(s-(sizeof(struct sdshdr)));
-
-        /* Make sure added region doesn't contain garbage */
-        totlen = sh->len+sh->free;
-        memset(s+sh->len,0,sh->free+1);
-        sh->len = reqlen;
-        sh->free = totlen-sh->len;
-    }
-
-    bit = 7 - (bit & 0x7);
-    on &= 0x1;
-    s[byte] |= on << bit;
-    s[byte] &= ~((!on) << bit);
-    return s;
-}
-
 sds sdscatvprintf(sds s, const char *fmt, va_list ap) {
     va_list cpy;
     char *buf, *t;

src/sds.h

@@ -49,11 +49,11 @@ size_t sdslen(const sds s);
 sds sdsdup(const sds s);
 void sdsfree(sds s);
 size_t sdsavail(sds s);
+sds sdsgrowsafe(sds s, size_t len);
 sds sdscatlen(sds s, void *t, size_t len);
 sds sdscat(sds s, char *t);
 sds sdscpylen(sds s, char *t, size_t len);
 sds sdscpy(sds s, char *t);
-sds sdssetbit(sds s, size_t bit, int on);
 
 sds sdscatvprintf(sds s, const char *fmt, va_list ap);
 #ifdef __GNUC__

src/t_string.c

@@ -103,22 +103,26 @@ static int getBitOffsetFromArgument(redisClient *c, robj *o, size_t *offset) {
 
 void setbitCommand(redisClient *c) {
     robj *o;
+    char *err = "bit is not an integer or out of range";
     size_t bitoffset;
-    int on;
+    long long bitvalue;
+    int byte, bit, on;
 
     if (getBitOffsetFromArgument(c,c->argv[2],&bitoffset) != REDIS_OK)
         return;
 
-    on = ((sds)c->argv[3]->ptr)[0] - '0';
-    if (sdslen(c->argv[3]->ptr) != 1 || (on & ~1)) {
-        addReplyError(c,"bit should be 0 or 1");
+    if (getLongLongFromObjectOrReply(c,c->argv[3],&bitvalue,err) != REDIS_OK)
+        return;
+
+    /* A bit can only be set to be on or off... */
+    if (bitvalue & ~1) {
+        addReplyError(c,err);
         return;
     }
 
     o = lookupKeyWrite(c->db,c->argv[1]);
     if (o == NULL) {
-        sds value = sdssetbit(sdsempty(),bitoffset,on);
-        o = createObject(REDIS_STRING,value);
+        o = createObject(REDIS_STRING,sdsempty());
         dbAdd(c->db,c->argv[1],o);
     } else {
         if (checkType(c,o,REDIS_STRING)) return;
@@ -130,9 +134,15 @@ void setbitCommand(redisClient *c) {
             decrRefCount(decoded);
             dbReplace(c->db,c->argv[1],o);
         }
-
-        o->ptr = sdssetbit(o->ptr,bitoffset,on);
     }
+
+    byte = bitoffset >> 3;
+    bit = 7 - (bitoffset & 0x7);
+    on = bitvalue & 0x1;
+    o->ptr = sdsgrowsafe(o->ptr,byte+1);
+    ((char*)o->ptr)[byte] |= on << bit;
+    ((char*)o->ptr)[byte] &= ~((!on) << bit);
+
     touchWatchedKey(c->db,c->argv[1]);
     server.dirty++;
     addReply(c,shared.cone);

tests/unit/basic.tcl

@@ -416,10 +416,10 @@ start_server {tags {"basic"}} {
 
     test "SETBIT with non-bit argument" {
         r del mykey
-        assert_error "*0 or 1*" {r setbit mykey 0 -1}
-        assert_error "*0 or 1*" {r setbit mykey 0  2}
-        assert_error "*0 or 1*" {r setbit mykey 0 10}
-        assert_error "*0 or 1*" {r setbit mykey 0 01}
+        assert_error "*out of range*" {r setbit mykey 0 -1}
+        assert_error "*out of range*" {r setbit mykey 0  2}
+        assert_error "*out of range*" {r setbit mykey 0 10}
+        assert_error "*out of range*" {r setbit mykey 0 20}
     }
 
     test "GETBIT against non-existing key" {