libjxl

padded_bytes.h (6338B)

---

 // Copyright (c) the JPEG XL Project Authors. All rights reserved.
 //
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
 #ifndef LIB_JXL_BASE_PADDED_BYTES_H_
 #define LIB_JXL_BASE_PADDED_BYTES_H_
 
 // std::vector replacement with padding to reduce bounds checks in WriteBits
 
 #include <stddef.h>
 #include <stdint.h>
 #include <string.h>  // memcpy
 
 #include <algorithm>  // max
 #include <initializer_list>
 #include <utility>  // swap
 
 #include "lib/jxl/base/compiler_specific.h"
 #include "lib/jxl/base/status.h"
 #include "lib/jxl/cache_aligned.h"
 
 namespace jxl {
 
 // Provides a subset of the std::vector interface with some differences:
 // - allows BitWriter to write 64 bits at a time without bounds checking;
 // - ONLY zero-initializes the first byte (required by BitWriter);
 // - ensures cache-line alignment.
 class PaddedBytes {
  public:
   // Required for output params.
   PaddedBytes() : size_(0), capacity_(0) {}
 
   explicit PaddedBytes(size_t size) : size_(size), capacity_(0) {
     reserve(size);
   }
 
   PaddedBytes(size_t size, uint8_t value) : size_(size), capacity_(0) {
     reserve(size);
     if (size_ != 0) {
       memset(data(), value, size);
     }
   }
 
   PaddedBytes(const PaddedBytes& other) : size_(other.size_), capacity_(0) {
     reserve(size_);
     if (data() != nullptr) memcpy(data(), other.data(), size_);
   }
   PaddedBytes& operator=(const PaddedBytes& other) {
     // Self-assignment is safe.
     resize(other.size());
     if (data() != nullptr) memmove(data(), other.data(), size_);
     return *this;
   }
 
   // default is not OK - need to set other.size_ to 0!
   PaddedBytes(PaddedBytes&& other) noexcept
       : size_(other.size_),
         capacity_(other.capacity_),
         data_(std::move(other.data_)) {
     other.size_ = other.capacity_ = 0;
   }
   PaddedBytes& operator=(PaddedBytes&& other) noexcept {
     size_ = other.size_;
     capacity_ = other.capacity_;
     data_ = std::move(other.data_);
 
     if (&other != this) {
       other.size_ = other.capacity_ = 0;
     }
     return *this;
   }
 
   void swap(PaddedBytes& other) noexcept {
     std::swap(size_, other.size_);
     std::swap(capacity_, other.capacity_);
     std::swap(data_, other.data_);
   }
 
   // If current capacity is greater than requested, then no-op. Otherwise
   // copies existing data to newly allocated "data_". If allocation fails,
   // data() == nullptr and size_ = capacity_ = 0.
   // The new capacity will be at least 1.5 times the old capacity. This ensures
   // that we avoid quadratic behaviour.
   void reserve(size_t capacity) {
     if (capacity <= capacity_) return;
 
     size_t new_capacity = std::max(capacity, 3 * capacity_ / 2);
     new_capacity = std::max<size_t>(64, new_capacity);
 
     // BitWriter writes up to 7 bytes past the end.
     CacheAlignedUniquePtr new_data = AllocateArray(new_capacity + 8);
     if (new_data == nullptr) {
       // Allocation failed, discard all data to ensure this is noticed.
       size_ = capacity_ = 0;
       return;
     }
 
     if (data_ == nullptr) {
       // First allocation: ensure first byte is initialized (won't be copied).
       new_data[0] = 0;
     } else {
       // Subsequent resize: copy existing data to new location.
       memcpy(new_data.get(), data_.get(), size_);
       // Ensure that the first new byte is initialized, to allow write_bits to
       // safely append to the newly-resized PaddedBytes.
       new_data[size_] = 0;
     }
 
     capacity_ = new_capacity;
     std::swap(new_data, data_);
   }
 
   // NOTE: unlike vector, this does not initialize the new data!
   // However, we guarantee that write_bits can safely append after
   // the resize, as we zero-initialize the first new byte of data.
   // If size < capacity(), does not invalidate the memory.
   void resize(size_t size) {
     reserve(size);
     size_ = (data() == nullptr) ? 0 : size;
   }
 
   // resize(size) plus explicit initialization of the new data with `value`.
   void resize(size_t size, uint8_t value) {
     size_t old_size = size_;
     resize(size);
     if (size_ > old_size) {
       memset(data() + old_size, value, size_ - old_size);
     }
   }
 
   // Amortized constant complexity due to exponential growth.
   void push_back(uint8_t x) {
     if (size_ == capacity_) {
       reserve(capacity_ + 1);
       if (data() == nullptr) return;
     }
 
     data_[size_++] = x;
   }
 
   size_t size() const { return size_; }
   size_t capacity() const { return capacity_; }
 
   uint8_t* data() { return data_.get(); }
   const uint8_t* data() const { return data_.get(); }
 
   // std::vector operations implemented in terms of the public interface above.
 
   void clear() { resize(0); }
   bool empty() const { return size() == 0; }
 
   void assign(std::initializer_list<uint8_t> il) {
     resize(il.size());
     memcpy(data(), il.begin(), il.size());
   }
 
   uint8_t* begin() { return data(); }
   const uint8_t* begin() const { return data(); }
   uint8_t* end() { return begin() + size(); }
   const uint8_t* end() const { return begin() + size(); }
 
   uint8_t& operator[](const size_t i) {
     BoundsCheck(i);
     return data()[i];
   }
   const uint8_t& operator[](const size_t i) const {
     BoundsCheck(i);
     return data()[i];
   }
 
   uint8_t& back() {
     JXL_ASSERT(size() != 0);
     return data()[size() - 1];
   }
   const uint8_t& back() const {
     JXL_ASSERT(size() != 0);
     return data()[size() - 1];
   }
 
   template <typename T>
   void append(const T& other) {
     append(reinterpret_cast<const uint8_t*>(other.data()),
            reinterpret_cast<const uint8_t*>(other.data()) + other.size());
   }
 
   void append(const uint8_t* begin, const uint8_t* end) {
     if (end - begin > 0) {
       size_t old_size = size();
       resize(size() + (end - begin));
       memcpy(data() + old_size, begin, end - begin);
     }
   }
 
  private:
   void BoundsCheck(size_t i) const {
     // <= is safe due to padding and required by BitWriter.
     JXL_ASSERT(i <= size());
   }
 
   size_t size_;
   size_t capacity_;
   CacheAlignedUniquePtr data_;
 };
 
 template <typename T>
 static inline void Append(const T& s, PaddedBytes* out,
                           size_t* JXL_RESTRICT byte_pos) {
   memcpy(out->data() + *byte_pos, s.data(), s.size());
   *byte_pos += s.size();
   JXL_CHECK(*byte_pos <= out->size());
 }
 
 }  // namespace jxl
 
 #endif  // LIB_JXL_BASE_PADDED_BYTES_H_