mirror of
https://github.com/facebook/rocksdb.git
synced 2026-07-07 14:47:40 +08:00
Add prefix varint codec and tests (#14692)
Summary: Add MLIR-compatible PrefixVarint32/64 helpers in util/prefix_varint.h. Document the format, split decode API, and hot paths, and cover the codec in coding_test with round-trip, disk-read, overflow, and truncation cases. I'm intending to use this in the new blog file format because you only need to read the first byte to know how many varint bytes to read total, and it might be more CPU efficient in other uses as well (to be determined in follow-up work). Pull Request resolved: https://github.com/facebook/rocksdb/pull/14692 Test Plan: Unit tests included Reviewed By: joshkang97 Differential Revision: D103245079 Pulled By: pdillinger fbshipit-source-id: dca435beccb6e666d864a31d8683ad19e2121c34
This commit is contained in:
Peter Dillinger
committed by
meta-codesync[bot]
meta-codesync[bot]
parent
a4045b0a5f
commit
c99d5ec06e
+366
-2
@@ -9,11 +9,15 @@
|
||||
|
||||
#include "util/coding.h"
|
||||
|
||||
#include <initializer_list>
|
||||
|
||||
#include "test_util/testharness.h"
|
||||
#include "util/prefix_varint.h"
|
||||
|
||||
namespace ROCKSDB_NAMESPACE {
|
||||
|
||||
class Coding {};
|
||||
|
||||
TEST(Coding, Fixed16) {
|
||||
std::string s;
|
||||
for (uint16_t v = 0; v < 0xFFFF; v++) {
|
||||
@@ -119,8 +123,8 @@ TEST(Coding, Varint64) {
|
||||
// Some special values
|
||||
values.push_back(0);
|
||||
values.push_back(100);
|
||||
values.push_back(~static_cast<uint64_t>(0));
|
||||
values.push_back(~static_cast<uint64_t>(0) - 1);
|
||||
values.push_back(~uint64_t{0});
|
||||
values.push_back(~uint64_t{0} - 1);
|
||||
for (uint32_t k = 0; k < 64; k++) {
|
||||
// Test values near powers of two
|
||||
const uint64_t power = 1ull << k;
|
||||
@@ -208,6 +212,366 @@ TEST(Coding, Strings) {
|
||||
ASSERT_EQ("", input.ToString());
|
||||
}
|
||||
|
||||
// Keep PrefixVarint-specific helpers and vectors with the PrefixVarint tests
|
||||
// below so the legacy coding tests stay grouped above.
|
||||
namespace {
|
||||
|
||||
template <typename T>
|
||||
struct PrefixVarintTestCase {
|
||||
T value;
|
||||
uint32_t length;
|
||||
std::string encoded;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
struct PrefixVarintTraits;
|
||||
|
||||
template <>
|
||||
struct PrefixVarintTraits<uint32_t> {
|
||||
static constexpr size_t kMaxLength = kMaxPrefixVarint32Length;
|
||||
|
||||
static uint32_t Length(uint32_t value) { return PrefixVarint32Length(value); }
|
||||
static char* Encode(char* dst, uint32_t value) {
|
||||
return EncodePrefixVarint32(dst, value);
|
||||
}
|
||||
static void Put(std::string* dst, uint32_t value) {
|
||||
PutPrefixVarint32(dst, value);
|
||||
}
|
||||
static const char* GetPtr(const char* p, const char* limit, uint32_t* value) {
|
||||
return GetPrefixVarint32Ptr(p, limit, value);
|
||||
}
|
||||
static bool Get(Slice* input, uint32_t* value) {
|
||||
return GetPrefixVarint32(input, value);
|
||||
}
|
||||
static uint32_t AddlByteCount(char first_byte) {
|
||||
return PrefixVarint32AddlByteCount(first_byte);
|
||||
}
|
||||
static bool Decode(char first_byte, const char* addl_bytes,
|
||||
size_t addl_byte_count, uint32_t* value) {
|
||||
return DecodePrefixVarint32(first_byte, addl_bytes, addl_byte_count, value);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct PrefixVarintTraits<uint64_t> {
|
||||
static constexpr size_t kMaxLength = kMaxPrefixVarint64Length;
|
||||
|
||||
static uint32_t Length(uint64_t value) { return PrefixVarint64Length(value); }
|
||||
static char* Encode(char* dst, uint64_t value) {
|
||||
return EncodePrefixVarint64(dst, value);
|
||||
}
|
||||
static void Put(std::string* dst, uint64_t value) {
|
||||
PutPrefixVarint64(dst, value);
|
||||
}
|
||||
static const char* GetPtr(const char* p, const char* limit, uint64_t* value) {
|
||||
return GetPrefixVarint64Ptr(p, limit, value);
|
||||
}
|
||||
static bool Get(Slice* input, uint64_t* value) {
|
||||
return GetPrefixVarint64(input, value);
|
||||
}
|
||||
static uint32_t AddlByteCount(char first_byte) {
|
||||
return PrefixVarint64AddlByteCount(first_byte);
|
||||
}
|
||||
static bool Decode(char first_byte, const char* addl_bytes,
|
||||
size_t addl_byte_count, uint64_t* value) {
|
||||
return DecodePrefixVarint64(first_byte, addl_bytes, addl_byte_count, value);
|
||||
}
|
||||
};
|
||||
|
||||
std::string ByteString(std::initializer_list<unsigned char> bytes) {
|
||||
std::string out;
|
||||
out.reserve(bytes.size());
|
||||
for (unsigned char b : bytes) {
|
||||
out.push_back(static_cast<char>(b));
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
std::string EncodePrefixVarintToString(T value) {
|
||||
char buf[PrefixVarintTraits<T>::kMaxLength];
|
||||
char* end = PrefixVarintTraits<T>::Encode(buf, value);
|
||||
return std::string(buf, static_cast<size_t>(end - buf));
|
||||
}
|
||||
|
||||
std::string EncodeInvalidPrefixVarint32Overflow() {
|
||||
char buf[kMaxPrefixVarint32Length];
|
||||
uint64_t encoded = (uint64_t{1} << 37) | 0x10u;
|
||||
for (size_t i = 0; i < sizeof(buf); ++i) {
|
||||
buf[i] = static_cast<char>(encoded >> (i * 8));
|
||||
}
|
||||
return std::string(buf, sizeof(buf));
|
||||
}
|
||||
|
||||
template <typename T, size_t N>
|
||||
void AssertPrefixVarintRoundTrip(const PrefixVarintTestCase<T> (&cases)[N]) {
|
||||
std::string encoded_values;
|
||||
for (const auto& tc : cases) {
|
||||
ASSERT_EQ(tc.length, PrefixVarintTraits<T>::Length(tc.value));
|
||||
ASSERT_EQ(tc.encoded, EncodePrefixVarintToString(tc.value));
|
||||
PrefixVarintTraits<T>::Put(&encoded_values, tc.value);
|
||||
}
|
||||
|
||||
const char* p = encoded_values.data();
|
||||
const char* limit = p + encoded_values.size();
|
||||
for (const auto& tc : cases) {
|
||||
T actual = 0;
|
||||
const char* start = p;
|
||||
p = PrefixVarintTraits<T>::GetPtr(p, limit, &actual);
|
||||
ASSERT_TRUE(p != nullptr);
|
||||
ASSERT_EQ(tc.value, actual);
|
||||
ASSERT_EQ(tc.length, p - start);
|
||||
}
|
||||
ASSERT_EQ(p, limit);
|
||||
|
||||
Slice input(encoded_values);
|
||||
for (const auto& tc : cases) {
|
||||
T actual = 0;
|
||||
ASSERT_TRUE(PrefixVarintTraits<T>::Get(&input, &actual));
|
||||
ASSERT_EQ(tc.value, actual);
|
||||
}
|
||||
ASSERT_TRUE(input.empty());
|
||||
}
|
||||
|
||||
template <typename T, size_t N>
|
||||
void AssertPrefixVarintDiskReadApi(const PrefixVarintTestCase<T> (&cases)[N]) {
|
||||
for (const auto& tc : cases) {
|
||||
const std::string encoded = EncodePrefixVarintToString(tc.value);
|
||||
ASSERT_EQ(tc.length, encoded.size());
|
||||
|
||||
const uint32_t addl_byte_count =
|
||||
PrefixVarintTraits<T>::AddlByteCount(encoded[0]);
|
||||
ASSERT_EQ(tc.length - 1, addl_byte_count);
|
||||
|
||||
T actual = 0;
|
||||
ASSERT_TRUE(PrefixVarintTraits<T>::Decode(encoded[0], encoded.data() + 1,
|
||||
addl_byte_count, &actual));
|
||||
ASSERT_EQ(tc.value, actual);
|
||||
if (addl_byte_count > 0) {
|
||||
ASSERT_FALSE(PrefixVarintTraits<T>::Decode(encoded[0], encoded.data() + 1,
|
||||
addl_byte_count - 1, &actual));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void AssertPrefixVarintTruncation(T value) {
|
||||
const std::string encoded = EncodePrefixVarintToString(value);
|
||||
T actual = 0;
|
||||
for (size_t len = 0; len + 1 < encoded.size(); ++len) {
|
||||
ASSERT_TRUE(PrefixVarintTraits<T>::GetPtr(
|
||||
encoded.data(), encoded.data() + len, &actual) == nullptr);
|
||||
}
|
||||
ASSERT_TRUE(PrefixVarintTraits<T>::GetPtr(encoded.data(),
|
||||
encoded.data() + encoded.size(),
|
||||
&actual) != nullptr);
|
||||
ASSERT_EQ(value, actual);
|
||||
}
|
||||
|
||||
const PrefixVarintTestCase<uint32_t> kPrefixVarint32TestCases[] = {
|
||||
{0, 1, ByteString({0x01})},
|
||||
{1, 1, ByteString({0x03})},
|
||||
{127, 1, ByteString({0xFF})},
|
||||
{128, 2, ByteString({0x02, 0x02})},
|
||||
{255, 2, ByteString({0xFE, 0x03})},
|
||||
{16383, 2, ByteString({0xFE, 0xFF})},
|
||||
{16384, 3, ByteString({0x04, 0x00, 0x02})},
|
||||
{(1u << 21) - 1, 3, ByteString({0xFC, 0xFF, 0xFF})},
|
||||
{(1u << 21), 4, ByteString({0x08, 0x00, 0x00, 0x02})},
|
||||
{(1u << 28) - 1, 4, ByteString({0xF8, 0xFF, 0xFF, 0xFF})},
|
||||
{(1u << 28), 5, ByteString({0x10, 0x00, 0x00, 0x00, 0x02})},
|
||||
{~uint32_t{0}, 5, ByteString({0xF0, 0xFF, 0xFF, 0xFF, 0x1F})},
|
||||
};
|
||||
|
||||
const PrefixVarintTestCase<uint64_t> kPrefixVarint64TestCases[] = {
|
||||
{0, 1, ByteString({0x01})},
|
||||
{1, 1, ByteString({0x03})},
|
||||
{127, 1, ByteString({0xFF})},
|
||||
{128, 2, ByteString({0x02, 0x02})},
|
||||
{16383, 2, ByteString({0xFE, 0xFF})},
|
||||
{(uint64_t{1} << 21), 4, ByteString({0x08, 0x00, 0x00, 0x02})},
|
||||
{(uint64_t{1} << 28), 5, ByteString({0x10, 0x00, 0x00, 0x00, 0x02})},
|
||||
{(uint64_t{1} << 35), 6, ByteString({0x20, 0x00, 0x00, 0x00, 0x00, 0x02})},
|
||||
{(uint64_t{1} << 42), 7,
|
||||
ByteString({0x40, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02})},
|
||||
{(uint64_t{1} << 49), 8,
|
||||
ByteString({0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02})},
|
||||
{(uint64_t{1} << 56) - 1, 8,
|
||||
ByteString({0x80, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF})},
|
||||
{(uint64_t{1} << 56), 9,
|
||||
ByteString({0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01})},
|
||||
{~uint64_t{0}, 9,
|
||||
ByteString({0x00, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF})},
|
||||
};
|
||||
|
||||
template <uint32_t kMinimumBytes>
|
||||
void TestPrefixVarint64MinimumBytes(const std::vector<uint64_t>& values) {
|
||||
for (uint64_t value : values) {
|
||||
SCOPED_TRACE("value=" + std::to_string(value) +
|
||||
" kMinimumBytes=" + std::to_string(kMinimumBytes));
|
||||
char buf[kMaxPrefixVarint64Length];
|
||||
char* end = EncodePrefixVarint64<kMinimumBytes>(buf, value);
|
||||
uint32_t encoded_len = static_cast<uint32_t>(end - buf);
|
||||
uint32_t natural_len = PrefixVarint64Length(value);
|
||||
ASSERT_EQ(encoded_len, std::max(natural_len, kMinimumBytes));
|
||||
|
||||
// Verify via split-decode API
|
||||
uint32_t addl = PrefixVarint64AddlByteCount(buf[0]);
|
||||
ASSERT_EQ(addl + 1, encoded_len);
|
||||
uint64_t decoded = ~uint64_t{0};
|
||||
ASSERT_TRUE(DecodePrefixVarint64(buf[0], buf + 1, addl, &decoded));
|
||||
ASSERT_EQ(value, decoded);
|
||||
|
||||
// Verify via GetPtr API
|
||||
decoded = ~uint64_t{0};
|
||||
const char* next = GetPrefixVarint64Ptr(buf, end, &decoded);
|
||||
ASSERT_EQ(next, end);
|
||||
ASSERT_EQ(value, decoded);
|
||||
|
||||
// Verify via Slice API
|
||||
decoded = ~uint64_t{0};
|
||||
Slice slice(buf, encoded_len);
|
||||
ASSERT_TRUE(GetPrefixVarint64(&slice, &decoded));
|
||||
ASSERT_EQ(value, decoded);
|
||||
ASSERT_TRUE(slice.empty());
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// Keep PrefixVarint tests after the legacy coding tests so the two varint
|
||||
// families stay visually separated.
|
||||
TEST(Coding, PrefixVarint32) {
|
||||
for (const auto& tc : kPrefixVarint32TestCases) {
|
||||
ASSERT_EQ(tc.length, VarintLength(tc.value));
|
||||
}
|
||||
AssertPrefixVarintRoundTrip(kPrefixVarint32TestCases);
|
||||
}
|
||||
|
||||
TEST(Coding, PrefixVarint32DiskReadApi) {
|
||||
AssertPrefixVarintDiskReadApi(kPrefixVarint32TestCases);
|
||||
ASSERT_EQ(kInvalidPrefixVarint32AddlByteCount,
|
||||
PrefixVarint32AddlByteCount('\0'));
|
||||
ASSERT_EQ(kInvalidPrefixVarint32AddlByteCount,
|
||||
PrefixVarint32AddlByteCount('\x20'));
|
||||
ASSERT_EQ(kInvalidPrefixVarint32AddlByteCount,
|
||||
PrefixVarint32AddlByteCount('\x40'));
|
||||
ASSERT_EQ(kInvalidPrefixVarint32AddlByteCount,
|
||||
PrefixVarint32AddlByteCount('\x80'));
|
||||
}
|
||||
|
||||
TEST(Coding, PrefixVarint64) {
|
||||
for (const auto& tc : kPrefixVarint64TestCases) {
|
||||
const uint32_t expected_length =
|
||||
tc.value < (uint64_t{1} << 56) ? VarintLength(tc.value) : 9;
|
||||
ASSERT_EQ(tc.length, expected_length);
|
||||
}
|
||||
AssertPrefixVarintRoundTrip(kPrefixVarint64TestCases);
|
||||
}
|
||||
|
||||
TEST(Coding, PrefixVarint64DiskReadApi) {
|
||||
AssertPrefixVarintDiskReadApi(kPrefixVarint64TestCases);
|
||||
ASSERT_EQ(8u, PrefixVarint64AddlByteCount('\0'));
|
||||
ASSERT_EQ(7u, PrefixVarint64AddlByteCount('\x80'));
|
||||
}
|
||||
|
||||
TEST(Coding, PrefixVarint32Overflow) {
|
||||
uint32_t result = 0;
|
||||
const std::string input = EncodeInvalidPrefixVarint32Overflow();
|
||||
ASSERT_FALSE(DecodePrefixVarint32(input[0], input.data() + 1,
|
||||
input.size() - 1, &result));
|
||||
ASSERT_TRUE(GetPrefixVarint32Ptr(input.data(), input.data() + input.size(),
|
||||
&result) == nullptr);
|
||||
}
|
||||
|
||||
TEST(Coding, PrefixVarint32Truncation) {
|
||||
AssertPrefixVarintTruncation(~uint32_t{0});
|
||||
}
|
||||
|
||||
TEST(Coding, PrefixVarint64Truncation) {
|
||||
AssertPrefixVarintTruncation(~uint64_t{0});
|
||||
}
|
||||
|
||||
TEST(Coding, PrefixVarint64ImproperEncoding) {
|
||||
// Values spanning each natural encoding length
|
||||
const std::vector<uint64_t> values = {
|
||||
0,
|
||||
1,
|
||||
63,
|
||||
127, // 1-byte natural
|
||||
128,
|
||||
255,
|
||||
16383, // 2-byte natural
|
||||
16384,
|
||||
(1ull << 21) - 1, // 3-byte natural
|
||||
(1ull << 21),
|
||||
(1ull << 28) - 1, // 4-byte natural
|
||||
(1ull << 28),
|
||||
(1ull << 35) - 1, // 5-byte natural
|
||||
(1ull << 35),
|
||||
(1ull << 42) - 1, // 6-byte natural
|
||||
(1ull << 42),
|
||||
(1ull << 49) - 1, // 7-byte natural
|
||||
(1ull << 49),
|
||||
(1ull << 56) - 1, // 8-byte natural
|
||||
(1ull << 56),
|
||||
~uint64_t{0}, // 9-byte natural
|
||||
};
|
||||
|
||||
TestPrefixVarint64MinimumBytes<2>(values);
|
||||
TestPrefixVarint64MinimumBytes<3>(values);
|
||||
TestPrefixVarint64MinimumBytes<4>(values);
|
||||
TestPrefixVarint64MinimumBytes<5>(values);
|
||||
TestPrefixVarint64MinimumBytes<6>(values);
|
||||
TestPrefixVarint64MinimumBytes<7>(values);
|
||||
TestPrefixVarint64MinimumBytes<8>(values);
|
||||
|
||||
// Verify specific byte patterns for some improper encodings
|
||||
auto check_bytes = [](const char* buf, const char* end,
|
||||
const std::string& expected) {
|
||||
ASSERT_EQ(expected, std::string(buf, static_cast<size_t>(end - buf)));
|
||||
};
|
||||
|
||||
// value=0 with kMinimumBytes=2: (0 << 2) | (1 << 1) = 0x02
|
||||
{
|
||||
char buf[kMaxPrefixVarint64Length];
|
||||
char* end = EncodePrefixVarint64<2>(buf, 0);
|
||||
check_bytes(buf, end, ByteString({0x02, 0x00}));
|
||||
}
|
||||
// value=0 with kMinimumBytes=3: (0 << 3) | (1 << 2) = 0x04
|
||||
{
|
||||
char buf[kMaxPrefixVarint64Length];
|
||||
char* end = EncodePrefixVarint64<3>(buf, 0);
|
||||
check_bytes(buf, end, ByteString({0x04, 0x00, 0x00}));
|
||||
}
|
||||
// value=1 with kMinimumBytes=2: (1 << 2) | (1 << 1) = 0x06
|
||||
{
|
||||
char buf[kMaxPrefixVarint64Length];
|
||||
char* end = EncodePrefixVarint64<2>(buf, 1);
|
||||
check_bytes(buf, end, ByteString({0x06, 0x00}));
|
||||
}
|
||||
// value=127 with kMinimumBytes=2: (127 << 2) | (1 << 1) = 0x01FE
|
||||
{
|
||||
char buf[kMaxPrefixVarint64Length];
|
||||
char* end = EncodePrefixVarint64<2>(buf, 127);
|
||||
check_bytes(buf, end, ByteString({0xFE, 0x01}));
|
||||
}
|
||||
// value=0 with kMinimumBytes=8: (0 << 8) | (1 << 7) = 0x80
|
||||
{
|
||||
char buf[kMaxPrefixVarint64Length];
|
||||
char* end = EncodePrefixVarint64<8>(buf, 0);
|
||||
check_bytes(buf, end,
|
||||
ByteString({0x80, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}));
|
||||
}
|
||||
// value=1 with kMinimumBytes=8: (1 << 8) | (1 << 7) = 0x0180
|
||||
{
|
||||
char buf[kMaxPrefixVarint64Length];
|
||||
char* end = EncodePrefixVarint64<8>(buf, 1);
|
||||
check_bytes(buf, end,
|
||||
ByteString({0x80, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}));
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace ROCKSDB_NAMESPACE
|
||||
|
||||
int main(int argc, char** argv) {
|
||||
|
||||
@@ -0,0 +1,320 @@
|
||||
// Copyright (c) Meta Platforms, Inc. and affiliates.
|
||||
// This source code is licensed under both the GPLv2 (found in the
|
||||
// COPYING file in the root directory) and Apache 2.0 License
|
||||
// (found in the LICENSE.Apache file in the root directory).
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
|
||||
#include "port/likely.h"
|
||||
#include "rocksdb/slice.h"
|
||||
#include "util/cast_util.h"
|
||||
#include "util/coding_lean.h"
|
||||
#include "util/math.h"
|
||||
|
||||
namespace ROCKSDB_NAMESPACE {
|
||||
|
||||
// Prefix varint is an established little-endian, low-bit-prefix alternative to
|
||||
// the continuation-bit varint/LEB128 format used by RocksDB's existing
|
||||
// Varint32/64 helpers. It exists for callers that naturally read one byte
|
||||
// first, such as disk APIs that can issue a second read for the remaining
|
||||
// bytes once the encoded length is known.
|
||||
//
|
||||
// Another motivation is CPU efficiency: for some in-memory workloads this
|
||||
// format may decode more efficiently than continuation-bit varint because the
|
||||
// length is known from the first byte and the payload bits become contiguous
|
||||
// after a shift. That tradeoff still needs evaluation in RocksDB.
|
||||
// TODO: Benchmark PrefixVarint against continuation-bit varint for both
|
||||
// disk-read and in-memory callers before treating it as a CPU optimization.
|
||||
//
|
||||
// This format family is not new. LLVM/MLIR bytecode uses the same
|
||||
// "PrefixVarInt" encoding, including the 9-byte `00000000 <fixed64>` form for
|
||||
// full-width 64-bit values.
|
||||
//
|
||||
// The low-order zero-bit count in the first byte determines the total encoded
|
||||
// length, so the first byte alone tells the caller how many additional bytes
|
||||
// are needed before decoding.
|
||||
//
|
||||
// Prefix-byte layout:
|
||||
// xxxxxxx1 -> 1 byte (7 payload bits)
|
||||
// xxxxxx10 -> 2 bytes (14 payload bits)
|
||||
// xxxxx100 -> 3 bytes (21 payload bits)
|
||||
// xxxx1000 -> 4 bytes (28 payload bits)
|
||||
// xxx10000 -> 5 bytes (35 payload bits)
|
||||
// xx100000 -> 6 bytes (42 payload bits)
|
||||
// x1000000 -> 7 bytes (49 payload bits)
|
||||
// 10000000 -> 8 bytes (56 payload bits)
|
||||
// 00000000 -> 9 bytes (64 payload bits via fixed64 payload bytes)
|
||||
//
|
||||
// PrefixVarint32 uses the first five rows of this table, and therefore uses
|
||||
// exactly the same encoded lengths as LEB128 varint.
|
||||
//
|
||||
// PrefixVarint64 uses the full MLIR PrefixVarInt table above. Relative to
|
||||
// LEB128, the encoded lengths are the same except in the extreme uint64_t case,
|
||||
// where PrefixVarint64 uses 9 bytes and LEB128 would use 10.
|
||||
constexpr uint32_t kMaxPrefixVarint32Length = 5;
|
||||
constexpr uint32_t kInvalidPrefixVarint32AddlByteCount =
|
||||
kMaxPrefixVarint32Length;
|
||||
|
||||
// PrefixVarint64 uses the same low-bit prefix for 1-8 byte encodings.
|
||||
// A zero first byte is reserved for the 9-byte form and is followed by a
|
||||
// little-endian fixed64 payload.
|
||||
constexpr uint32_t kMaxPrefixVarint64Length = 9;
|
||||
|
||||
namespace detail {
|
||||
|
||||
template <typename T>
|
||||
inline uint32_t PrefixVarintLengthImpl(T value) {
|
||||
const uint32_t bits = static_cast<uint32_t>(FloorLog2(value | T{1}) + 1);
|
||||
return (bits + 6) / 7;
|
||||
}
|
||||
|
||||
// Reassembles the already-encoded bytes into a little-endian word so callers
|
||||
// can strip the prefix bits and recover the payload with shifts/masks.
|
||||
inline uint64_t LoadPrefixVarintEncodedWord(char first_byte,
|
||||
const char* addl_bytes,
|
||||
uint32_t addl_byte_count) {
|
||||
uint64_t encoded_word = static_cast<unsigned char>(first_byte);
|
||||
// TODO: Benchmark an unaligned word load plus masking here. MLIR's
|
||||
// parseMultiByteVarInt uses a bulk little-endian load for this path; that
|
||||
// may beat the current byte-at-a-time assembly for hot in-memory callers.
|
||||
for (uint32_t i = 0; i < addl_byte_count; ++i) {
|
||||
encoded_word |=
|
||||
static_cast<uint64_t>(static_cast<unsigned char>(addl_bytes[i]))
|
||||
<< ((i + 1) * 8);
|
||||
}
|
||||
return encoded_word;
|
||||
}
|
||||
|
||||
// Stores the low `byte_count` bytes of an already-formed encoded word.
|
||||
inline char* StorePrefixVarintEncodedWord(char* dst, uint64_t encoded_word,
|
||||
uint32_t byte_count) {
|
||||
unsigned char* ptr = lossless_cast<unsigned char*>(dst);
|
||||
// TODO: Benchmark an unaligned word store here. MLIR's emitMultiByteVarInt
|
||||
// writes from an already-formed little-endian word, which may beat the
|
||||
// current byte-at-a-time write for hot in-memory callers.
|
||||
for (uint32_t i = 0; i < byte_count; ++i) {
|
||||
ptr[i] = static_cast<unsigned char>(encoded_word >> (i * 8));
|
||||
}
|
||||
return lossless_cast<char*>(ptr + byte_count);
|
||||
}
|
||||
|
||||
} // namespace detail
|
||||
|
||||
inline uint32_t PrefixVarint32Length(uint32_t value) {
|
||||
return detail::PrefixVarintLengthImpl(value);
|
||||
}
|
||||
|
||||
inline uint32_t PrefixVarint64Length(uint64_t value) {
|
||||
const uint32_t len = detail::PrefixVarintLengthImpl(value);
|
||||
return std::min(len, kMaxPrefixVarint64Length);
|
||||
}
|
||||
|
||||
inline char* EncodePrefixVarint32(char* dst, uint32_t value) {
|
||||
const uint32_t num_bytes = PrefixVarint32Length(value);
|
||||
const uint64_t encoded_word = (static_cast<uint64_t>(value) << num_bytes) |
|
||||
(uint64_t{1} << (num_bytes - 1));
|
||||
return detail::StorePrefixVarintEncodedWord(dst, encoded_word, num_bytes);
|
||||
}
|
||||
|
||||
// Encode `value` as a PrefixVarint64 to the buffer at `dst`, which should in
|
||||
// general be at least kMaxPrefixVarint64Length bytes. Returns a pointer to the
|
||||
// byte after the last encoded byte. kMinimumBytes is rarely used, to support
|
||||
// "improper" (non-minimal) encoding.
|
||||
template <uint32_t kMinimumBytes = 0>
|
||||
inline char* EncodePrefixVarint64(char* dst, uint64_t value) {
|
||||
uint32_t num_bytes = PrefixVarint64Length(value);
|
||||
if constexpr (kMinimumBytes > 1) {
|
||||
static_assert(kMinimumBytes < kMaxPrefixVarint64Length);
|
||||
num_bytes = std::max(num_bytes, kMinimumBytes);
|
||||
}
|
||||
if (UNLIKELY(num_bytes == kMaxPrefixVarint64Length)) {
|
||||
dst[0] = 0;
|
||||
EncodeFixed64(dst + 1, value);
|
||||
return dst + kMaxPrefixVarint64Length;
|
||||
}
|
||||
|
||||
const uint64_t encoded_word =
|
||||
(value << num_bytes) | (uint64_t{1} << (num_bytes - 1));
|
||||
return detail::StorePrefixVarintEncodedWord(dst, encoded_word, num_bytes);
|
||||
}
|
||||
|
||||
inline void PutPrefixVarint32(std::string* dst, uint32_t value) {
|
||||
char buf[kMaxPrefixVarint32Length];
|
||||
char* ptr = EncodePrefixVarint32(buf, value);
|
||||
dst->append(buf, static_cast<size_t>(ptr - buf));
|
||||
}
|
||||
|
||||
inline void PutPrefixVarint64(std::string* dst, uint64_t value) {
|
||||
char buf[kMaxPrefixVarint64Length];
|
||||
char* ptr = EncodePrefixVarint64(buf, value);
|
||||
dst->append(buf, static_cast<size_t>(ptr - buf));
|
||||
}
|
||||
|
||||
// Split-decode helper for callers that already have the first byte and need to
|
||||
// know how many additional bytes to fetch before calling
|
||||
// DecodePrefixVarint32(). Returns kInvalidPrefixVarint32AddlByteCount if
|
||||
// `first_byte` cannot start a valid PrefixVarint32.
|
||||
inline uint32_t PrefixVarint32AddlByteCount(char first_byte) {
|
||||
const uint32_t bits = static_cast<unsigned char>(first_byte);
|
||||
if (UNLIKELY(bits == 0)) {
|
||||
return kInvalidPrefixVarint32AddlByteCount;
|
||||
}
|
||||
const uint32_t addl_byte_count =
|
||||
static_cast<uint32_t>(CountTrailingZeroBits(bits));
|
||||
return addl_byte_count < kMaxPrefixVarint32Length
|
||||
? addl_byte_count
|
||||
: kInvalidPrefixVarint32AddlByteCount;
|
||||
}
|
||||
|
||||
// Split-decode helper for callers that already have the first byte and need to
|
||||
// know how many additional bytes to fetch before calling
|
||||
// DecodePrefixVarint64(). For the extreme 9-byte form, this returns 8.
|
||||
inline uint32_t PrefixVarint64AddlByteCount(char first_byte) {
|
||||
const uint32_t bits = static_cast<unsigned char>(first_byte);
|
||||
return bits == 0 ? kMaxPrefixVarint64Length - 1
|
||||
: static_cast<uint32_t>(CountTrailingZeroBits(bits));
|
||||
}
|
||||
|
||||
// Decodes a PrefixVarint32 when the caller already has the first byte and the
|
||||
// additional bytes requested by PrefixVarint32AddlByteCount(first_byte), which
|
||||
// is useful for disk-read APIs. Returns false on invalid first byte, too few
|
||||
// additional bytes, or overflow beyond uint32_t.
|
||||
inline bool DecodePrefixVarint32(char first_byte, const char* addl_bytes,
|
||||
size_t addl_byte_count, uint32_t* value) {
|
||||
const uint32_t first = static_cast<unsigned char>(first_byte);
|
||||
// Optimize the overwhelmingly common single-byte case.
|
||||
if (LIKELY((first & 0x01u) != 0)) {
|
||||
*value = first >> 1;
|
||||
return true;
|
||||
}
|
||||
|
||||
const uint32_t required_addl_byte_count =
|
||||
PrefixVarint32AddlByteCount(first_byte);
|
||||
if (required_addl_byte_count == kInvalidPrefixVarint32AddlByteCount ||
|
||||
addl_byte_count < required_addl_byte_count) {
|
||||
return false;
|
||||
}
|
||||
|
||||
const uint64_t decoded =
|
||||
detail::LoadPrefixVarintEncodedWord(first_byte, addl_bytes,
|
||||
required_addl_byte_count) >>
|
||||
(required_addl_byte_count + 1);
|
||||
if ((decoded >> 32) != 0) {
|
||||
return false;
|
||||
}
|
||||
*value = static_cast<uint32_t>(decoded);
|
||||
return true;
|
||||
}
|
||||
|
||||
// Decodes a PrefixVarint64 when the caller already has the first byte and the
|
||||
// additional bytes requested by PrefixVarint64AddlByteCount(first_byte).
|
||||
// Returns false only when too few additional bytes are provided.
|
||||
inline bool DecodePrefixVarint64(char first_byte, const char* addl_bytes,
|
||||
size_t addl_byte_count, uint64_t* value) {
|
||||
const uint32_t first = static_cast<unsigned char>(first_byte);
|
||||
// Optimize the overwhelmingly common single-byte case.
|
||||
if (LIKELY((first & 0x01u) != 0)) {
|
||||
*value = first >> 1;
|
||||
return true;
|
||||
}
|
||||
|
||||
const uint32_t required_addl_byte_count =
|
||||
PrefixVarint64AddlByteCount(first_byte);
|
||||
if (addl_byte_count < required_addl_byte_count) {
|
||||
return false;
|
||||
}
|
||||
if (UNLIKELY(required_addl_byte_count == kMaxPrefixVarint64Length - 1)) {
|
||||
*value = DecodeFixed64(addl_bytes);
|
||||
return true;
|
||||
}
|
||||
|
||||
*value = detail::LoadPrefixVarintEncodedWord(first_byte, addl_bytes,
|
||||
required_addl_byte_count) >>
|
||||
(required_addl_byte_count + 1);
|
||||
return true;
|
||||
}
|
||||
|
||||
// Parses one PrefixVarint32 from [p, limit). On success stores the decoded
|
||||
// value in *value and returns the first byte after the encoding. On failure
|
||||
// returns nullptr. It never reads past `limit`.
|
||||
inline const char* GetPrefixVarint32Ptr(const char* p, const char* limit,
|
||||
uint32_t* value) {
|
||||
if (p >= limit) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
const uint32_t first = static_cast<unsigned char>(*p);
|
||||
// Optimize the overwhelmingly common single-byte case.
|
||||
if (LIKELY((first & 0x01u) != 0)) {
|
||||
*value = first >> 1;
|
||||
return p + 1;
|
||||
}
|
||||
|
||||
const uint32_t addl_byte_count = PrefixVarint32AddlByteCount(*p);
|
||||
if (addl_byte_count == kInvalidPrefixVarint32AddlByteCount ||
|
||||
static_cast<size_t>(limit - p) < addl_byte_count + 1) {
|
||||
return nullptr;
|
||||
}
|
||||
if (!DecodePrefixVarint32(*p, p + 1, addl_byte_count, value)) {
|
||||
return nullptr;
|
||||
}
|
||||
return p + addl_byte_count + 1;
|
||||
}
|
||||
|
||||
// Parses one PrefixVarint64 from [p, limit). On success stores the decoded
|
||||
// value in *value and returns the first byte after the encoding. On failure
|
||||
// returns nullptr. It never reads past `limit`.
|
||||
inline const char* GetPrefixVarint64Ptr(const char* p, const char* limit,
|
||||
uint64_t* value) {
|
||||
if (p >= limit) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
const uint32_t first = static_cast<unsigned char>(*p);
|
||||
// Optimize the overwhelmingly common single-byte case.
|
||||
if (LIKELY((first & 0x01u) != 0)) {
|
||||
*value = first >> 1;
|
||||
return p + 1;
|
||||
}
|
||||
|
||||
const uint32_t addl_byte_count = PrefixVarint64AddlByteCount(*p);
|
||||
if (static_cast<size_t>(limit - p) < addl_byte_count + 1) {
|
||||
return nullptr;
|
||||
}
|
||||
if (!DecodePrefixVarint64(*p, p + 1, addl_byte_count, value)) {
|
||||
return nullptr;
|
||||
}
|
||||
return p + addl_byte_count + 1;
|
||||
}
|
||||
|
||||
// Slice adapter around GetPrefixVarint32Ptr(). Advances `input` only on
|
||||
// success.
|
||||
inline bool GetPrefixVarint32(Slice* input, uint32_t* value) {
|
||||
const char* p = input->data();
|
||||
const char* limit = p + input->size();
|
||||
const char* q = GetPrefixVarint32Ptr(p, limit, value);
|
||||
if (q == nullptr) {
|
||||
return false;
|
||||
}
|
||||
*input = Slice(q, static_cast<size_t>(limit - q));
|
||||
return true;
|
||||
}
|
||||
|
||||
// Slice adapter around GetPrefixVarint64Ptr(). Advances `input` only on
|
||||
// success.
|
||||
inline bool GetPrefixVarint64(Slice* input, uint64_t* value) {
|
||||
const char* p = input->data();
|
||||
const char* limit = p + input->size();
|
||||
const char* q = GetPrefixVarint64Ptr(p, limit, value);
|
||||
if (q == nullptr) {
|
||||
return false;
|
||||
}
|
||||
*input = Slice(q, static_cast<size_t>(limit - q));
|
||||
return true;
|
||||
}
|
||||
|
||||
} // namespace ROCKSDB_NAMESPACE
|
||||
Reference in New Issue
Block a user