Add commands SETBIT/GETBIT

src/redis.c

@@ -78,6 +78,8 @@ struct redisCommand readonlyCommandTable[] = {
     {"strlen",strlenCommand,2,0,NULL,1,1,1},
     {"del",delCommand,-2,0,NULL,0,0,0},
     {"exists",existsCommand,2,0,NULL,1,1,1},
+    {"setbit",setbitCommand,4,REDIS_CMD_DENYOOM,NULL,1,1,1},
+    {"getbit",getbitCommand,3,0,NULL,1,1,1},
     {"incr",incrCommand,2,REDIS_CMD_DENYOOM,NULL,1,1,1},
     {"decr",decrCommand,2,REDIS_CMD_DENYOOM,NULL,1,1,1},
     {"mget",mgetCommand,-2,0,NULL,1,-1,1},

src/redis.h

@@ -887,6 +887,8 @@ void setexCommand(redisClient *c);
 void getCommand(redisClient *c);
 void delCommand(redisClient *c);
 void existsCommand(redisClient *c);
+void setbitCommand(redisClient *c);
+void getbitCommand(redisClient *c);
 void incrCommand(redisClient *c);
 void decrCommand(redisClient *c);
 void incrbyCommand(redisClient *c);

src/sds.c

@@ -155,6 +155,32 @@ 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

@@ -53,6 +53,7 @@ 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

+#include <limits.h>
 #include "redis.h"
 
 /*-----------------------------------------------------------------------------
@@ -80,6 +81,87 @@ void getsetCommand(redisClient *c) {
     removeExpire(c->db,c->argv[1]);
 }
 
+static int getBitOffsetFromArgument(redisClient *c, robj *o, size_t *offset) {
+    long long loffset;
+    char *err = "bit offset is not an integer or out of range";
+
+    if (getLongLongFromObjectOrReply(c,o,&loffset,err) != REDIS_OK)
+        return REDIS_ERR;
+
+    /* Limit offset to SIZE_T_MAX or 1GB in bytes */
+    if ((loffset < 0) ||
+        ((unsigned long long)loffset >= (unsigned)SIZE_T_MAX) ||
+        ((unsigned long long)loffset >> 3) >= (1024*1024*1024))
+    {
+        addReplyError(c,err);
+        return REDIS_ERR;
+    }
+
+    *offset = (size_t)loffset;
+    return REDIS_OK;
+}
+
+void setbitCommand(redisClient *c) {
+    robj *o;
+    size_t bitoffset;
+    int 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");
+        return;
+    }
+
+    o = lookupKeyWrite(c->db,c->argv[1]);
+    if (o == NULL) {
+        sds value = sdssetbit(sdsempty(),bitoffset,on);
+        o = createObject(REDIS_STRING,value);
+        dbAdd(c->db,c->argv[1],o);
+    } else {
+        if (checkType(c,o,REDIS_STRING)) return;
+
+        /* Create a copy when the object is shared or encoded. */
+        if (o->refcount != 1 || o->encoding != REDIS_ENCODING_RAW) {
+            robj *decoded = getDecodedObject(o);
+            o = createStringObject(decoded->ptr, sdslen(decoded->ptr));
+            decrRefCount(decoded);
+            dbReplace(c->db,c->argv[1],o);
+        }
+
+        o->ptr = sdssetbit(o->ptr,bitoffset,on);
+    }
+    touchWatchedKey(c->db,c->argv[1]);
+    server.dirty++;
+    addReply(c,shared.cone);
+}
+
+void getbitCommand(redisClient *c) {
+    robj *o;
+    size_t bitoffset, byte, bitmask;
+    int on = 0;
+    char llbuf[32];
+
+    if (getBitOffsetFromArgument(c,c->argv[2],&bitoffset) != REDIS_OK)
+        return;
+
+    if ((o = lookupKeyReadOrReply(c,c->argv[1],shared.czero)) == NULL ||
+        checkType(c,o,REDIS_STRING)) return;
+
+    byte = bitoffset >> 3;
+    bitmask = 1 << (7 - (bitoffset & 0x7));
+    if (o->encoding != REDIS_ENCODING_RAW) {
+        if (byte < (size_t)ll2string(llbuf,sizeof(llbuf),(long)o->ptr))
+            on = llbuf[byte] & bitmask;
+    } else {
+        if (byte < sdslen(o->ptr))
+            on = ((sds)o->ptr)[byte] & bitmask;
+    }
+    addReply(c, on ? shared.cone : shared.czero);
+}
+
 void mgetCommand(redisClient *c) {
     int j;
 

tests/unit/basic.tcl

@@ -374,4 +374,88 @@ start_server {tags {"basic"}} {
         r set mystring "foozzz0123456789 baz"
         r strlen mystring
     }
+
+    test "SETBIT against non-existing key" {
+        r del mykey
+
+        # Setting 2nd bit to on is integer 64, ascii "@"
+        assert_equal 1 [r setbit mykey 1 1]
+        assert_equal "@" [r get mykey]
+    }
+
+    test "SETBIT against string-encoded key" {
+        # Single byte with 2nd bit set
+        r set mykey "@"
+
+        # 64 + 32 = 96 => ascii "`" (backtick)
+        assert_equal 1 [r setbit mykey 2 1]
+        assert_equal "`" [r get mykey]
+    }
+
+    test "SETBIT against integer-encoded key" {
+        r set mykey 1
+        assert_encoding int mykey
+
+        # Ascii "1" is integer 49 = 00 11 00 01
+        # Setting 7th bit = 51 => ascii "3"
+        assert_equal 1 [r setbit mykey 6 1]
+        assert_equal "3" [r get mykey]
+    }
+
+    test "SETBIT against key with wrong type" {
+        r del mykey
+        r lpush mykey "foo"
+        assert_error "*wrong kind*" {r setbit mykey 0 1}
+    }
+
+    test "SETBIT with out of range bit offset" {
+        r del mykey
+        assert_error "*out of range*" {r setbit mykey [expr 8*1024*1024*1024] 1}
+        assert_error "*out of range*" {r setbit mykey -1 1}
+    }
+
+    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}
+    }
+
+    test "GETBIT against non-existing key" {
+        r del mykey
+        assert_equal 0 [r getbit mykey 0]
+    }
+
+    test "GETBIT against string-encoded key" {
+        # Single byte with 2nd and 3rd bit set
+        r set mykey "`"
+
+        # In-range
+        assert_equal 0 [r getbit mykey 0]
+        assert_equal 1 [r getbit mykey 1]
+        assert_equal 1 [r getbit mykey 2]
+        assert_equal 0 [r getbit mykey 3]
+
+        # Out-range
+        assert_equal 0 [r getbit mykey 8]
+        assert_equal 0 [r getbit mykey 100]
+        assert_equal 0 [r getbit mykey 10000]
+    }
+
+    test "GETBIT against integer-encoded key" {
+        r set mykey 1
+        assert_encoding int mykey
+
+        # Ascii "1" is integer 49 = 00 11 00 01
+        assert_equal 0 [r getbit mykey 0]
+        assert_equal 0 [r getbit mykey 1]
+        assert_equal 1 [r getbit mykey 2]
+        assert_equal 1 [r getbit mykey 3]
+
+        # Out-range
+        assert_equal 0 [r getbit mykey 8]
+        assert_equal 0 [r getbit mykey 100]
+        assert_equal 0 [r getbit mykey 10000]
+    }
 }