diff --git a/CMakeLists.txt b/CMakeLists.txt index 7ed3dc7..4dc9b59 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -133,6 +133,16 @@ target_include_directories(dbps_common_lib PUBLIC ) target_link_libraries(dbps_common_lib PUBLIC tcb_span) +# Byte buffer processing library +add_library(dbps_byte_buffer_lib STATIC + src/processing/byte_buffer.cpp +) +target_include_directories(dbps_byte_buffer_lib PUBLIC + src/processing + src/common +) +target_link_libraries(dbps_byte_buffer_lib PUBLIC tcb_span) + # Server components library add_library(dbps_server_lib STATIC src/processing/encryption_sequencer.cpp @@ -205,7 +215,7 @@ target_include_directories(dbps_local_lib PUBLIC ) # Ensure PIC on static libs that are linked into shared libraries -set_target_properties(dbps_common_lib dbps_server_lib dbps_local_lib +set_target_properties(dbps_common_lib dbps_byte_buffer_lib dbps_server_lib dbps_local_lib PROPERTIES POSITION_INDEPENDENT_CODE ON) # ============================================================================= @@ -325,6 +335,14 @@ if(BUILD_TESTS) ) target_include_directories(compression_utils_test PRIVATE src/processing) + # Byte buffer tests + add_executable(byte_buffer_test src/processing/byte_buffer_test.cpp) + target_link_libraries(byte_buffer_test + dbps_byte_buffer_lib + gtest_main + ) + target_include_directories(byte_buffer_test PRIVATE src/processing src/common) + # Basic encryptor tests add_executable(basic_encryptor_test src/processing/encryptors/basic_encryptor_test.cpp) target_link_libraries(basic_encryptor_test @@ -482,6 +500,7 @@ endif() add_custom_target(libraries DEPENDS dbps_common_lib + dbps_byte_buffer_lib dbps_server_lib dbps_client_lib dbps_remote_lib @@ -517,6 +536,7 @@ if(BUILD_TESTS) parquet_utils_test bytes_utils_test compression_utils_test + byte_buffer_test basic_encryptor_test auth_utils_test dbpa_interface_test @@ -539,6 +559,7 @@ if(BUILD_TESTS) gtest_discover_tests(parquet_utils_test) gtest_discover_tests(bytes_utils_test) gtest_discover_tests(compression_utils_test) + gtest_discover_tests(byte_buffer_test) gtest_discover_tests(basic_encryptor_test) gtest_discover_tests(auth_utils_test) gtest_discover_tests(dbpa_interface_test) diff --git a/src/common/bytes_utils.h b/src/common/bytes_utils.h index 9847f2d..93d9047 100644 --- a/src/common/bytes_utils.h +++ b/src/common/bytes_utils.h @@ -25,15 +25,18 @@ #include #include #include +#include #include "exceptions.h" // Utility functions for little-endian number reading and writing. inline void append_u32_le(std::vector& out, uint32_t v) { - out.push_back(static_cast(v & 0xFF)); - out.push_back(static_cast((v >> 8) & 0xFF)); - out.push_back(static_cast((v >> 16) & 0xFF)); - out.push_back(static_cast((v >> 24) & 0xFF)); + const size_t offset = out.size(); + out.resize(offset + 4); + out[offset + 0] = static_cast(v & 0xFF); + out[offset + 1] = static_cast((v >> 8) & 0xFF); + out[offset + 2] = static_cast((v >> 16) & 0xFF); + out[offset + 3] = static_cast((v >> 24) & 0xFF); } inline void append_i32_le(std::vector& out, int32_t v) { @@ -41,14 +44,16 @@ inline void append_i32_le(std::vector& out, int32_t v) { } inline void append_u64_le(std::vector& out, uint64_t v) { - out.push_back(static_cast(v & 0xFF)); - out.push_back(static_cast((v >> 8) & 0xFF)); - out.push_back(static_cast((v >> 16) & 0xFF)); - out.push_back(static_cast((v >> 24) & 0xFF)); - out.push_back(static_cast((v >> 32) & 0xFF)); - out.push_back(static_cast((v >> 40) & 0xFF)); - out.push_back(static_cast((v >> 48) & 0xFF)); - out.push_back(static_cast((v >> 56) & 0xFF)); + const size_t offset = out.size(); + out.resize(offset + 8); + out[offset + 0] = static_cast(v & 0xFF); + out[offset + 1] = static_cast((v >> 8) & 0xFF); + out[offset + 2] = static_cast((v >> 16) & 0xFF); + out[offset + 3] = static_cast((v >> 24) & 0xFF); + out[offset + 4] = static_cast((v >> 32) & 0xFF); + out[offset + 5] = static_cast((v >> 40) & 0xFF); + out[offset + 6] = static_cast((v >> 48) & 0xFF); + out[offset + 7] = static_cast((v >> 56) & 0xFF); } inline void append_i64_le(std::vector& out, int64_t v) { @@ -74,6 +79,13 @@ inline uint32_t read_u32_le(const std::vector& in, size_t offset) { (static_cast(in[offset + 3]) << 24); } +inline uint32_t read_u32_le(tcb::span in, size_t offset) { + return static_cast(in[offset]) | + (static_cast(in[offset + 1]) << 8) | + (static_cast(in[offset + 2]) << 16) | + (static_cast(in[offset + 3]) << 24); +} + // Utility functions for splitting and joining byte vectors. struct SplitBytesPair { diff --git a/src/processing/byte_buffer.cpp b/src/processing/byte_buffer.cpp new file mode 100644 index 0000000..75eda74 --- /dev/null +++ b/src/processing/byte_buffer.cpp @@ -0,0 +1,506 @@ +// 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 "byte_buffer.h" + +#include "bytes_utils.h" +#include "exceptions.h" +#include +#include +#include +#include + +namespace dbps::processing { + +namespace { +// Reads the element size at the given offset. +inline size_t ReadSizeAt(tcb::span bytes, size_t offset) { + return static_cast(read_u32_le(bytes, offset)); +} +} + +// ----------------------------------------------------------------------------- +// Constructors and initializers for read-only buffer +// ----------------------------------------------------------------------------- + +// Constructor for read-only buffer with fixed-size elements. +ByteBuffer::ByteBuffer( + tcb::span elements_span, + size_t element_size, + size_t prefix_size) + : elements_span_(elements_span), + num_elements_(0), + has_fixed_sized_elements_(true), + element_size_(element_size), + prefix_size_(prefix_size), + is_initialized_from_span_(false) {} + +// Constructor for read-only buffer with variable-size elements. +ByteBuffer::ByteBuffer( + tcb::span elements_span, + size_t prefix_size) + : elements_span_(elements_span), + num_elements_(0), + has_fixed_sized_elements_(false), + element_size_(0), + prefix_size_(prefix_size), + is_initialized_from_span_(false) {} + +// Initializes `num_elements_` and `offsets_` from the span. +// Called in a lazy manner when the buffer is accessed with GetElement, avoiding unnecessary initialization. +void ByteBuffer::InitializeFromSpan() const { + if (elements_span_.size() < prefix_size_) { + throw InvalidInputException("Malformed buffer: prefix_size exceeds span size"); + } + + const size_t readable_size = elements_span_.size() - prefix_size_; + + // No elements to index. Initialize with empty values. + if (readable_size == 0) { + num_elements_ = 0; + offsets_.clear(); + is_initialized_from_span_ = true; + return; + } + + // Fixed-size layout has implicit offsets from index * element_size. + // We validate shape and derive element count. No need to store offsets. + if (has_fixed_sized_elements_) { + if (element_size_ <= 0) { + throw InvalidInputException("Invalid fixed-size buffer: element_size must be greater than zero"); + } + if ((readable_size % element_size_) != 0) { + throw InvalidInputException("Malformed fixed-size buffer: buffer does not align with element_size"); + } + num_elements_ = readable_size / element_size_; + offsets_.clear(); + is_initialized_from_span_ = true; + return; + } + + // Variable-size layout stores [u32 size][element value] back-to-back. + // Single pass validates shape and captures per-element prefix offsets. + offsets_.clear(); + offsets_.reserve(EstimateOffsetsReserveCountFromSample(elements_span_.subspan(prefix_size_))); + size_t cursor = prefix_size_; + while (cursor < elements_span_.size()) { + if (elements_span_.size() - cursor < kSizePrefixBytes) { + throw InvalidInputException("Malformed variable-size buffer: truncated length prefix"); + } + offsets_.push_back(cursor); + const size_t current_element_size = ReadSizeAt(elements_span_, cursor); + cursor += kSizePrefixBytes; + if (elements_span_.size() - cursor < current_element_size) { + throw InvalidInputException("Malformed variable-size buffer: truncated element payload"); + } + cursor += current_element_size; + } + num_elements_ = offsets_.size(); + is_initialized_from_span_ = true; +} + +void ByteBuffer::EnsureInitializedFromSpan() const { + // If the span is already initialized, skip it. + if (is_initialized_from_span_) { + return; + } + // If the write buffer is initialized, we don't need to initialize from the span. + if (is_write_buffer_initialized_) { + return; + } + InitializeFromSpan(); +} + +size_t ByteBuffer::EstimateOffsetsReserveCountFromSample(tcb::span bytes) { + if (bytes.empty()) + return 0; + + // Sample the first 10 elements to estimate the total element count. + size_t cursor = 0; + size_t sampled_elements = 0; + while (cursor < bytes.size() && sampled_elements < 10) { + if (bytes.size() - cursor < kSizePrefixBytes) { + throw InvalidInputException("Malformed variable-size buffer: truncated length prefix"); + } + const size_t current_element_size = ReadSizeAt(bytes, cursor); + cursor += kSizePrefixBytes; + if (bytes.size() - cursor < current_element_size) { + throw InvalidInputException("Malformed variable-size buffer: truncated element payload"); + } + cursor += current_element_size; + ++sampled_elements; + } + + // If no elements were sampled or the cursor didn't move, it means the buffer is empty. So return 0. + if (sampled_elements == 0 || cursor == 0) + return 0; + + // If sampling consumed the full buffer (<= sample window), we already know the exact count. + if (cursor == bytes.size()) + return sampled_elements; + + // Estimate total element count from sample density: (sampled_elements / sampled_bytes) * total_bytes. + // - sampled_elements / sampled_bytes gives "elements per byte" in the sampled prefix, + // - then multiplying by total_bytes extrapolates a full-buffer estimate. + const long double estimated = + (static_cast(bytes.size()) * static_cast(sampled_elements)) / + static_cast(cursor); + const size_t estimated_count = static_cast(std::ceil(estimated)); + const size_t estimated_with_headroom = + static_cast(std::ceil(static_cast(estimated_count) * 1.1L)); + + // If the count of sampled elements is more than the estimated, conservatively return actual count. + return std::max(estimated_with_headroom, sampled_elements); +} + +// ----------------------------------------------------------------------------- +// Element span reader methods +// ----------------------------------------------------------------------------- + +size_t ByteBuffer::CalculateOffsetOfElement(size_t position) const { + EnsureInitializedFromSpan(); + if (position >= num_elements_) { + throw InvalidInputException("Element position out of range during CalculateOffsetOfElement"); + } + if (has_fixed_sized_elements_) { + return prefix_size_ + (position * element_size_); + } + return offsets_[position]; +} + +tcb::span ByteBuffer::GetElement(size_t position) const { + EnsureInitializedFromSpan(); + if (position >= num_elements_) { + throw InvalidInputException("Element position out of range during GetElement"); + } + const size_t offset = CalculateOffsetOfElement(position); + + // For fixed-size elemments are stored contiguously. + if (has_fixed_sized_elements_) { + return elements_span_.subspan(offset, element_size_); + } + + // For variable-size elements, we need to read the size first [u32 size][element]. + if (offset == kUnsetVariableElementOffset) { + throw InvalidInputException("Element position has not been written yet"); + } + const size_t element_size = ReadSizeAt(elements_span_, offset); + return elements_span_.subspan(offset + kSizePrefixBytes, element_size); +} + +// ----------------------------------------------------------------------------- +// Elemment span iterator +// +// Allows an alternative read of elements_span_ without need for lazy initialization of offsets_, +// so saving execution time when the traversal of the buffer is strictly sequential. +// This is the most common behavior when reading elements in single threaded mode. +// ----------------------------------------------------------------------------- + +ByteBuffer::ConstIterator::ConstIterator(const ByteBuffer* buffer, size_t cursor_offset) + : buffer_(buffer), cursor_offset_(cursor_offset) {} + +void ByteBuffer::ConstIterator::ValidateFixedSizeElementAtCursor() const { + if (buffer_->element_size_ <= 0) { + throw InvalidInputException("Invalid fixed-size buffer: element_size must be greater than zero"); + } + if ((buffer_->elements_span_.size() - cursor_offset_) < buffer_->element_size_) { + throw InvalidInputException("Malformed fixed-size buffer: truncated element payload"); + } +} + +size_t ByteBuffer::ConstIterator::ReadAndValidateVariableElementSizeAtCursor() const { + if ((buffer_->elements_span_.size() - cursor_offset_) < kSizePrefixBytes) { + throw InvalidInputException("Malformed variable-size buffer: truncated length prefix"); + } + const size_t current_element_size = ReadSizeAt(buffer_->elements_span_, cursor_offset_); + const size_t payload_offset = cursor_offset_ + kSizePrefixBytes; + if ((buffer_->elements_span_.size() - payload_offset) < current_element_size) { + throw InvalidInputException("Malformed variable-size buffer: truncated element payload"); + } + return current_element_size; +} + +ByteBuffer::ConstIterator::value_type ByteBuffer::ConstIterator::operator*() const { + if (buffer_ == nullptr || cursor_offset_ >= buffer_->elements_span_.size()) { + throw InvalidInputException("Cannot dereference ByteBuffer iterator at end position"); + } + if (buffer_->has_fixed_sized_elements_) { + ValidateFixedSizeElementAtCursor(); + return buffer_->elements_span_.subspan(cursor_offset_, buffer_->element_size_); + } + + const size_t current_element_size = ReadAndValidateVariableElementSizeAtCursor(); + const size_t payload_offset = cursor_offset_ + kSizePrefixBytes; + return buffer_->elements_span_.subspan(payload_offset, current_element_size); +} + +ByteBuffer::ConstIterator& ByteBuffer::ConstIterator::operator++() { + if (buffer_ == nullptr || cursor_offset_ >= buffer_->elements_span_.size()) { + return *this; + } + if (buffer_->has_fixed_sized_elements_) { + ValidateFixedSizeElementAtCursor(); + cursor_offset_ += buffer_->element_size_; + return *this; + } + + const size_t current_element_size = ReadAndValidateVariableElementSizeAtCursor(); + cursor_offset_ += (kSizePrefixBytes + current_element_size); + return *this; +} + +bool ByteBuffer::ConstIterator::operator==(const ConstIterator& other) const { + return buffer_ == other.buffer_ && cursor_offset_ == other.cursor_offset_; +} + +bool ByteBuffer::ConstIterator::operator!=(const ConstIterator& other) const { + return !(*this == other); +} + +void ByteBuffer::ValidateIteratorReadPreconditions() const { + if (is_write_buffer_initialized_) { + throw InvalidInputException("Iterator is only available for read buffers"); + } + if (elements_span_.size() < prefix_size_) { + throw InvalidInputException("Malformed buffer: prefix_size exceeds span size"); + } + if (has_fixed_sized_elements_) { + if (element_size_ <= 0) { + throw InvalidInputException("Invalid fixed-size buffer: element_size must be greater than zero"); + } + const size_t readable_size = elements_span_.size() - prefix_size_; + if ((readable_size % element_size_) != 0) { + throw InvalidInputException("Malformed fixed-size buffer: buffer does not align with element_size"); + } + } +} + +ByteBuffer::ConstIterator ByteBuffer::begin() const { + ValidateIteratorReadPreconditions(); + return ConstIterator(this, prefix_size_); +} + +ByteBuffer::ConstIterator ByteBuffer::end() const { + ValidateIteratorReadPreconditions(); + return ConstIterator(this, elements_span_.size()); +} + +// ----------------------------------------------------------------------------- +// Constructors and initializers for write buffer +// ----------------------------------------------------------------------------- + +// Constructor for a new write buffer with fixed-size elements. +ByteBuffer::ByteBuffer( + size_t num_elements, + size_t element_size, + size_t prefix_size) + : num_elements_(num_elements), + has_fixed_sized_elements_(true), + element_size_(element_size), + prefix_size_(prefix_size) { + InitializeForWriteBuffer(0); +} + +// Constructor for a new write buffer with variable-size elements. +ByteBuffer::ByteBuffer( + size_t num_elements, + size_t reserved_bytes_hint, + bool use_reserve_hint, + size_t prefix_size) + : num_elements_(num_elements), + has_fixed_sized_elements_(false), + element_size_(0), + prefix_size_(prefix_size) { + InitializeForWriteBuffer(use_reserve_hint ? reserved_bytes_hint : 0); +} + +// Initializes `write_buffer_`, `offsets_` and `elements_span_` +void ByteBuffer::InitializeForWriteBuffer(size_t variable_size_reserved_bytes_hint) { + // Fixed-size elements + if (has_fixed_sized_elements_) { + if (element_size_ <= 0) { + throw InvalidInputException("Invalid fixed-size buffer: element_size must be greater than zero"); + } + + // write_buffer can be allocated to precise size since the element size and number of elements are known. + // We initialize it to 0s to have random-ish access during writes. + const size_t fixed_size_total_bytes = prefix_size_ + (num_elements_ * element_size_); + write_buffer_.clear(); + write_buffer_.resize(fixed_size_total_bytes, static_cast(0)); + is_write_buffer_initialized_ = true; + + // offsets_ are not used for fixed-size elements. + offsets_.clear(); + + // elements_span_ is re-bound to the write buffer. + RebindSpanToWriteBuffer(); + return; + } + + // Variable-size elements + + // Reserve write_buffer to at least the size of the prefix [u32 size] bytes for all elements, + // and use a larger reserved-bytes hint if given as a best guess to reduce reallocations. + // write_buffer is not initialized to anything since we will be appending to it during SetElement, just reserving capacity. + const size_t min_required_record_bytes = num_elements_ * kSizePrefixBytes; + const size_t variable_size_reserved_bytes = + prefix_size_ + std::max(variable_size_reserved_bytes_hint, min_required_record_bytes); + write_buffer_.clear(); + write_buffer_.resize(prefix_size_, static_cast(0)); + write_buffer_.reserve(variable_size_reserved_bytes); + is_write_buffer_initialized_ = true; + + // offsets_ is initialized so the vector is fully allocated and have random-ish access during writes. + offsets_.clear(); + offsets_.resize(num_elements_, kUnsetVariableElementOffset); + + // next_expected_sequential_position_ is initialized to 0 for sequential write checking. + next_expected_write_position_ = 0; + + // elements_span_ is re-bound to the write buffer. + RebindSpanToWriteBuffer(); +} + +// ----------------------------------------------------------------------------- +// Buffer writer methods +// ----------------------------------------------------------------------------- + +void ByteBuffer::SetElement(size_t position, tcb::span element) { + if (!is_write_buffer_initialized_) { + throw InvalidInputException("Cannot SetElement: write buffer is not initialized."); + } + + if (is_write_buffer_finalized_) { + throw InvalidInputException("Cannot SetElement: write buffer has been finalized"); + } + + if (position >= num_elements_) { + throw InvalidInputException("Element position out of range during SetElement"); + } + + // For fixed-size elements, we write the element to buffer at the offset. No need to re-bind the span. + if (has_fixed_sized_elements_) { + if (element.size() != element_size_) { + throw InvalidInputException("Fixed-size element payload size mismatch"); + } + const size_t offset = CalculateOffsetOfElement(position); + std::memcpy(write_buffer_.data() + offset, element.data(), element_size_); + return; + } + + // Defensive check for unlikely extremely large element size that exceeds uint32. + if (element.size() > static_cast(std::numeric_limits::max())) { + throw InvalidInputException("Variable-size element payload exceeds uint32 capacity.. Woohhh!!"); + } + + // For variable-size elements, we append the element to the write buffer and update offsets_. + // + // This can result on orphaned bytes if a position is set multiple times or positions written out of order. + // This is intentional to allow random writes of elements while the buffer is built. + // During FinalizeAndTakeBuffer, the buffer is rebuilt to be sequential and orphaned bytes are removed. + const size_t offset = write_buffer_.size(); + offsets_[position] = offset; + append_u32_le(write_buffer_, static_cast(element.size())); + write_buffer_.insert(write_buffer_.end(), element.begin(), element.end()); // Appends at the end of the buffer. + + // Update next_expected_write_position_ for sequential write checking. + if (next_expected_write_position_ != kUnsetVariableElementOffset) { + if (position == next_expected_write_position_) { + next_expected_write_position_ += 1; + } else { + next_expected_write_position_ = kUnsetVariableElementOffset; + } + } + + RebindSpanToWriteBuffer(); +} + +std::vector ByteBuffer::FinalizeAndTakeBuffer() { + if (is_write_buffer_finalized_) { + throw InvalidInputException("FinalizeAndTakeBuffer: write buffer has already been finalized"); + } + + if (!is_write_buffer_initialized_) { + throw InvalidInputException("FinalizeAndTakeBuffer: write buffer is not initialized"); + } + + // Fixed-size: write_buffer_ is always in element order, transfer ownership directly. + if (has_fixed_sized_elements_) { + is_write_buffer_finalized_ = true; + return std::move(write_buffer_); + } + + // For variable-size when all elements were written exactly once and in sequential order, + // we can skip out-of-order or fragmentation checks. This is the fast path. + // This is the most common behavior when writing elements in single threaded mode. + if (next_expected_write_position_ == num_elements_) { + if (num_elements_ > 0) { + const size_t last_element_offset = offsets_[num_elements_ - 1]; + const size_t last_element_size = ReadSizeAt(elements_span_, last_element_offset); + const size_t logical_size = last_element_offset + kSizePrefixBytes + last_element_size; + if (logical_size != write_buffer_.size()) { + throw InvalidInputException("FinalizeAndTakeBuffer: trailing bytes detected beyond last element"); + } + } + is_write_buffer_finalized_ = true; + return std::move(write_buffer_); + } + + // For variable-size, when elements are written out of order, assume the buffer is fragmented and potentially with orphaned bytes + // The buffer is validated and rebuilt into an ordered compact buffer in one pass. + std::vector result; + result.reserve(write_buffer_.size()); + // Copy the prefix bytes at the beginning of the result. + result.insert( + result.end(), + write_buffer_.begin(), + write_buffer_.begin() + static_cast(prefix_size_)); + for (size_t i = 0; i < num_elements_; ++i) { + const size_t element_offset = offsets_[i]; + if (element_offset == kUnsetVariableElementOffset) { + throw InvalidInputException("Cannot finalize variable-size buffer: not all elements were written"); + } + if (element_offset > write_buffer_.size() || (write_buffer_.size() - element_offset) < kSizePrefixBytes) { + throw InvalidInputException("Cannot finalize variable-size buffer: invalid element offset"); + } + + const size_t element_size = ReadSizeAt(elements_span_, element_offset); + if (element_size > (write_buffer_.size() - element_offset - kSizePrefixBytes)) { + throw InvalidInputException("Cannot finalize variable-size buffer: malformed element payload"); + } + + const size_t record_size = kSizePrefixBytes + element_size; + result.insert( + result.end(), + write_buffer_.data() + element_offset, + write_buffer_.data() + element_offset + record_size); + } + + // Defrag path returns a new buffer; release the original fragmented write buffer. + write_buffer_.clear(); + write_buffer_.shrink_to_fit(); + is_write_buffer_initialized_ = false; + is_write_buffer_finalized_ = true; + + return result; +} + +void ByteBuffer::RebindSpanToWriteBuffer() { + elements_span_ = tcb::span(write_buffer_.data(), write_buffer_.size()); +} + +} // namespace dbps::processing diff --git a/src/processing/byte_buffer.h b/src/processing/byte_buffer.h new file mode 100644 index 0000000..96d8fd1 --- /dev/null +++ b/src/processing/byte_buffer.h @@ -0,0 +1,141 @@ +// 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 +#include +#include +#include +#include + +#include + +namespace dbps::processing { + +class ByteBuffer { +public: + // Constructor for read-only buffer with fixed-size elements. + // Elements are stored contiguously in the span. + ByteBuffer( + tcb::span elements_span, + size_t element_size, + size_t prefix_size); + + // Constructor for read-only buffer with variable-size elements. + // Elements are encoded as [u32 size][element value] contiguously in the span. + ByteBuffer( + tcb::span elements_span, + size_t prefix_size = 0); + + // Constructor for a new write buffer with fixed-size elements. + ByteBuffer( + size_t num_elements, + size_t element_size, + size_t prefix_size = 0); + + // Constructor for a new write buffer with variable-size elements. + ByteBuffer( + size_t num_elements, + size_t reserved_bytes_hint, + bool use_reserve_hint, + size_t prefix_size = 0); + + // Get and set elements by position + tcb::span GetElement(size_t position) const; + void SetElement(size_t position, tcb::span element); + + // Finalizes the write path and transfers the resulting buffer ownership. + std::vector FinalizeAndTakeBuffer(); + + // Iterator for read-only elements. + class ConstIterator { + public: + // Iterator traits consumed indirectly by STL iterator machinery. + using iterator_category = std::forward_iterator_tag; + using value_type = tcb::span; + using difference_type = std::ptrdiff_t; + using pointer = void; + using reference = value_type; + + // Basic forward-iterator operations over encoded elements in elements_span_. + ConstIterator(const ByteBuffer* buffer, size_t cursor_offset); + value_type operator*() const; + ConstIterator& operator++(); + bool operator==(const ConstIterator& other) const; + bool operator!=(const ConstIterator& other) const; + + private: + void ValidateFixedSizeElementAtCursor() const; + size_t ReadAndValidateVariableElementSizeAtCursor() const; + + const ByteBuffer* buffer_ = nullptr; + size_t cursor_offset_ = 0; + }; + + // Methods used by the STL iterator machinery to iterate over the buffer. + ConstIterator begin() const; + ConstIterator end() const; + +protected: + // Helper for reserve heuristics in variable-size parsing. + static size_t EstimateOffsetsReserveCountFromSample(tcb::span bytes); + + // Helper for calculating the offset of an element by position. + size_t CalculateOffsetOfElement(size_t position) const; + + // Helper to validate the preconditions for reading the buffer with an iterator. + void ValidateIteratorReadPreconditions() const; + + // Variables for span elements reading + tcb::span elements_span_; + mutable size_t num_elements_; + bool has_fixed_sized_elements_; + + // Variables for determining offset of elements. + size_t prefix_size_ = 0; + size_t element_size_; // for fixed-size elements + mutable std::vector offsets_; // for variable-size elements + + // Variables for write buffer. + std::vector write_buffer_; + + // Variable for sequential variable-size writes. + // Tracks next expected position for sequential variable-size writes. + // Value is invalidated to kUnsetVariableElementOffset once order is violated. + size_t next_expected_write_position_ = 0; + +private: + // Initialization methods and flags for read-only buffer + void InitializeFromSpan() const; + void EnsureInitializedFromSpan() const; + mutable bool is_initialized_from_span_ = false; + + // Initialization methods and flags for write buffer + void InitializeForWriteBuffer(size_t variable_size_reserved_bytes_hint); + void RebindSpanToWriteBuffer(); + bool is_write_buffer_initialized_ = false; + bool is_write_buffer_finalized_ = false; +}; + +// Constant to mark an offset value as unset. +inline constexpr size_t kUnsetVariableElementOffset = std::numeric_limits::max(); + +// Constant for the size of the [u32 size] prefix for variable-size elements. +inline constexpr size_t kSizePrefixBytes = sizeof(uint32_t); + +} // namespace dbps::processing diff --git a/src/processing/byte_buffer_test.cpp b/src/processing/byte_buffer_test.cpp new file mode 100644 index 0000000..bdf16cd --- /dev/null +++ b/src/processing/byte_buffer_test.cpp @@ -0,0 +1,971 @@ +// 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 "byte_buffer.h" + +#include +#include +#include + +#include "bytes_utils.h" +#include "exceptions.h" + +using dbps::processing::ByteBuffer; +using dbps::processing::kUnsetVariableElementOffset; + +class ByteBufferTestProxy : public ByteBuffer { +public: + ByteBufferTestProxy( + tcb::span elements_span) + : ByteBuffer(elements_span) {} + ByteBufferTestProxy( + tcb::span elements_span, + size_t element_size, + size_t prefix_size = 0) + : ByteBuffer(elements_span, element_size, prefix_size) {} + ByteBufferTestProxy(size_t num_elements) + : ByteBuffer(num_elements, 0, false) {} + ByteBufferTestProxy( + size_t num_elements, + size_t reserved_bytes_hint, + bool use_reserve_hint, + size_t prefix_size = 0) + : ByteBuffer(num_elements, reserved_bytes_hint, use_reserve_hint, prefix_size) {} + ByteBufferTestProxy(size_t num_elements, size_t element_size, size_t prefix_size = 0) + : ByteBuffer(num_elements, element_size, prefix_size) {} + using ByteBuffer::EstimateOffsetsReserveCountFromSample; + + size_t GetNumElements() const { return num_elements_; } + bool GetHasFixedSizedElements() const { return has_fixed_sized_elements_; } + size_t GetElementSize() const { return element_size_; } + const std::vector& GetOffsets() const { return offsets_; } + const std::vector& GetWriteBuffer() const { return write_buffer_; } + void SetWriteBufferByteForTest(size_t idx, uint8_t value) { + write_buffer_[idx] = value; + } + void AppendTrailingBytesForTest(tcb::span bytes) { + write_buffer_.insert(write_buffer_.end(), bytes.begin(), bytes.end()); + } +}; + +namespace { + +void ExpectCommonState( + const ByteBufferTestProxy& buffer, + size_t expected_num_elements, + bool expected_fixed_size_flag, + size_t expected_element_size) { + EXPECT_EQ(buffer.GetNumElements(), expected_num_elements); + EXPECT_EQ(buffer.GetHasFixedSizedElements(), expected_fixed_size_flag); + EXPECT_EQ(buffer.GetElementSize(), expected_element_size); +} + +std::vector MakePayload(size_t size, uint8_t seed) { + std::vector payload(size); + for (size_t i = 0; i < size; ++i) { + payload[i] = static_cast(seed + static_cast(i)); + } + return payload; +} + +} // namespace + +TEST(ByteBufferTest, ConstructFixedSize_ValidBuffer_InitializesExpectedState) { + std::vector bytes = {0x01, 0x02, 0x03, 0x04, 0x05, 0x06}; + ByteBufferTestProxy buffer(tcb::span(bytes), 2); + // Trigger lazy span initialization. + EXPECT_NO_THROW((void)buffer.GetElement(0)); + ExpectCommonState(buffer, 3u, true, 2u); + EXPECT_TRUE(buffer.GetOffsets().empty()); +} + +TEST(ByteBufferTest, GetElement_FixedSize_ReturnsExpectedSlices) { + std::vector bytes = {0x10, 0x11, 0x20, 0x21, 0x30, 0x31}; + ByteBufferTestProxy buffer(tcb::span(bytes), 2); + + const auto first = buffer.GetElement(0); + const auto second = buffer.GetElement(1); + const auto third = buffer.GetElement(2); + + ASSERT_EQ(first.size(), 2u); + ASSERT_EQ(second.size(), 2u); + ASSERT_EQ(third.size(), 2u); + EXPECT_EQ(first[0], 0x10); + EXPECT_EQ(first[1], 0x11); + EXPECT_EQ(second[0], 0x20); + EXPECT_EQ(second[1], 0x21); + EXPECT_EQ(third[0], 0x30); + EXPECT_EQ(third[1], 0x31); +} + +TEST(ByteBufferTest, GetElement_FixedSize_WithPrefixSize_SkipsPrefix) { + std::vector bytes = { + 0xFE, 0xFD, 0xFC, // prefix + 0x10, 0x11, 0x20, 0x21, 0x30, 0x31 + }; + ByteBufferTestProxy buffer(tcb::span(bytes), 2u, 3u); + + const auto first = buffer.GetElement(0); + const auto second = buffer.GetElement(1); + const auto third = buffer.GetElement(2); + + EXPECT_EQ(std::vector(first.begin(), first.end()), (std::vector{0x10, 0x11})); + EXPECT_EQ(std::vector(second.begin(), second.end()), (std::vector{0x20, 0x21})); + EXPECT_EQ(std::vector(third.begin(), third.end()), (std::vector{0x30, 0x31})); +} + +TEST(ByteBufferTest, GetElement_VariableSize_WithPrefixSize_SkipsPrefix) { + std::vector bytes = { + 0xEE, 0xDD, // prefix + 0x02, 0x00, 0x00, 0x00, 0x41, 0x42, + 0x03, 0x00, 0x00, 0x00, 0x78, 0x79, 0x7A + }; + ByteBuffer buffer(tcb::span(bytes), 2u); + + const auto first = buffer.GetElement(0); + const auto second = buffer.GetElement(1); + + EXPECT_EQ(std::vector(first.begin(), first.end()), (std::vector{0x41, 0x42})); + EXPECT_EQ(std::vector(second.begin(), second.end()), (std::vector{0x78, 0x79, 0x7A})); +} + +TEST(ByteBufferTest, ConstructFixedSize_ZeroElementSize_Throws) { + std::vector bytes = {0x01, 0x02, 0x03, 0x04}; + ByteBufferTestProxy buffer(tcb::span(bytes), 0); + EXPECT_THROW((void)buffer.GetElement(0), InvalidInputException); +} + +TEST(ByteBufferTest, ConstructFixedSize_NonDivisibleSize_Throws) { + std::vector bytes = {0x01, 0x02, 0x03}; + ByteBufferTestProxy buffer(tcb::span(bytes), 2); + EXPECT_THROW((void)buffer.GetElement(0), InvalidInputException); +} + +TEST(ByteBufferTest, ConstructVariableSize_ValidEncodedBuffer_InitializesExpectedState) { + // [len=5]["ABCDE"][len=7]["1234567"] + std::vector bytes = { + 0x05, 0x00, 0x00, 0x00, 0x41, 0x42, 0x43, 0x44, 0x45, + 0x07, 0x00, 0x00, 0x00, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37 + }; + ByteBufferTestProxy buffer{tcb::span(bytes)}; + // Trigger lazy variable-size index parsing. + EXPECT_NO_THROW((void)buffer.GetElement(0)); + ExpectCommonState(buffer, 2u, false, 0u); + ASSERT_EQ(buffer.GetOffsets().size(), 2u); + EXPECT_EQ(buffer.GetOffsets()[0], 0u); + EXPECT_EQ(buffer.GetOffsets()[1], 9u); // 4 bytes length prefix + 5 bytes first payload. +} + +TEST(ByteBufferTest, EstimateOffsetsReserveCountFromSample_MultipleCases) { + const auto make_variable_size_bytes = [](const std::vector& sizes) { + std::vector bytes; + for (size_t idx = 0; idx < sizes.size(); ++idx) { + append_u32_le(bytes, static_cast(sizes[idx])); + for (size_t j = 0; j < sizes[idx]; ++j) { + bytes.push_back(static_cast((idx + j) & 0xFF)); + } + } + return bytes; + }; + + // Empty buffer. + { + const std::vector sizes = {}; + const std::vector bytes = make_variable_size_bytes(sizes); + const size_t estimated = ByteBufferTestProxy::EstimateOffsetsReserveCountFromSample(tcb::span(bytes)); + EXPECT_EQ(estimated, 0u); + } + + // Buffer with less than sample size (5 vs 10): exact count, no extrapolation/headroom needed. + { + const std::vector sizes = {1u, 2u, 50u, 4u, 7u}; + const std::vector bytes = make_variable_size_bytes(sizes); + const size_t estimated = ByteBufferTestProxy::EstimateOffsetsReserveCountFromSample(tcb::span(bytes)); + EXPECT_EQ(estimated, sizes.size()); + } + + // Buffer with uniform record sizes. Base estimate is exact, then +10% headroom is applied. + { + // 30 elements, each 6 bytes long. + const size_t num_elements = 30u; + const size_t payload_size = 6u; + std::vector bytes; + for (size_t i = 0; i < num_elements; ++i) { + append_u32_le(bytes, static_cast(payload_size)); + for (size_t j = 0; j < payload_size; ++j) { + bytes.push_back(static_cast((i + j) & 0xFF)); + } + } + const size_t estimate = ByteBufferTestProxy::EstimateOffsetsReserveCountFromSample(tcb::span(bytes)); + EXPECT_EQ(estimate, 33u); //33 = 30 elements + 10% headroom. + } + + // Buffer to estimate intentionally overshoots the true element count. + { + // First 10 elements are very small, tail elements are very large: + const std::vector sizes = {1u, 1u, 1u, 1u, 1u, 1u, 1u, 1u, 1u, 1u, 100u, 100u}; + const std::vector bytes = make_variable_size_bytes(sizes); + const size_t estimate = ByteBufferTestProxy::EstimateOffsetsReserveCountFromSample(tcb::span(bytes)); + EXPECT_GT(estimate, sizes.size()); + } +} + +TEST(ByteBufferTest, GetElement_VariableSize_ReturnsExpectedPayload) { + // [len=5]["ABCDE"][len=7]["1234567"] + std::vector bytes = { + 0x05, 0x00, 0x00, 0x00, 0x41, 0x42, 0x43, 0x44, 0x45, + 0x07, 0x00, 0x00, 0x00, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37 + }; + ByteBufferTestProxy buffer{tcb::span(bytes)}; + + const auto first = buffer.GetElement(0); + const auto second = buffer.GetElement(1); + + ASSERT_EQ(first.size(), 5u); + ASSERT_EQ(second.size(), 7u); + EXPECT_EQ(first[0], 0x41); + EXPECT_EQ(first[4], 0x45); + EXPECT_EQ(second[0], 0x31); + EXPECT_EQ(second[6], 0x37); +} + +TEST(ByteBufferTest, Iterate_ReadOnlyVariableSize_WithPrefixSize_SkipsPrefix) { + std::vector bytes = { + 0xEE, 0xDD, // prefix + 0x02, 0x00, 0x00, 0x00, 0x41, 0x42, + 0x03, 0x00, 0x00, 0x00, 0x78, 0x79, 0x7A + }; + ByteBuffer buffer(tcb::span(bytes), 2u); + + std::vector> collected; + for (const auto element : buffer) { + collected.push_back(std::vector(element.begin(), element.end())); + } + + ASSERT_EQ(collected.size(), 2u); + EXPECT_EQ(collected[0], (std::vector{0x41, 0x42})); + EXPECT_EQ(collected[1], (std::vector{0x78, 0x79, 0x7A})); +} + +TEST(ByteBufferTest, Iterate_ReadOnlyFixedSize_WithPrefixSize_SkipsPrefix) { + std::vector bytes = { + 0xFE, 0xFD, 0xFC, // prefix + 0x10, 0x11, 0x20, 0x21, 0x30, 0x31 + }; + ByteBufferTestProxy buffer(tcb::span(bytes), 2u, 3u); + + std::vector> collected; + for (const auto element : buffer) { + collected.push_back(std::vector(element.begin(), element.end())); + } + + ASSERT_EQ(collected.size(), 3u); + EXPECT_EQ(collected[0], (std::vector{0x10, 0x11})); + EXPECT_EQ(collected[1], (std::vector{0x20, 0x21})); + EXPECT_EQ(collected[2], (std::vector{0x30, 0x31})); +} + +TEST(ByteBufferTest, Iterate_ReadOnlyFixedSize_TraversesInOrder) { + std::vector bytes = {0x10, 0x11, 0x20, 0x21, 0x30, 0x31}; + ByteBufferTestProxy buffer(tcb::span(bytes), 2); + + std::vector> collected; + for (const auto element : buffer) { + collected.push_back(std::vector(element.begin(), element.end())); + } + + ASSERT_EQ(collected.size(), 3u); + EXPECT_EQ(collected[0], (std::vector{0x10, 0x11})); + EXPECT_EQ(collected[1], (std::vector{0x20, 0x21})); + EXPECT_EQ(collected[2], (std::vector{0x30, 0x31})); + + // Check InitializeFromSpan was not called yet. + EXPECT_TRUE(buffer.GetOffsets().empty()); + EXPECT_EQ(buffer.GetNumElements(), 0); + + // Now read the elemennts through GetElement. It calls the InitializeFromSpan method. + std::vector> collected_with_get_element; + for (size_t i = 0; i < collected.size(); ++i) { + const auto element = buffer.GetElement(i); + collected_with_get_element.push_back(std::vector(element.begin(), element.end())); + } + + // Check elements read with GetElement are the same as the ones collected with the iterator. + EXPECT_EQ(collected_with_get_element, collected); + + // Offsets should still be empty after the iteration because these fixed-size elements. + // However, the number of elements is now known. + EXPECT_TRUE(buffer.GetOffsets().empty()); + EXPECT_EQ(buffer.GetNumElements(), 3u); + +} + +TEST(ByteBufferTest, Iterate_ReadOnlyVariableSize_TraversesInOrder) { + // [len=5]["ABCDE"][len=7]["1234567"] + std::vector bytes = { + 0x05, 0x00, 0x00, 0x00, 0x41, 0x42, 0x43, 0x44, 0x45, + 0x07, 0x00, 0x00, 0x00, 0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37 + }; + ByteBufferTestProxy buffer{tcb::span(bytes)}; + + std::vector> collected; + for (const auto element : buffer) { + collected.push_back(std::vector(element.begin(), element.end())); + } + + ASSERT_EQ(collected.size(), 2u); + EXPECT_EQ(collected[0], (std::vector{0x41, 0x42, 0x43, 0x44, 0x45})); + EXPECT_EQ(collected[1], (std::vector{0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37})); + + // Check InitializeFromSpan was not called yet. + EXPECT_TRUE(buffer.GetOffsets().empty()); + EXPECT_EQ(buffer.GetNumElements(), 0); + + // Now read the elemennts through GetElement. It calls the InitializeFromSpan method. + std::vector> collected_with_get_element; + for (size_t i = 0; i < collected.size(); ++i) { + const auto element = buffer.GetElement(i); + collected_with_get_element.push_back(std::vector(element.begin(), element.end())); + } + + // Check elements read with GetElement are the same as the ones collected with the iterator. + EXPECT_EQ(collected_with_get_element, collected); + + // Offsets and number of elements should now be initialized. + ASSERT_EQ(buffer.GetOffsets().size(), 2u); + EXPECT_EQ(buffer.GetOffsets()[0], 0u); + EXPECT_EQ(buffer.GetOffsets()[1], 9u); // 4 bytes length prefix + 5 bytes first payload. + EXPECT_EQ(buffer.GetNumElements(), 2u); + +} + +TEST(ByteBufferTest, TransformFixedSize_ReadIterateWriteFinalize_RoundTrip) { + constexpr size_t kElementSize = 2u; + std::vector source_bytes = { + 0x10, 0x11, + 0x20, 0x21, + 0x30, 0x31 + }; + + // Create a source buffer to read the elements from. + ByteBufferTestProxy source_buffer(tcb::span(source_bytes), kElementSize); + + // Create a new buffer to write the transformed elements. + ByteBufferTestProxy transformed_buffer(3u, kElementSize); + size_t position = 0; + for (const auto element : source_buffer) { + std::vector source_element(element.begin(), element.end()); + std::vector transformed_element = source_element; + transformed_element[0] = static_cast(transformed_element[0] + 1u); + transformed_buffer.SetElement(position, tcb::span(transformed_element)); + ++position; + } + + ASSERT_EQ(position, 3u); + + // Compare source and transformed buffers before finalization. + for (size_t i = 0; i < position; ++i) { + const auto source_element = source_buffer.GetElement(i); + const auto transformed_element = transformed_buffer.GetElement(i); + + std::vector expected_transformed(source_element.begin(), source_element.end()); + expected_transformed[0] = static_cast(expected_transformed[0] + 1u); + + EXPECT_EQ(std::vector(transformed_element.begin(), transformed_element.end()), expected_transformed); + // Check that transformed element differs from source after applying the mutation. + EXPECT_NE(std::vector(source_element.begin(), source_element.end()), + std::vector(transformed_element.begin(), transformed_element.end())); + } + + // Now finalize the transformed buffer and populate a third buffer to read the elements from. + std::vector finalized_bytes = transformed_buffer.FinalizeAndTakeBuffer(); + ByteBufferTestProxy finalized_read_buffer(tcb::span(finalized_bytes), kElementSize); + + // Compare source and finalized read buffer using the same transformation rule. + for (size_t i = 0; i < position; ++i) { + const auto source_element = source_buffer.GetElement(i); + const auto finalized_element = finalized_read_buffer.GetElement(i); + + std::vector expected_transformed(source_element.begin(), source_element.end()); + expected_transformed[0] = static_cast(expected_transformed[0] + 1u); + + EXPECT_EQ(std::vector(finalized_element.begin(), finalized_element.end()), expected_transformed); + EXPECT_NE(std::vector(source_element.begin(), source_element.end()), + std::vector(finalized_element.begin(), finalized_element.end())); + } +} + +TEST(ByteBufferTest, TransformVariableSize_ReadIterateWriteFinalize_RoundTrip) { + std::vector source_bytes; + append_u32_le(source_bytes, 5u); + source_bytes.insert(source_bytes.end(), {0x41, 0x42, 0x43, 0x44, 0x45}); // "ABCDE" + append_u32_le(source_bytes, 7u); + source_bytes.insert(source_bytes.end(), {0x31, 0x32, 0x33, 0x34, 0x35, 0x36, 0x37}); // "1234567" + append_u32_le(source_bytes, 3u); + source_bytes.insert(source_bytes.end(), {0x61, 0x62, 0x63}); // "abc" + + ByteBufferTestProxy source_buffer{tcb::span(source_bytes)}; + ByteBufferTestProxy transformed_buffer(3u, source_bytes.size(), true); + size_t position = 0; + + for (const auto element : source_buffer) { + std::vector source_element(element.begin(), element.end()); + std::vector transformed_element = source_element; + transformed_element[0] = static_cast(transformed_element[0] + 1u); + transformed_buffer.SetElement(position, tcb::span(transformed_element)); + ++position; + } + + ASSERT_EQ(position, 3u); + + // Compare source and transformed buffers before finalization. + for (size_t i = 0; i < position; ++i) { + const auto source_element = source_buffer.GetElement(i); + const auto transformed_element = transformed_buffer.GetElement(i); + + std::vector expected_transformed(source_element.begin(), source_element.end()); + expected_transformed[0] = static_cast(expected_transformed[0] + 1u); + + EXPECT_EQ(std::vector(transformed_element.begin(), transformed_element.end()), expected_transformed); + EXPECT_NE(std::vector(source_element.begin(), source_element.end()), + std::vector(transformed_element.begin(), transformed_element.end())); + } + + std::vector finalized_bytes = transformed_buffer.FinalizeAndTakeBuffer(); + ByteBufferTestProxy finalized_read_buffer{tcb::span(finalized_bytes)}; + + // Compare source and finalized read buffer using the same transformation rule. + for (size_t i = 0; i < position; ++i) { + const auto source_element = source_buffer.GetElement(i); + const auto finalized_element = finalized_read_buffer.GetElement(i); + + std::vector expected_transformed(source_element.begin(), source_element.end()); + expected_transformed[0] = static_cast(expected_transformed[0] + 1u); + + EXPECT_EQ(std::vector(finalized_element.begin(), finalized_element.end()), expected_transformed); + EXPECT_NE(std::vector(source_element.begin(), source_element.end()), + std::vector(finalized_element.begin(), finalized_element.end())); + } +} + +TEST(ByteBufferTest, Iterate_ReadOnlyEmptySpan_VisitsNoElements) { + std::vector empty_bytes; + + // Fixed-size empty span. + ByteBufferTestProxy fixed_buffer(tcb::span(empty_bytes), 2); + size_t fixed_count = 0; + for (const auto element : fixed_buffer) { + (void)element; + ++fixed_count; + } + EXPECT_EQ(fixed_count, 0u); + + // Variable-size empty span. + ByteBufferTestProxy variable_buffer{tcb::span(empty_bytes)}; + size_t variable_count = 0; + for (const auto element : variable_buffer) { + (void)element; + ++variable_count; + } + EXPECT_EQ(variable_count, 0u); +} + +TEST(ByteBufferTest, GetElement_OutOfRange_Throws) { + std::vector bytes = {0x01, 0x02, 0x03, 0x04}; + ByteBufferTestProxy buffer(tcb::span(bytes), 2); + EXPECT_THROW((void)buffer.GetElement(2), InvalidInputException); +} + +TEST(ByteBufferTest, ConstructWithNumElements_FixedSize_AllocatesAndSets) { + ByteBufferTestProxy buffer(3u, 2u); + EXPECT_EQ(buffer.GetNumElements(), 3u); + EXPECT_TRUE(buffer.GetHasFixedSizedElements()); + EXPECT_EQ(buffer.GetElementSize(), 2u); + EXPECT_TRUE(buffer.GetOffsets().empty()); + + std::vector first = {0xAA, 0xBB}; + std::vector third = {0xCC, 0xDD}; + buffer.SetElement(0, tcb::span(first)); + buffer.SetElement(2, tcb::span(third)); + + const auto read_first = buffer.GetElement(0); + const auto read_third = buffer.GetElement(2); + ASSERT_EQ(read_first.size(), 2u); + ASSERT_EQ(read_third.size(), 2u); + EXPECT_EQ(read_first[0], 0xAA); + EXPECT_EQ(read_first[1], 0xBB); + EXPECT_EQ(read_third[0], 0xCC); + EXPECT_EQ(read_third[1], 0xDD); +} + +TEST(ByteBufferTest, ConstructWithNumElements_FixedSize_WithPrefixSize_PreservesLeadingBytes) { + ByteBufferTestProxy buffer(2u, 3u, (size_t) 7u); + + std::vector first = {0xAA, 0xBB, 0xBC}; + std::vector second = {0xCC, 0xDD, 0xDE}; + + // Write elements out of order to trigger the defragmentation path. + buffer.SetElement(1, tcb::span(second)); + buffer.SetElement(0, tcb::span(first)); + + // Manually annotate prefix bytes and verify they are preserved after finalize. + buffer.SetWriteBufferByteForTest(0u, 0xF0); + buffer.SetWriteBufferByteForTest(1u, 0x0D); + buffer.SetWriteBufferByteForTest(2u, 0xBE); + + + std::vector final_buffer = buffer.FinalizeAndTakeBuffer(); + ASSERT_EQ(final_buffer.size(), 13u); + EXPECT_EQ(final_buffer[0], 0xF0); + EXPECT_EQ(final_buffer[1], 0x0D); + EXPECT_EQ(final_buffer[2], 0xBE); + EXPECT_EQ(final_buffer[7], 0xAA); + EXPECT_EQ(final_buffer[8], 0xBB); + EXPECT_EQ(final_buffer[9], 0xBC); + EXPECT_EQ(final_buffer[10], 0xCC); + EXPECT_EQ(final_buffer[11], 0xDD); + EXPECT_EQ(final_buffer[12], 0xDE); +} + +TEST(ByteBufferTest, ConstructWithNumElements_VariableSize_AllocatesAndSets) { + ByteBufferTestProxy buffer(2u, 8u, true); + EXPECT_EQ(buffer.GetNumElements(), 2u); + EXPECT_FALSE(buffer.GetHasFixedSizedElements()); + EXPECT_EQ(buffer.GetElementSize(), 0u); + ASSERT_EQ(buffer.GetOffsets().size(), 2u); + EXPECT_EQ(buffer.GetOffsets()[0], kUnsetVariableElementOffset); + EXPECT_EQ(buffer.GetOffsets()[1], kUnsetVariableElementOffset); + + std::vector first = {0x10, 0x11, 0x12, 0x13, 0x14}; + std::vector second = {0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26}; + buffer.SetElement(0, tcb::span(first)); + buffer.SetElement(1, tcb::span(second)); + + const auto read_first = buffer.GetElement(0); + const auto read_second = buffer.GetElement(1); + ASSERT_EQ(read_first.size(), first.size()); + ASSERT_EQ(read_second.size(), second.size()); + EXPECT_EQ(read_first[0], 0x10); + EXPECT_EQ(read_first[4], 0x14); + EXPECT_EQ(read_second[0], 0x20); + EXPECT_EQ(read_second[6], 0x26); +} + +TEST(ByteBufferTest, ConstructWithNumElements_VariableSize_WithPrefixSize_PreservesLeadingBytes) { + size_t prefix_size = 3u; + ByteBufferTestProxy buffer(2u, 8u, true, prefix_size); + + std::vector first = {0x10, 0x11}; + std::vector second = {0x20, 0x21, 0x22, 0x23, 0x24}; + + // Write elements out of order to trigger the defragmentation path. + buffer.SetElement(1, tcb::span(second)); + buffer.SetElement(0, tcb::span(first)); + + // Manually annotate prefix bytes and verify they are preserved after finalize. + buffer.SetWriteBufferByteForTest(0u, 0xF0); + buffer.SetWriteBufferByteForTest(1u, 0x0D); + buffer.SetWriteBufferByteForTest(2u, 0xBE); + + std::vector final_buffer = buffer.FinalizeAndTakeBuffer(); + ASSERT_EQ(final_buffer.size(), prefix_size + 4 + first.size() + 4 + second.size()); + EXPECT_EQ(final_buffer[0], 0xF0); + EXPECT_EQ(final_buffer[1], 0x0D); + EXPECT_EQ(final_buffer[2], 0xBE); + + // First record starts right after the configured prefix_size. + EXPECT_EQ(read_u32_le(final_buffer, prefix_size), first.size()); + EXPECT_EQ(final_buffer[7], 0x10); + EXPECT_EQ(final_buffer[8], 0x11); + + // Second record follows contiguously. + EXPECT_EQ(read_u32_le(final_buffer, prefix_size + 4 + first.size()), second.size()); + EXPECT_EQ(final_buffer[13], 0x20); + EXPECT_EQ(final_buffer[14], 0x21); + EXPECT_EQ(final_buffer[15], 0x22); + EXPECT_EQ(final_buffer[16], 0x23); + EXPECT_EQ(final_buffer[17], 0x24); +} + +TEST(ByteBufferTest, GetElement_VariableSize_UnsetPosition_Throws) { + ByteBufferTestProxy buffer(2u, 8u, true); + std::vector second = {0x20, 0x21, 0x22}; + buffer.SetElement(1, tcb::span(second)); + + EXPECT_THROW((void)buffer.GetElement(0), InvalidInputException); +} + +TEST(ByteBufferTest, SetElement_VariableSize_OutOfOrderAndOverwrite_ReturnsLatestValues) { + ByteBufferTestProxy buffer(3u, 32u, true); + + std::vector v2_first = {0xA0, 0xA1}; + std::vector v0_first = {0xB0, 0xB1, 0xB2}; + std::vector v1_first = {0xC0}; + std::vector v0_second = {0xD0, 0xD1, 0xD2, 0xD3}; + std::vector v2_second = {0xE0, 0xE1, 0xE2}; + + // Write in non-sequential order. + buffer.SetElement(2, tcb::span(v2_first)); + buffer.SetElement(0, tcb::span(v0_first)); + buffer.SetElement(1, tcb::span(v1_first)); + + const auto e0_before = buffer.GetElement(0); + const auto e1_before = buffer.GetElement(1); + const auto e2_before = buffer.GetElement(2); + ASSERT_EQ(e0_before.size(), v0_first.size()); + ASSERT_EQ(e1_before.size(), v1_first.size()); + ASSERT_EQ(e2_before.size(), v2_first.size()); + for (size_t i = 0; i < v0_first.size(); ++i) { + EXPECT_EQ(e0_before[i], v0_first[i]); + } + for (size_t i = 0; i < v1_first.size(); ++i) { + EXPECT_EQ(e1_before[i], v1_first[i]); + } + for (size_t i = 0; i < v2_first.size(); ++i) { + EXPECT_EQ(e2_before[i], v2_first[i]); + } + + const size_t offset0_before_overwrite = buffer.GetOffsets()[0]; + const size_t offset2_before_overwrite = buffer.GetOffsets()[2]; + + // Overwrite previously written positions; latest append should win. + buffer.SetElement(0, tcb::span(v0_second)); + buffer.SetElement(2, tcb::span(v2_second)); + + const size_t expected_write_buffer_size = + (4u + v2_first.size()) + + (4u + v0_first.size()) + + (4u + v1_first.size()) + + (4u + v0_second.size()) + + (4u + v2_second.size()); + // Check that append-only writes include overwritten records too. + EXPECT_EQ(buffer.GetWriteBuffer().size(), expected_write_buffer_size); + + EXPECT_GT(buffer.GetOffsets()[0], offset0_before_overwrite); + EXPECT_GT(buffer.GetOffsets()[2], offset2_before_overwrite); + + const auto e0 = buffer.GetElement(0); + const auto e1 = buffer.GetElement(1); + const auto e2 = buffer.GetElement(2); + + ASSERT_EQ(e0.size(), v0_second.size()); + ASSERT_EQ(e1.size(), v1_first.size()); + ASSERT_EQ(e2.size(), v2_second.size()); + + for (size_t i = 0; i < v0_second.size(); ++i) { + EXPECT_EQ(e0[i], v0_second[i]); + } + for (size_t i = 0; i < v1_first.size(); ++i) { + EXPECT_EQ(e1[i], v1_first[i]); + } + for (size_t i = 0; i < v2_second.size(); ++i) { + EXPECT_EQ(e2[i], v2_second[i]); + } +} + +TEST(ByteBufferTest, VariableSizeWrite_ExactHint_NoReallocationAndExactUsedSize) { + const std::vector payload_sizes = {37u, 11u, 47u, 23u, 43u, 17u, 31u}; + const size_t num_elements = payload_sizes.size(); + + size_t total_payload_bytes = 0; + for (size_t size : payload_sizes) { + total_payload_bytes += size; + } + const size_t exact_hint_bytes = (num_elements * 4u) + total_payload_bytes; + + ByteBufferTestProxy buffer(num_elements, exact_hint_bytes, true); + const size_t initial_capacity = buffer.GetWriteBuffer().capacity(); + const uint8_t* const initial_data_ptr = buffer.GetWriteBuffer().data(); + + std::vector> payloads; + payloads.reserve(num_elements); + for (size_t i = 0; i < num_elements; ++i) { + payloads.push_back(MakePayload(payload_sizes[i], static_cast(0x10 + i))); + buffer.SetElement(i, tcb::span(payloads.back())); + } + + EXPECT_EQ(buffer.GetWriteBuffer().size(), exact_hint_bytes); + EXPECT_EQ(buffer.GetWriteBuffer().capacity(), initial_capacity); + EXPECT_EQ(buffer.GetWriteBuffer().capacity(), exact_hint_bytes); + EXPECT_EQ(buffer.GetWriteBuffer().data(), initial_data_ptr); + + for (size_t i = 0; i < num_elements; ++i) { + const auto value = buffer.GetElement(i); + ASSERT_EQ(value.size(), payloads[i].size()); + for (size_t j = 0; j < payloads[i].size(); ++j) { + EXPECT_EQ(value[j], payloads[i][j]); + } + } + + const uint8_t* const data_ptr_before_finalize = buffer.GetWriteBuffer().data(); + std::vector final_buffer = buffer.FinalizeAndTakeBuffer(); + EXPECT_EQ(final_buffer.data(), data_ptr_before_finalize); // Same allocation: finalize returned write_buffer_ as-is. +} + +TEST(ByteBufferTest, VariableSizeWrite_ExceedsHint_ReallocatesBuffer) { + const size_t num_elements = 7u; + ByteBufferTestProxy buffer(num_elements, 32u, true); + + const size_t initial_capacity = buffer.GetWriteBuffer().capacity(); + const uint8_t* const initial_data_ptr = buffer.GetWriteBuffer().data(); + + std::vector> payloads; + payloads.reserve(num_elements); + for (size_t i = 0; i < num_elements; ++i) { + payloads.push_back(MakePayload(64u + (i * 9u), static_cast(0x40 + i))); + buffer.SetElement(i, tcb::span(payloads.back())); + } + + EXPECT_GT(buffer.GetWriteBuffer().size(), initial_capacity); + EXPECT_GT(buffer.GetWriteBuffer().capacity(), initial_capacity); + EXPECT_NE(buffer.GetWriteBuffer().data(), initial_data_ptr); + + for (size_t i = 0; i < num_elements; ++i) { + const auto value = buffer.GetElement(i); + ASSERT_EQ(value.size(), payloads[i].size()); + for (size_t j = 0; j < payloads[i].size(); ++j) { + EXPECT_EQ(value[j], payloads[i][j]); + } + } +} + +TEST(ByteBufferTest, FinalizeAndTakeBuffer_VariableSize_PartialWrite_ThrowsAndAllowsRetry) { + ByteBufferTestProxy buffer(2u, 8u, true); + std::vector first = {0x10, 0x11}; + std::vector second = {0x20, 0x21, 0x22}; + buffer.SetElement(1, tcb::span(second)); + + EXPECT_THROW((void)buffer.FinalizeAndTakeBuffer(), InvalidInputException); + + buffer.SetElement(0, tcb::span(first)); + const uint8_t* const data_ptr_before = buffer.GetWriteBuffer().data(); + std::vector final_buffer = buffer.FinalizeAndTakeBuffer(); + EXPECT_NE(final_buffer.data(), data_ptr_before); // Different allocation (defragmented after retry). + + ByteBufferTestProxy read_back{tcb::span(final_buffer)}; + const auto r0 = read_back.GetElement(0); + const auto r1 = read_back.GetElement(1); + ASSERT_EQ(r0.size(), first.size()); + ASSERT_EQ(r1.size(), second.size()); + EXPECT_EQ(r0[0], 0x10); + EXPECT_EQ(r0[1], 0x11); + EXPECT_EQ(r1[0], 0x20); + EXPECT_EQ(r1[2], 0x22); +} + +TEST(ByteBufferTest, FinalizeAndTakeBuffer_VariableSize_Sequential_ReturnsAsIs) { + ByteBufferTestProxy buffer(3u, 64u, true); + std::vector first = {0x10, 0x11}; + std::vector second = {0x20, 0x21, 0x22}; + std::vector third = {0x30}; + + buffer.SetElement(0, tcb::span(first)); + buffer.SetElement(1, tcb::span(second)); + buffer.SetElement(2, tcb::span(third)); + + const std::vector raw_before_finalize = buffer.GetWriteBuffer(); + const uint8_t* const data_ptr_before = buffer.GetWriteBuffer().data(); + std::vector final_buffer = buffer.FinalizeAndTakeBuffer(); + + EXPECT_EQ(final_buffer, raw_before_finalize); // Same byte content. + EXPECT_EQ(final_buffer.data(), data_ptr_before); // Same allocation (moved, not copied). +} + +TEST(ByteBufferTest, FinalizeAndTakeBuffer_VariableSize_SequentialWithTrailingBytes_Throws) { + ByteBufferTestProxy buffer(3u, 64u, true); + std::vector first = {0x10, 0x11}; + std::vector second = {0x20, 0x21, 0x22}; + std::vector third = {0x30}; + + buffer.SetElement(0, tcb::span(first)); + buffer.SetElement(1, tcb::span(second)); + buffer.SetElement(2, tcb::span(third)); + + const size_t expected_trimmed_size = + (4u + first.size()) + (4u + second.size()) + (4u + third.size()); + + std::vector trailing = {0xEE, 0xEF, 0xF0}; + buffer.AppendTrailingBytesForTest(tcb::span(trailing)); + EXPECT_EQ(buffer.GetWriteBuffer().size(), expected_trimmed_size + trailing.size()); + + EXPECT_THROW((void)buffer.FinalizeAndTakeBuffer(), InvalidInputException); +} + +TEST(ByteBufferTest, FinalizeAndTakeBuffer_VariableSize_OutOfOrder_Defragments) { + ByteBufferTestProxy buffer(3u, 64u, true); + std::vector first = {0x10, 0x11}; + std::vector second = {0x20, 0x21, 0x22}; + std::vector third = {0x30}; + + buffer.SetElement(2, tcb::span(third)); + buffer.SetElement(0, tcb::span(first)); + buffer.SetElement(1, tcb::span(second)); + + const std::vector raw_before_finalize = buffer.GetWriteBuffer(); + const uint8_t* const data_ptr_before = buffer.GetWriteBuffer().data(); + std::vector final_buffer = buffer.FinalizeAndTakeBuffer(); + + EXPECT_NE(final_buffer, raw_before_finalize); // Different byte content. + EXPECT_NE(final_buffer.data(), data_ptr_before); // Different allocation (defragmented copy). + + ByteBufferTestProxy read_back{tcb::span(final_buffer)}; + const auto r0 = read_back.GetElement(0); + const auto r1 = read_back.GetElement(1); + const auto r2 = read_back.GetElement(2); + ASSERT_EQ(r0.size(), first.size()); + ASSERT_EQ(r1.size(), second.size()); + ASSERT_EQ(r2.size(), third.size()); + for (size_t i = 0; i < first.size(); ++i) { + EXPECT_EQ(r0[i], first[i]); + } + for (size_t i = 0; i < second.size(); ++i) { + EXPECT_EQ(r1[i], second[i]); + } + for (size_t i = 0; i < third.size(); ++i) { + EXPECT_EQ(r2[i], third[i]); + } +} + +TEST(ByteBufferTest, FinalizeAndTakeBuffer_VariableSize_Fragmented_Defragments) { + ByteBufferTestProxy buffer(2u, 64u, true); + std::vector first_initial = {0x10, 0x11}; + std::vector second = {0x20, 0x21, 0x22}; + std::vector first_overwrite = {0x30, 0x31, 0x32, 0x33}; + + buffer.SetElement(0, tcb::span(first_initial)); + buffer.SetElement(1, tcb::span(second)); + buffer.SetElement(0, tcb::span(first_overwrite)); + + const size_t raw_size_before_finalize = buffer.GetWriteBuffer().size(); + const std::vector raw_before_finalize = buffer.GetWriteBuffer(); + const uint8_t* const data_ptr_before = buffer.GetWriteBuffer().data(); + std::vector final_buffer = buffer.FinalizeAndTakeBuffer(); + + EXPECT_LT(final_buffer.size(), raw_size_before_finalize); + EXPECT_NE(final_buffer, raw_before_finalize); + EXPECT_NE(final_buffer.data(), data_ptr_before); // Different allocation (defragmented copy). + + ByteBufferTestProxy read_back{tcb::span(final_buffer)}; + const auto r0 = read_back.GetElement(0); + const auto r1 = read_back.GetElement(1); + ASSERT_EQ(r0.size(), first_overwrite.size()); + ASSERT_EQ(r1.size(), second.size()); + for (size_t i = 0; i < first_overwrite.size(); ++i) { + EXPECT_EQ(r0[i], first_overwrite[i]); + } + for (size_t i = 0; i < second.size(); ++i) { + EXPECT_EQ(r1[i], second[i]); + } +} + +TEST(ByteBufferTest, SetElement_FixedSize_OutOfOrderAndOverwrite_ReturnsLatestValues) { + ByteBufferTestProxy buffer(3u, 2u); + + std::vector v2_first = {0xA0, 0xA1}; + std::vector v0_first = {0xB0, 0xB1}; + std::vector v1_first = {0xC0, 0xC1}; + std::vector v0_second = {0xD0, 0xD1}; + std::vector v2_second = {0xE0, 0xE1}; + + // Write in non-sequential order. + buffer.SetElement(2, tcb::span(v2_first)); + buffer.SetElement(0, tcb::span(v0_first)); + buffer.SetElement(1, tcb::span(v1_first)); + + const auto e0_before = buffer.GetElement(0); + const auto e1_before = buffer.GetElement(1); + const auto e2_before = buffer.GetElement(2); + ASSERT_EQ(e0_before.size(), 2u); + ASSERT_EQ(e1_before.size(), 2u); + ASSERT_EQ(e2_before.size(), 2u); + EXPECT_EQ(e0_before[0], v0_first[0]); + EXPECT_EQ(e0_before[1], v0_first[1]); + EXPECT_EQ(e1_before[0], v1_first[0]); + EXPECT_EQ(e1_before[1], v1_first[1]); + EXPECT_EQ(e2_before[0], v2_first[0]); + EXPECT_EQ(e2_before[1], v2_first[1]); + + // Overwrite previously written positions; latest fixed-size bytes should win. + buffer.SetElement(0, tcb::span(v0_second)); + buffer.SetElement(2, tcb::span(v2_second)); + + const auto e0 = buffer.GetElement(0); + const auto e1 = buffer.GetElement(1); + const auto e2 = buffer.GetElement(2); + + ASSERT_EQ(e0.size(), 2u); + ASSERT_EQ(e1.size(), 2u); + ASSERT_EQ(e2.size(), 2u); + EXPECT_EQ(e0[0], v0_second[0]); + EXPECT_EQ(e0[1], v0_second[1]); + EXPECT_EQ(e1[0], v1_first[0]); + EXPECT_EQ(e1[1], v1_first[1]); + EXPECT_EQ(e2[0], v2_second[0]); + EXPECT_EQ(e2[1], v2_second[1]); + + const uint8_t* const data_ptr_before_finalize = buffer.GetWriteBuffer().data(); + std::vector final_buffer = buffer.FinalizeAndTakeBuffer(); + EXPECT_EQ(final_buffer.data(), data_ptr_before_finalize); // Fixed-size finalize should move write_buffer_ directly. + ASSERT_EQ(final_buffer.size(), 6u); + EXPECT_EQ(final_buffer[0], v0_second[0]); + EXPECT_EQ(final_buffer[1], v0_second[1]); + EXPECT_EQ(final_buffer[2], v1_first[0]); + EXPECT_EQ(final_buffer[3], v1_first[1]); + EXPECT_EQ(final_buffer[4], v2_second[0]); + EXPECT_EQ(final_buffer[5], v2_second[1]); +} + +TEST(ByteBufferTest, SetElement_FixedSize_WrongPayloadSize_Throws) { + ByteBufferTestProxy buffer(1u, 4u); + std::vector wrong = {0x01, 0x02}; + EXPECT_THROW((void)buffer.SetElement(0, tcb::span(wrong)), InvalidInputException); +} + +TEST(ByteBufferTest, SetElement_OnReadOnlyBuffer_Throws) { + std::vector bytes = {0x10, 0x11, 0x20, 0x21}; + ByteBufferTestProxy buffer(tcb::span(bytes), 2); + std::vector replacement = {0xAA, 0xBB}; + EXPECT_THROW((void)buffer.SetElement(0, tcb::span(replacement)), InvalidInputException); +} + +TEST(ByteBufferTest, Iterate_OnWriteBuffer_Throws) { + ByteBufferTestProxy buffer(3u, 2u); + EXPECT_THROW((void)buffer.begin(), InvalidInputException); + EXPECT_THROW((void)buffer.end(), InvalidInputException); +} + +TEST(ByteBufferTest, ConstructVariableSize_EmptyBuffer_InitializesEmptyState) { + std::vector bytes; + ByteBufferTestProxy buffer{tcb::span(bytes)}; + ExpectCommonState(buffer, 0u, false, 0u); + EXPECT_TRUE(buffer.GetOffsets().empty()); +} + +TEST(ByteBufferTest, ConstructVariableSize_TruncatedLengthPrefix_Throws) { + std::vector bytes = {0x01, 0x00, 0x00}; // only 3 bytes + ByteBufferTestProxy buffer{tcb::span(bytes)}; + EXPECT_THROW((void)buffer.GetElement(0), InvalidInputException); +} + +TEST(ByteBufferTest, ConstructVariableSize_TruncatedPayload_Throws) { + // Declares payload length 5, but provides only 2 bytes. + std::vector bytes = { + 0x05, 0x00, 0x00, 0x00, 0xAA, 0xBB + }; + ByteBufferTestProxy buffer{tcb::span(bytes)}; + EXPECT_THROW((void)buffer.GetElement(0), InvalidInputException); +}