Support for the BITFIELD command

src/commands/cmd_bit.cc

@@ -225,10 +225,136 @@ class CommandBitOp : public Commander {
   BitOpFlags op_flag_;
 };
 
+class CommandBitfield : public Commander {
+ public:
+  Status Parse(const std::vector<std::string> &args) override {
+    BitfieldOperation cmd;
+
+    read_only_ = true;
+    // BITFIELD <key> [commands...]
+    for (auto group = args.begin() + 2; group != args.end();) {
+      auto size = std::distance(group, args.end());
+
+      std::string opcode = util::ToLower(group[0]);
+      if (opcode == "get") {
+        cmd.type = BitfieldOperation::Type::kGet;
+      } else if (opcode == "set") {
+        cmd.type = BitfieldOperation::Type::kSet;
+        read_only_ = false;
+      } else if (opcode == "incrby") {
+        cmd.type = BitfieldOperation::Type::kIncrBy;
+        read_only_ = false;
+      } else if (opcode == "overflow") {
+        constexpr auto kOverflowCmdSize = 2;
+        if (size < kOverflowCmdSize) {
+          return {Status::RedisParseErr, errWrongNumOfArguments};
+        }
+        std::string overflow = util::ToLower(group[1]);
+        if (overflow == "wrap") {
+          cmd.overflow = BitfieldOverflowBehavior::kWrap;
+        } else if (overflow == "sat") {
+          cmd.overflow = BitfieldOverflowBehavior::kSat;
+        } else if (overflow == "fail") {
+          cmd.overflow = BitfieldOverflowBehavior::kFail;
+        } else {
+          return {Status::RedisParseErr, errUnknownSubcommandOrWrongArguments};
+        }
+        group += kOverflowCmdSize;
+        continue;
+      } else {
+        return {Status::RedisParseErr, errUnknownSubcommandOrWrongArguments};
+      }
+
+      if (size < 3) {
+        return {Status::RedisParseErr, errWrongNumOfArguments};
+      }
+
+      // parse encoding
+      auto encoding = parseBitfieldEncoding(group[1]);
+      if (!encoding.IsOK()) {
+        return encoding.ToStatus();
+      }
+      cmd.encoding = encoding.GetValue();
+
+      // parse offset
+      if (!GetBitOffsetFromArgument(group[2], &cmd.offset).IsOK()) {
+        return {Status::RedisParseErr, "bit offset is not an integer or out of range"};
+      }
+
+      if (cmd.type != BitfieldOperation::Type::kGet) {
+        if (size < 4) {
+          return {Status::RedisParseErr, errWrongNumOfArguments};
+        }
+
+        auto value = ParseInt<int64_t>(group[3], 10);
+        if (!value.IsOK()) {
+          return value.ToStatus();
+        }
+        cmd.value = value.GetValue();
+      }
+
+      cmds_.push_back(cmd);
+    }
+
+    return Commander::Parse(args);
+  }
+
+  Status Execute(Server *srv, Connection *conn, std::string *output) override {
+    redis::Bitmap bitmap_db(srv->storage, conn->GetNamespace());
+    std::vector<std::string> rets;
+    rocksdb::Status s;
+    if (read_only_) {
+      s = bitmap_db.BitfieldReadOnly(args_[1], cmds_, &rets);
+    } else {
+      s = bitmap_db.Bitfield(args_[1], cmds_, &rets);
+    }
+    *output = redis::Array(rets);
+    return Status::OK();
+  }
+
+ private:
+  static StatusOr<BitfieldEncoding> parseBitfieldEncoding(const std::string &token) {
+    if (token.empty()) {
+      return {Status::RedisParseErr, errUnknownSubcommandOrWrongArguments};
+    }
+
+    auto sign = std::tolower(token[0]);
+    if (sign != 'u' && sign != 'i') {
+      return {Status::RedisParseErr, errUnknownSubcommandOrWrongArguments};
+    }
+
+    BitfieldEncoding enc;
+    enc.is_signed = false;
+    // u64 is unsupported for redis now.
+    // see also https://redis.io/commands/bitfield/
+    int16_t max_bits = 63;
+
+    if (sign == 'i') {
+      enc.is_signed = true;
+      max_bits = 64;
+    }
+
+    auto res = ParseInt<uint8_t>(token.substr(1), 10);
+    if (!res.IsOK()) {
+      return res.ToStatus();
+    }
+
+    enc.bits = res.GetValue();
+    if (enc.bits < 1 || max_bits < enc.bits) {
+      return {Status::RedisParseErr, errUnknownSubcommandOrWrongArguments};
+    }
+    return enc;
+  }
+
+  std::vector<BitfieldOperation> cmds_;
+  bool read_only_;
+};
+
 REDIS_REGISTER_COMMANDS(MakeCmdAttr<CommandGetBit>("getbit", 3, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandSetBit>("setbit", 4, "write", 1, 1, 1),
                         MakeCmdAttr<CommandBitCount>("bitcount", -2, "read-only", 1, 1, 1),
                         MakeCmdAttr<CommandBitPos>("bitpos", -3, "read-only", 1, 1, 1),
-                        MakeCmdAttr<CommandBitOp>("bitop", -4, "write", 2, -1, 1), )
+                        MakeCmdAttr<CommandBitOp>("bitop", -4, "write", 2, -1, 1),
+                        MakeCmdAttr<CommandBitfield>("bitfield", -2, "write", 1, 1, 1), )
 
 }  // namespace redis

src/common/bitfield_util.cc

@@ -0,0 +1,121 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ *
+ */
+
+#include "bitfield_util.h"
+
+namespace detail {
+namespace {
+// safe cast from unsigned to signed.
+// see also "Integral conversions" on https://en.cppreference.com/w/cpp/language/implicit_conversion
+// If the destination type is signed, the result when overflow is implementation-defined until C++20
+union SafeInt64 {
+  int64_t i;
+  uint64_t u;
+
+  explicit SafeInt64(int64_t i) noexcept : i(i) {}
+  explicit SafeInt64(uint64_t u) noexcept : u(u) {}
+};
+}  // namespace
+
+bool SignedBitfieldPlus(uint64_t value, int64_t incr, uint8_t bits, BitfieldOverflowBehavior overflow, uint64_t *dst) {
+  auto max = std::numeric_limits<int64_t>::max();
+  if (bits != 64) {
+    max = (static_cast<int64_t>(1) << (bits - 1)) - 1;
+  }
+  int64_t min = -max - 1;
+
+  SafeInt64 safe_value(value);
+
+  SafeInt64 max_incr(static_cast<uint64_t>(max) - safe_value.u);
+  SafeInt64 min_incr(min - safe_value.i);
+
+  auto wrap = [](uint64_t value, int64_t incr, uint8_t bits) {
+    uint64_t res = value + static_cast<uint64_t>(incr);
+    if (bits < 64) {
+      auto mask = std::numeric_limits<uint64_t>::max() << bits;
+      if ((res & (1 << (bits - 1))) != 0) {
+        res |= mask;
+      } else {
+        res &= ~mask;
+      }
+    }
+    return res;
+  };
+
+  if (safe_value.i > max || (bits != 64 && incr > max_incr.i) ||
+      (safe_value.i >= 0 && incr >= 0 && incr > max_incr.i)) {
+    if (overflow == BitfieldOverflowBehavior::kWrap) {
+      *dst = wrap(safe_value.u, incr, bits);
+    } else {
+      *dst = max;
+    }
+    return true;
+  } else if (safe_value.i < min || (bits != 64 && incr < min_incr.i) ||
+             (safe_value.i < 0 && incr < 0 && incr < min_incr.i)) {
+    if (overflow == BitfieldOverflowBehavior::kWrap) {
+      *dst = wrap(safe_value.u, incr, bits);
+    } else {
+      *dst = min;
+    }
+    return true;
+  }
+
+  *dst = safe_value.i + incr;
+  return false;
+}
+
+// return true if overflow.
+bool UnsignedBitfieldPlus(uint64_t value, int64_t incr, uint8_t bits, BitfieldOverflowBehavior overflow,
+                          uint64_t *dst) {
+  auto max = std::numeric_limits<uint64_t>::max();
+  if (bits != 64) {
+    max = (static_cast<uint64_t>(1) << bits) - 1;
+  }
+
+  SafeInt64 max_incr(max - value);
+  SafeInt64 min_incr(-value);
+
+  auto wrap = [](uint64_t value, int64_t incr, uint8_t bits) {
+    uint64_t mask = std::numeric_limits<uint64_t>::max() << bits;
+    uint64_t res = value + incr;
+    res &= ~mask;
+    return res;
+  };
+
+  if (value > max || (incr > 0 && incr > max_incr.i)) {
+    if (overflow == BitfieldOverflowBehavior::kWrap) {
+      *dst = wrap(value, incr, bits);
+    } else if (overflow == BitfieldOverflowBehavior::kSat) {
+      *dst = max;
+    }
+    return true;
+  } else if (incr < 0 && incr < min_incr.i) {
+    if (overflow == BitfieldOverflowBehavior::kWrap) {
+      *dst = wrap(value, incr, bits);
+    } else if (overflow == BitfieldOverflowBehavior::kSat) {
+      *dst = 0;
+    }
+    return true;
+  }
+
+  *dst = value + incr;
+  return false;
+}
+}  // namespace detail

src/common/bitfield_util.h

@@ -0,0 +1,123 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ *
+ */
+
+#pragma once
+
+#include <cstdint>
+#include <limits>
+#include <type_traits>
+#include <utility>
+
+enum class BitfieldOverflowBehavior : uint8_t { kWrap, kSat, kFail };
+
+struct BitfieldEncoding {
+  bool is_signed;
+  // 1 ~ 63 if unsigned, 1 ~ 64 if signed.
+  uint8_t bits;
+};
+
+struct BitfieldOperation {
+  // see https://redis.io/commands/bitfield/ to get more details.
+  enum class Type : uint8_t { kGet, kSet, kIncrBy };
+
+  Type type;
+  BitfieldOverflowBehavior overflow{BitfieldOverflowBehavior::kWrap};
+  BitfieldEncoding encoding;
+  uint32_t offset;
+  // INCRBY amount or SET value
+  int64_t value;
+};
+
+namespace detail {
+// return true if overflow
+bool SignedBitfieldPlus(uint64_t value, int64_t incr, uint8_t bits, BitfieldOverflowBehavior overflow, uint64_t *dst);
+
+// return true if overflow.
+bool UnsignedBitfieldPlus(uint64_t value, int64_t incr, uint8_t bits, BitfieldOverflowBehavior overflow, uint64_t *dst);
+}  // namespace detail
+
+// return true if overflow.
+inline auto BitfieldPlus(uint64_t value, int64_t incr, BitfieldEncoding enc, BitfieldOverflowBehavior overflow,
+                         uint64_t *dst) {
+  if (enc.is_signed) {
+    return detail::SignedBitfieldPlus(value, incr, enc.bits, overflow, dst);
+  }
+  return detail::UnsignedBitfieldPlus(value, incr, enc.bits, overflow, dst);
+}
+
+// return true if successful
+inline bool BitfieldOp(BitfieldOperation op, uint64_t old_value, uint64_t *new_value) {
+  if (op.type == BitfieldOperation::Type::kGet) {
+    *new_value = old_value;
+    return true;
+  }
+
+  bool overflow = false;
+  if (op.type == BitfieldOperation::Type::kSet) {
+    overflow = BitfieldPlus(op.value, 0, op.encoding, op.overflow, new_value);
+  } else {
+    overflow = BitfieldPlus(old_value, op.value, op.encoding, op.overflow, new_value);
+  }
+
+  return op.overflow != BitfieldOverflowBehavior::kFail || !overflow;
+}
+
+template <class BitmapView>
+inline void SetBitfield(BitmapView view, uint32_t offset, uint32_t bits, uint64_t value) {
+  while (bits--) {
+    bool v = (value & (uint64_t(1) << bits)) != 0;
+    uint32_t byte = offset >> 3;
+    uint32_t bit = 7 - (offset & 7);
+    uint32_t byteval = view[byte];
+    byteval &= ~(1 << bit);
+    byteval |= int(v) << bit;
+    view[byte] = char(byteval);
+    offset++;
+  }
+}
+
+template <class BitmapView>
+inline uint64_t GetUnsignedBitfield(BitmapView &&view, uint64_t offset, uint64_t bits) {
+  uint64_t value = 0;
+  while (bits--) {
+    uint32_t byte = offset >> 3;
+    uint32_t bit = 7 - (offset & 0x7);
+    uint32_t byteval = view[byte];
+    uint32_t bitval = (byteval >> bit) & 1;
+    value = (value << 1) | bitval;
+    offset++;
+  }
+  return value;
+}
+
+template <class BitmapView>
+inline int64_t GetSignedBitfield(BitmapView &&view, uint64_t offset, uint64_t bits) {
+  union {
+    uint64_t u;
+    int64_t i;
+  } conv;
+  conv.u = GetUnsignedBitfield(std::forward<BitmapView>(view), offset, bits);
+  int64_t value = conv.i;
+
+  if ((value & (static_cast<int64_t>(1) << (bits - 1))) != 0) {
+    value |= static_cast<int64_t>(-1) << bits;
+  }
+  return value;
+}