chore: Refactor CompactObject encoding (#5342)

Refactor CompactObject encoding to use StrEncoding helper object

Author: dranikpg

src/core/compact_object.cc

@@ -786,9 +786,7 @@ size_t CompactObj::Size() const {
     switch (taglen_) {
       case SMALL_TAG:
         raw_size = u_.small_str.size();
-        if (mask_bits_.encoding == HUFFMAN_ENC) {
-          return DecodedLen(raw_size, u_.small_str.first_byte());
-        }
+        first_byte = u_.small_str.first_byte();
         break;
       case INT_TAG: {
         absl::AlphaNum an(u_.ival);
@@ -801,9 +799,7 @@ size_t CompactObj::Size() const {
         break;
       case ROBJ_TAG:
         raw_size = u_.r_obj.Size();
-        if (mask_bits_.encoding == HUFFMAN_ENC) {
-          return DecodedLen(raw_size, *(uint8_t*)u_.r_obj.inner_obj());
-        }
+        first_byte = *(uint8_t*)u_.r_obj.inner_obj();
         break;
       case JSON_TAG:
         DCHECK_EQ(mask_bits_.encoding, NONE_ENC);
@@ -821,7 +817,7 @@ size_t CompactObj::Size() const {
         LOG(DFATAL) << "Should not reach " << int(taglen_);
     }
   }
-  return mask_bits_.encoding ? DecodedLen(raw_size, first_byte) : raw_size;
+  return GetStrEncoding().DecodedSize(raw_size, first_byte);
 }
 
 uint64_t CompactObj::HashCode() const {
@@ -848,16 +844,8 @@ uint64_t CompactObj::HashCode() const {
 
   if (IsInline()) {
     char buf[kInlineLen * 3];  // should suffice for most huffman decodings.
-    size_t decoded_len = DecodedLen(taglen_, u_.inline_str[0]);
-    if (mask_bits_.encoding == HUFFMAN_ENC) {
-      if (decoded_len <= sizeof(buf) &&
-          tl.huff_keys.decoder.Decode({u_.inline_str + 1, size_t(taglen_ - 1)}, decoded_len, buf)) {
-        return XXH3_64bits_withSeed(buf, decoded_len, kHashSeed);
-      }
-    } else {
-      detail::ascii_unpack(to_byte(u_.inline_str), decoded_len, buf);
-      return XXH3_64bits_withSeed(buf, decoded_len, kHashSeed);
-    }
+    size_t decoded_len = GetStrEncoding().Decode(string_view{u_.inline_str, taglen_}, buf);
+    return XXH3_64bits_withSeed(buf, decoded_len, kHashSeed);
   }
 
   string_view sv = GetSlice(&tl.tmp_str);
@@ -1114,21 +1102,7 @@ void CompactObj::GetString(char* dest) const {
   CHECK(!IsExternal());
 
   if (IsInline()) {
-    switch (mask_bits_.encoding) {
-      case ASCII2_ENC:
-        DCHECK_EQ(taglen_ + 2u, ascii_len(taglen_));
-        detail::ascii_unpack(to_byte(u_.inline_str), taglen_ + 2, dest);
-        break;
-      case HUFFMAN_ENC:
-        tl.huff_keys.decoder.Decode({u_.inline_str + 1, size_t(taglen_ - 1)},
-                                    u_.inline_str[0] + taglen_ - 1, dest);
-        break;
-      case NONE_ENC:
-        memcpy(dest, u_.inline_str, taglen_);
-        break;
-      default:
-        DLOG(FATAL) << "should not reach " << int(mask_bits_.encoding);
-    }
+    GetStrEncoding().Decode({u_.inline_str, taglen_}, dest);
     return;
   }
 
@@ -1142,19 +1116,15 @@ void CompactObj::GetString(char* dest) const {
     if (taglen_ == ROBJ_TAG) {
       CHECK_EQ(OBJ_STRING, u_.r_obj.type());
       DCHECK_EQ(OBJ_ENCODING_RAW, u_.r_obj.encoding());
-      size_t decoded_len = DecodedLen(u_.r_obj.Size(), *(const uint8_t*)u_.r_obj.inner_obj());
-      if (mask_bits_.encoding == HUFFMAN_ENC) {
-        CHECK(tl.huff_keys.decoder.Decode(
-            {(const char*)u_.r_obj.inner_obj() + 1, u_.r_obj.Size() - 1}, decoded_len, dest));
-        return;
-      }
-      detail::ascii_unpack_simd(to_byte(u_.r_obj.inner_obj()), decoded_len, dest);
+      string_view blob{(const char*)u_.r_obj.inner_obj(), u_.r_obj.Size()};
+      GetStrEncoding().Decode(blob, dest);
+      return;
     } else {
       CHECK_EQ(SMALL_TAG, taglen_);
       string_view slices[2];
       unsigned num = u_.small_str.GetV(slices);
       DCHECK_EQ(2u, num);
-      size_t decoded_len = DecodedLen(u_.small_str.size(), slices[0][0]);
+      size_t decoded_len = GetStrEncoding().DecodedSize(u_.small_str.size(), slices[0][0]);
 
       if (mask_bits_.encoding == HUFFMAN_ENC) {
         tl.tmp_buf.resize(slices[0].size() + slices[1].size() - 1);
@@ -1575,15 +1545,6 @@ StringOrView CompactObj::GetRawString() const {
   return {};
 }
 
-size_t CompactObj::DecodedLen(size_t sz, uint8_t b) const {
-  DCHECK(mask_bits_.encoding);
-  if (mask_bits_.encoding == HUFFMAN_ENC) {
-    return sz + b - 1;
-  }
-  unsigned delta = (mask_bits_.encoding == ASCII1_ENC) ? 1 : 0;
-  return ascii_len(sz) - delta;
-}
-
 MemoryResource* CompactObj::memory_resource() {
   return tl.local_mr;
 }
@@ -1613,4 +1574,52 @@ CompactObjType ObjTypeFromString(std::string_view sv) {
   return kInvalidCompactObjType;
 }
 
+size_t CompactObj::StrEncoding::DecodedSize(string_view blob) const {
+  return DecodedSize(blob.size(), blob[0]);
+}
+
+size_t CompactObj::StrEncoding::DecodedSize(size_t blob_size, uint8_t first_byte) const {
+  switch (enc_) {
+    case NONE_ENC:
+      return blob_size;
+    case ASCII1_ENC:
+    case ASCII2_ENC:
+      return ascii_len(blob_size) - (enc_ == ASCII1_ENC);
+    case HUFFMAN_ENC:
+      return blob_size + int(first_byte) - 1;
+  };
+  return 0;
+}
+
+size_t CompactObj::StrEncoding::Decode(std::string_view blob, char* dest) const {
+  size_t decoded_len = DecodedSize(blob);
+  switch (enc_) {
+    case NONE_ENC:
+      memcpy(dest, blob.data(), blob.size());
+      break;
+    case ASCII1_ENC:
+    case ASCII2_ENC:
+      detail::ascii_unpack(reinterpret_cast<const uint8_t*>(blob.data()), decoded_len, dest);
+      break;
+    case HUFFMAN_ENC:
+      tl.huff_keys.decoder.Decode(blob.substr(1), decoded_len, dest);
+      break;
+  };
+  return decoded_len;
+}
+
+StringOrView CompactObj::StrEncoding::Decode(std::string_view blob) const {
+  switch (enc_) {
+    case NONE_ENC:
+      return StringOrView::FromView(blob);
+    default: {
+      string out;
+      out.resize(DecodedSize(blob));
+      Decode(blob, out.data());
+      return StringOrView::FromString(std::move(out));
+    }
+  }
+  return {};
+}
+
 }  // namespace dfly

src/core/compact_object.h

@@ -130,14 +130,30 @@ class CompactObj {
   // Therefore, in order to know the original length we introduce 2 states that
   // correct the length upon decoding. ASCII1_ENC rounds down the decoded length,
   // while ASCII2_ENC rounds it up. See DecodedLen implementation for more info.
-  enum Encoding : uint8_t {
+  enum EncodingEnum : uint8_t {
     NONE_ENC = 0,
     ASCII1_ENC = 1,
     ASCII2_ENC = 2,
     HUFFMAN_ENC = 3,  // TBD
   };
 
  public:
+  // Utility class for working with different string encodings (ascii, huffman, etc)
+  struct StrEncoding {
+    size_t DecodedSize(std::string_view blob) const;         // Size of decoded blob
+    size_t Decode(std::string_view blob, char* dest) const;  // Decode into dest, return size
+    StringOrView Decode(std::string_view blob) const;
+
+   private:
+    friend class CompactObj;
+    explicit StrEncoding(uint8_t enc) : enc_(static_cast<EncodingEnum>(enc)) {
+    }
+
+    size_t DecodedSize(size_t compr_size, uint8_t first_byte) const;
+
+    EncodingEnum enc_;
+  };
+
   using PrefixArray = std::vector<std::string_view>;
   using MemoryResource = detail::RobjWrapper::MemoryResource;
 
@@ -406,6 +422,10 @@ class CompactObj {
   // Precondition: the object is a non-inline string.
   StringOrView GetRawString() const;
 
+  StrEncoding GetStrEncoding() const {
+    return StrEncoding{mask_bits_.encoding};
+  }
+
   bool HasAllocated() const;
 
   bool TagAllowsEmptyValue() const;
@@ -416,7 +436,6 @@ class CompactObj {
 
  private:
   void EncodeString(std::string_view str);
-  size_t DecodedLen(size_t sz, uint8_t firstb) const;
 
   bool EqualNonInline(std::string_view sv) const;
 

src/core/compact_object_test.cc

@@ -8,6 +8,7 @@
 #include <mimalloc.h>
 #include <xxhash.h>
 
+#include <cstddef>
 #include <random>
 
 #include "base/gtest.h"
@@ -16,6 +17,7 @@
 #include "core/flat_set.h"
 #include "core/huff_coder.h"
 #include "core/mi_memory_resource.h"
+#include "core/string_or_view.h"
 #include "core/string_set.h"
 
 extern "C" {
@@ -593,34 +595,28 @@ TEST_F(CompactObjectTest, DefragSet) {
   ASSERT_FALSE(cobj_.DefragIfNeeded(0.8));
 }
 
-TEST_F(CompactObjectTest, RawInterface) {
-  string str(50, 'a'), tmp, owned;
-  cobj_.SetString(str);
-  {
-    auto raw_blob = cobj_.GetRawString();
-    EXPECT_LT(raw_blob.view().size(), str.size());
+TEST_F(CompactObjectTest, StrEncodingAndMaterialize) {
+  for (bool ascii : {true, false}) {
+    for (size_t len : {64, 128, 256, 512, 1024}) {
+      string test_str(len, 'a');
+      for (size_t i = 0; i < len; i++)
+        test_str[i] = char('a' + (i % 10));
+      if (!ascii)
+        test_str.push_back(char(200));  // non-ascii
 
-    raw_blob.MakeOwned();
-    cobj_.SetExternal(0, 10);  // dummy external pointer
-    cobj_.Materialize(raw_blob.view(), true);
+      CompactObj obj;
+      obj.SetString(test_str);
 
-    EXPECT_EQ(str, cobj_.GetSlice(&tmp));
-  }
-
-  str.assign(50, char(200));  // non ascii
-  cobj_.SetString(str);
-  ASSERT_EQ(str, cobj_.GetSlice(&tmp));
-
-  {
-    auto raw_blob = cobj_.GetRawString();
+      // Test StrEncoding helper
+      string raw_str = obj.GetRawString().Take();
+      CompactObj::StrEncoding enc = obj.GetStrEncoding();
+      EXPECT_EQ(test_str, enc.Decode(raw_str).Take());
 
-    EXPECT_EQ(raw_blob.view(), str);
-
-    raw_blob.MakeOwned();
-    cobj_.SetExternal(0, 10);  // dummy external pointer
-    cobj_.Materialize(raw_blob.view(), true);
-
-    EXPECT_EQ(str, cobj_.GetSlice(&tmp));
+      // Test Materialize
+      obj.SetExternal(0, 0);  // dummy values
+      obj.Materialize(raw_str, true);
+      EXPECT_EQ(test_str, obj.ToString());
+    }
   }
 }
 

src/core/string_or_view.h

@@ -65,6 +65,11 @@ class StringOrView {
       val_ = std::string{std::get<std::string_view>(val_)};
   }
 
+  std::string Take() && {
+    MakeOwned();
+    return std::move(std::get<std::string>(val_));
+  }
+
   bool empty() const {
     return visit([](const auto& s) { return s.empty(); }, val_);
   }

src/server/tiered_storage.cc

@@ -60,16 +60,6 @@ void RecordDeleted(const PrimeValue& pv, size_t tiered_len, DbTableStats* stats)
   stats->tiered_used_bytes -= tiered_len;
 }
 
-string DecodeString(bool is_raw, string_view str, PrimeValue decoder) {
-  if (is_raw) {
-    decoder.Materialize(str, true);
-    string tmp;
-    decoder.GetString(&tmp);
-    return tmp;
-  }
-  return string{str};
-}
-
 tiering::DiskSegment FromCoolItem(const PrimeValue::CoolItem& item) {
   return {item.record->page_index * tiering::kPageSize + item.page_offset, item.serialized_size};
 }
@@ -346,13 +336,9 @@ void TieredStorage::Read(DbIndex dbid, std::string_view key, const PrimeValue& v
                          std::function<void(const std::string&)> readf) {
   DCHECK(value.IsExternal());
   DCHECK(!value.IsCool());
-
-  PrimeValue decoder;
-  decoder.ImportExternal(value);
-
-  auto cb = [readf = std::move(readf), decoder = std::move(decoder)](
+  auto cb = [readf = std::move(readf), enc = value.GetStrEncoding()](
                 bool is_raw, const string* raw_val) mutable {
-    readf(DecodeString(is_raw, *raw_val, std::move(decoder)));
+    readf(is_raw ? enc.Decode(*raw_val).Take() : *raw_val);
     return false;
   };
   op_manager_->Enqueue(KeyRef(dbid, key), value.GetExternalSlice(), std::move(cb));
@@ -365,14 +351,11 @@ util::fb2::Future<T> TieredStorage::Modify(DbIndex dbid, std::string_view key,
   DCHECK(value.IsExternal());
 
   util::fb2::Future<T> future;
-  PrimeValue decoder;
-  decoder.ImportExternal(value);
-
-  auto cb = [future, modf = std::move(modf), decoder = std::move(decoder)](
+  auto cb = [future, modf = std::move(modf), enc = value.GetStrEncoding()](
                 bool is_raw, std::string* raw_val) mutable {
     if (is_raw) {
-      decoder.Materialize(*raw_val, true);
-      decoder.GetString(raw_val);
+      raw_val->resize(enc.DecodedSize(*raw_val));
+      enc.Decode(*raw_val, raw_val->data());
     }
     future.Resolve(modf(raw_val));
     return true;