enc_bit_writer.h (4308B)
1 // Copyright (c) the JPEG XL Project Authors. All rights reserved. 2 // 3 // Use of this source code is governed by a BSD-style 4 // license that can be found in the LICENSE file. 5 6 #ifndef LIB_JXL_ENC_BIT_WRITER_H_ 7 #define LIB_JXL_ENC_BIT_WRITER_H_ 8 9 // BitWriter class: unbuffered writes using unaligned 64-bit stores. 10 11 #include <stddef.h> 12 #include <stdint.h> 13 14 #include <utility> 15 #include <vector> 16 17 #include "lib/jxl/base/common.h" 18 #include "lib/jxl/base/compiler_specific.h" 19 #include "lib/jxl/base/span.h" 20 #include "lib/jxl/base/status.h" 21 #include "lib/jxl/padded_bytes.h" 22 23 namespace jxl { 24 25 struct AuxOut; 26 27 struct BitWriter { 28 // Upper bound on `n_bits` in each call to Write. We shift a 64-bit word by 29 // 7 bits (max already valid bits in the last byte) and at least 1 bit is 30 // needed to zero-initialize the bit-stream ahead (i.e. if 7 bits are valid 31 // and we write 57 bits, then the next write will access a byte that was not 32 // yet zero-initialized). 33 static constexpr size_t kMaxBitsPerCall = 56; 34 35 BitWriter() : bits_written_(0) {} 36 37 // Disallow copying - may lead to bugs. 38 BitWriter(const BitWriter&) = delete; 39 BitWriter& operator=(const BitWriter&) = delete; 40 BitWriter(BitWriter&&) = default; 41 BitWriter& operator=(BitWriter&&) = default; 42 43 size_t BitsWritten() const { return bits_written_; } 44 45 Span<const uint8_t> GetSpan() const { 46 // Callers must ensure byte alignment to avoid uninitialized bits. 47 JXL_ASSERT(bits_written_ % kBitsPerByte == 0); 48 return Bytes(storage_.data(), bits_written_ / kBitsPerByte); 49 } 50 51 // Example usage: bytes = std::move(writer).TakeBytes(); Useful for the 52 // top-level encoder which returns PaddedBytes, not a BitWriter. 53 // *this must be an rvalue reference and is invalid afterwards. 54 PaddedBytes&& TakeBytes() && { 55 // Callers must ensure byte alignment to avoid uninitialized bits. 56 JXL_ASSERT(bits_written_ % kBitsPerByte == 0); 57 storage_.resize(bits_written_ / kBitsPerByte); 58 return std::move(storage_); 59 } 60 61 // Must be byte-aligned before calling. 62 void AppendByteAligned(const Span<const uint8_t>& span); 63 64 // NOTE: no allotment needed, the other BitWriters have already been charged. 65 void AppendByteAligned(const BitWriter& other); 66 void AppendByteAligned(const std::vector<std::unique_ptr<BitWriter>>& others); 67 void AppendByteAligned(const std::vector<BitWriter>& others); 68 69 void AppendUnaligned(const BitWriter& other); 70 71 class Allotment { 72 public: 73 // Expands a BitWriter's storage. Must happen before calling Write or 74 // ZeroPadToByte. Must call ReclaimUnused after writing to reclaim the 75 // unused storage so that BitWriter memory use remains tightly bounded. 76 Allotment(BitWriter* JXL_RESTRICT writer, size_t max_bits); 77 ~Allotment(); 78 79 size_t MaxBits() const { return max_bits_; } 80 81 // Call after writing a histogram, but before ReclaimUnused. 82 void FinishedHistogram(BitWriter* JXL_RESTRICT writer); 83 84 size_t HistogramBits() const { 85 JXL_ASSERT(called_); 86 return histogram_bits_; 87 } 88 89 void ReclaimAndCharge(BitWriter* JXL_RESTRICT writer, size_t layer, 90 AuxOut* JXL_RESTRICT aux_out); 91 92 private: 93 void PrivateReclaim(BitWriter* JXL_RESTRICT writer, 94 size_t* JXL_RESTRICT used_bits, 95 size_t* JXL_RESTRICT unused_bits); 96 97 size_t prev_bits_written_; 98 const size_t max_bits_; 99 size_t histogram_bits_ = 0; 100 bool called_ = false; 101 Allotment* parent_; 102 }; 103 104 // Writes bits into bytes in increasing addresses, and within a byte 105 // least-significant-bit first. 106 // 107 // The function can write up to 56 bits in one go. 108 void Write(size_t n_bits, uint64_t bits); 109 110 // This should only rarely be used - e.g. when the current location will be 111 // referenced via byte offset (TOCs point to groups), or byte-aligned reading 112 // is required for speed. 113 void ZeroPadToByte() { 114 const size_t remainder_bits = 115 RoundUpBitsToByteMultiple(bits_written_) - bits_written_; 116 if (remainder_bits == 0) return; 117 Write(remainder_bits, 0); 118 JXL_ASSERT(bits_written_ % kBitsPerByte == 0); 119 } 120 121 private: 122 size_t bits_written_; 123 PaddedBytes storage_; 124 Allotment* current_allotment_ = nullptr; 125 }; 126 127 } // namespace jxl 128 129 #endif // LIB_JXL_ENC_BIT_WRITER_H_