Add commands SETBIT/GETBIT

[redis.git] / src / t_string.c

diff --git a/src/t_string.c b/src/t_string.c
index 3b8a39bbec7100d20622411c2376eda49565a2fa..4b6fe79208736c98940c6cf6ef44c99fb4670c6a 100644 (file)
--- a/src/t_string.c
+++ b/src/t_string.c
@@ -1,3 +1,4 @@
+#include <limits.h>
 #include "redis.h"
 
 /*-----------------------------------------------------------------------------
@@ -12,7 +13,7 @@ void setGenericCommand(redisClient *c, int nx, robj *key, robj *val, robj *expir
         if (getLongFromObjectOrReply(c, expire, &seconds, NULL) != REDIS_OK)
             return;
         if (seconds <= 0) {
-            addReplySds(c,sdsnew("-ERR invalid expire time in SETEX\r\n"));
+            addReplyError(c,"invalid expire time in SETEX");
             return;
         }
     }
@@ -37,14 +38,17 @@ void setGenericCommand(redisClient *c, int nx, robj *key, robj *val, robj *expir
 }
 
 void setCommand(redisClient *c) {
+    c->argv[2] = tryObjectEncoding(c->argv[2]);
     setGenericCommand(c,0,c->argv[1],c->argv[2],NULL);
 }
 
 void setnxCommand(redisClient *c) {
+    c->argv[2] = tryObjectEncoding(c->argv[2]);
     setGenericCommand(c,1,c->argv[1],c->argv[2],NULL);
 }
 
 void setexCommand(redisClient *c) {
+    c->argv[3] = tryObjectEncoding(c->argv[3]);
     setGenericCommand(c,0,c->argv[1],c->argv[3],c->argv[2]);
 }
 
@@ -69,6 +73,7 @@ void getCommand(redisClient *c) {
 
 void getsetCommand(redisClient *c) {
     if (getGenericCommand(c) == REDIS_ERR) return;
+    c->argv[2] = tryObjectEncoding(c->argv[2]);
     dbReplace(c->db,c->argv[1],c->argv[2]);
     incrRefCount(c->argv[2]);
     touchWatchedKey(c->db,c->argv[1]);
@@ -76,10 +81,91 @@ 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;
 
-    addReplySds(c,sdscatprintf(sdsempty(),"*%d\r\n",c->argc-1));
+    addReplyMultiBulkLen(c,c->argc-1);
     for (j = 1; j < c->argc; j++) {
         robj *o = lookupKeyRead(c->db,c->argv[j]);
         if (o == NULL) {
@@ -98,7 +184,7 @@ void msetGenericCommand(redisClient *c, int nx) {
     int j, busykeys = 0;
 
     if ((c->argc % 2) == 0) {
-        addReplySds(c,sdsnew("-ERR wrong number of arguments for MSET\r\n"));
+        addReplyError(c,"wrong number of arguments for MSET");
         return;
     }
     /* Handle the NX flag. The MSETNX semantic is to return zero and don't
@@ -180,6 +266,7 @@ void appendCommand(redisClient *c) {
     robj *o;
 
     o = lookupKeyWrite(c->db,c->argv[1]);
+    c->argv[2] = tryObjectEncoding(c->argv[2]);
     if (o == NULL) {
         /* Create the key */
         retval = dbAdd(c->db,c->argv[1],c->argv[2]);
@@ -211,7 +298,7 @@ void appendCommand(redisClient *c) {
     }
     touchWatchedKey(c->db,c->argv[1]);
     server.dirty++;
-    addReplySds(c,sdscatprintf(sdsempty(),":%lu\r\n",(unsigned long)totlen));
+    addReplyLongLong(c,totlen);
 }
 
 void substrCommand(redisClient *c) {