Add CHOP command to rxstrings

README.md

@@ -66,6 +66,18 @@ Note: the key shouldn't be repeated for the executed command.
 **Reply:** Null if not equal or for non existing key when the `XX` flag is used.
 On success, the reply depends on the actual command executed.
 
+## `CHOP key count`
+
+Removes characters from a value of a String key.
+
+* If 'count' is positive, it removes from the right of the string;
+* if 'count' is negative, it removes from the left of the string.
+* If |count| is greater than the string length, the value is set as an empty string.
+* If 'count' is zero, then the value remains unchanged.
+* It is an error if the value is not a String.
+
+**Reply:** Integer, the length of the string after the chop operation.
+
 ## `PREPEND key value`
 
 > Time complexity: O(1). The amortized time complexity is O(1) assuming the prepended value is small and the already present value is of any size, since the dynamic string library used by Redis will double the free space available on every reallocation.

src/rxstrings.c

@@ -135,6 +135,76 @@ int CheckAndCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
   return REDISMODULE_OK;
 }
 
+/*
+* CHOP key count
+* Removes characters from a value of a String key.
+* If 'count' is positive, it removes from the right of the string;
+* if 'count' is negative, it removes from the left of the string.
+* If |count| is greater than the string length, the value is set as an empty string.
+* If 'count' is zero, then the value remains unchanged.
+* It is an error if the value is not a String.
+* Integer Reply: the length of the string after the chop operation.
+*/
+int ChopCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
+  if (argc != 3) {
+    return RedisModule_WrongArity(ctx);
+  }
+  RedisModule_AutoMemory(ctx);
+
+  /* Obtain key and extract offset argument. */
+  RedisModuleKey *key =
+      RedisModule_OpenKey(ctx, argv[1], REDISMODULE_READ | REDISMODULE_WRITE);
+  long long count;
+  if (RedisModule_StringToLongLong(argv[2], &count) != REDISMODULE_OK) {
+    RedisModule_ReplyWithError(ctx, "ERR invalid count");
+    return REDISMODULE_ERR;
+  }
+
+  /* Key must be a string. */
+  if (RedisModule_KeyType(key) != REDISMODULE_KEYTYPE_STRING) {
+    RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE);
+    return REDISMODULE_ERR;
+  }
+
+  /* Calculate new length, protecting against overchop */
+  size_t valLen = RedisModule_ValueLength(key);
+  if (count == 0) {
+    RedisModule_ReplyWithLongLong(ctx, valLen);
+    return REDISMODULE_OK;
+  }
+  size_t absCount = (count < 0) ? -count : count;
+  size_t newLen = (valLen < absCount) ? 0 : (valLen - absCount);
+
+  /* Work around https://github.com/antirez/redis/issues/3717
+   * We use RM_StringSet to an empty string instead */
+  if (newLen == 0) {
+    RedisModuleString* emptyStr = RedisModule_CreateString(ctx, "", 0);
+    if (RedisModule_StringSet(key, emptyStr) != REDISMODULE_OK) {
+      RedisModule_ReplyWithError(ctx, "ERR RM_SetString empty failed");
+      return REDISMODULE_ERR;
+    }
+    RedisModule_ReplyWithLongLong(ctx, 0);
+    return REDISMODULE_OK;
+  }
+
+  /* If we chop from the left, we need to shift the characters,
+   * but only bother if we aren't zeroing the value */
+  if (count < 0 && newLen != 0) {
+    size_t dmaLen;
+    char* valPtr = RedisModule_StringDMA(key, &dmaLen, REDISMODULE_READ|REDISMODULE_WRITE);
+    memmove(valPtr, valPtr + absCount, newLen);
+  }
+
+  /* Truncate to new length */
+  if (RedisModule_StringTruncate(key, newLen) != REDISMODULE_OK) {
+    RedisModule_ReplyWithError(ctx, "ERR RM_StringTruncate failed");
+    return REDISMODULE_ERR;
+  }
+
+  RedisModule_ReplyWithLongLong(ctx, newLen);
+  return REDISMODULE_OK;
+}
+
 /*
 * PREPEND key value
 * Prepends a value to a String key.
@@ -167,7 +237,7 @@ int PrependCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
   /* Otherwise key must be a string. */
   if (RedisModule_KeyType(key) != REDISMODULE_KEYTYPE_STRING) {
     RedisModule_ReplyWithError(ctx, REDISMODULE_ERRORMSG_WRONGTYPE);
-    return REDISMODULE_OK;
+    return REDISMODULE_ERR;
   }
 
   /* Prepend the string: 1) expand string, 2) shift oldVal via memmove, 3) prepend arg */
@@ -409,6 +479,52 @@ int testCheckAnd(RedisModuleCtx *ctx) {
   return 0;
 }
 
+int testChop(RedisModuleCtx *ctx) {
+  RedisModuleCallReply *r;
+
+  r = RedisModule_Call(ctx, "set", "cc", "foo", "abcde");
+  RMUtil_AssertReplyEquals(r,"OK");
+  r = RedisModule_Call(ctx, "chop", "cc", "foo", "-2");
+  RMUtil_Assert(RedisModule_CallReplyInteger(r) == 3);
+  r = RedisModule_Call(ctx, "get", "c", "foo");
+  RMUtil_AssertReplyEquals(r,"cde");
+  r = RedisModule_Call(ctx, "FLUSHALL", "");
+
+  r = RedisModule_Call(ctx, "set", "cc", "foo", "abcde");
+  RMUtil_AssertReplyEquals(r,"OK");
+  r = RedisModule_Call(ctx, "chop", "cc", "foo", "2");
+  RMUtil_Assert(RedisModule_CallReplyInteger(r) == 3);
+  r = RedisModule_Call(ctx, "get", "c", "foo");
+  RMUtil_AssertReplyEquals(r,"abc");
+  r = RedisModule_Call(ctx, "FLUSHALL", "");
+
+  r = RedisModule_Call(ctx, "set", "cc", "foo", "abcde");
+  RMUtil_AssertReplyEquals(r,"OK");
+  r = RedisModule_Call(ctx, "chop", "cc", "foo", "-10");
+  RMUtil_Assert(RedisModule_CallReplyInteger(r) == 0);
+  r = RedisModule_Call(ctx, "get", "c", "foo");
+  RMUtil_AssertReplyEquals(r,"");
+  r = RedisModule_Call(ctx, "FLUSHALL", "");
+
+  r = RedisModule_Call(ctx, "set", "cc", "foo", "abcde");
+  RMUtil_AssertReplyEquals(r,"OK");
+  r = RedisModule_Call(ctx, "chop", "cc", "foo", "10");
+  RMUtil_Assert(RedisModule_CallReplyInteger(r) == 0);
+  r = RedisModule_Call(ctx, "get", "c", "foo");
+  RMUtil_AssertReplyEquals(r,"");
+  r = RedisModule_Call(ctx, "FLUSHALL", "");
+
+  r = RedisModule_Call(ctx, "set", "cc", "foo", "abcde");
+  RMUtil_AssertReplyEquals(r,"OK");
+  r = RedisModule_Call(ctx, "chop", "cc", "foo", "0");
+  RMUtil_Assert(RedisModule_CallReplyInteger(r) == 5);
+  r = RedisModule_Call(ctx, "get", "c", "foo");
+  RMUtil_AssertReplyEquals(r,"abcde");
+  r = RedisModule_Call(ctx, "FLUSHALL", "");
+
+  return 0;
+}
+
 int testPrepend(RedisModuleCtx *ctx) {
   RedisModuleCallReply *r;
 
@@ -445,6 +561,7 @@ int TestModule(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
   }
 
   RMUtil_Test(testCheckAnd);
+  RMUtil_Test(testChop);
   RMUtil_Test(testPrepend);
   RMUtil_Test(testSetRangeRand);
 
@@ -460,6 +577,10 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx) {
   if (RedisModule_CreateCommand(ctx, "checkand", CheckAndCommand,
                                 "write deny-oom", 1, 1, 1) == REDISMODULE_ERR)
     return REDISMODULE_ERR;
+  if (RedisModule_CreateCommand(ctx, "chop", ChopCommand,
+                                "write fast deny-oom", 1, 1,
+                                1) == REDISMODULE_ERR)
+    return REDISMODULE_ERR;
   if (RedisModule_CreateCommand(ctx, "prepend", PrependCommand,
                                 "write fast deny-oom", 1, 1,
                                 1) == REDISMODULE_ERR)
@@ -473,4 +594,4 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx) {
     return REDISMODULE_ERR;
 
   return REDISMODULE_OK;
-}
+}