enc_bit_writer.h (5046B)
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 <jxl/memory_manager.h> 12 13 #include <cstddef> 14 #include <cstdint> 15 #include <functional> 16 #include <memory> 17 #include <utility> 18 #include <vector> 19 20 #include "lib/jxl/base/common.h" 21 #include "lib/jxl/base/compiler_specific.h" 22 #include "lib/jxl/base/span.h" 23 #include "lib/jxl/base/status.h" 24 #include "lib/jxl/padded_bytes.h" 25 26 namespace jxl { 27 28 struct AuxOut; 29 enum class LayerType : uint8_t; 30 31 struct BitWriter { 32 // Upper bound on `n_bits` in each call to Write. We shift a 64-bit word by 33 // 7 bits (max already valid bits in the last byte) and at least 1 bit is 34 // needed to zero-initialize the bit-stream ahead (i.e. if 7 bits are valid 35 // and we write 57 bits, then the next write will access a byte that was not 36 // yet zero-initialized). 37 static constexpr size_t kMaxBitsPerCall = 56; 38 39 explicit BitWriter(JxlMemoryManager* memory_manager) 40 : bits_written_(0), storage_(memory_manager) {} 41 42 // Disallow copying - may lead to bugs. 43 BitWriter(const BitWriter&) = delete; 44 BitWriter& operator=(const BitWriter&) = delete; 45 BitWriter(BitWriter&&) = default; 46 BitWriter& operator=(BitWriter&&) = default; 47 48 size_t BitsWritten() const { return bits_written_; } 49 50 JxlMemoryManager* memory_manager() const { return storage_.memory_manager(); } 51 52 Span<const uint8_t> GetSpan() const { 53 // Callers must ensure byte alignment to avoid uninitialized bits. 54 JXL_DASSERT(bits_written_ % kBitsPerByte == 0); 55 return Bytes(storage_.data(), DivCeil(bits_written_, kBitsPerByte)); 56 } 57 58 // Example usage: bytes = std::move(writer).TakeBytes(); Useful for the 59 // top-level encoder which returns PaddedBytes, not a BitWriter. 60 // *this must be an rvalue reference and is invalid afterwards. 61 PaddedBytes&& TakeBytes() && { 62 // Callers must ensure byte alignment to avoid uninitialized bits. 63 JXL_DASSERT(bits_written_ % kBitsPerByte == 0); 64 Status status = storage_.resize(DivCeil(bits_written_, kBitsPerByte)); 65 JXL_DASSERT(status); 66 // Can never fail, because we are resizing to a lower size. 67 (void)status; 68 return std::move(storage_); 69 } 70 71 // Must be byte-aligned before calling. 72 Status AppendByteAligned(const Span<const uint8_t>& span); 73 74 // NOTE: no allotment needed, the other BitWriters have already been charged. 75 Status AppendByteAligned( 76 const std::vector<std::unique_ptr<BitWriter>>& others); 77 78 Status AppendUnaligned(const BitWriter& other); 79 80 // Writes bits into bytes in increasing addresses, and within a byte 81 // least-significant-bit first. 82 // 83 // The function can write up to 56 bits in one go. 84 void Write(size_t n_bits, uint64_t bits); 85 86 // This should only rarely be used - e.g. when the current location will be 87 // referenced via byte offset (TOCs point to groups), or byte-aligned reading 88 // is required for speed. 89 void ZeroPadToByte() { 90 const size_t remainder_bits = 91 RoundUpBitsToByteMultiple(bits_written_) - bits_written_; 92 if (remainder_bits == 0) return; 93 Write(remainder_bits, 0); 94 JXL_DASSERT(bits_written_ % kBitsPerByte == 0); 95 } 96 97 Status WithMaxBits(size_t max_bits, LayerType layer, 98 AuxOut* JXL_RESTRICT aux_out, 99 const std::function<Status()>& function, 100 bool finished_histogram = false); 101 102 private: 103 class Allotment { 104 public: 105 explicit Allotment(size_t max_bits); 106 ~Allotment(); 107 108 Allotment(const Allotment& other) = delete; 109 Allotment(Allotment&& other) = delete; 110 Allotment& operator=(const Allotment&) = delete; 111 Allotment& operator=(Allotment&&) = delete; 112 113 // Call after writing a histogram, but before ReclaimUnused. 114 Status FinishedHistogram(BitWriter* JXL_RESTRICT writer); 115 116 size_t HistogramBits() const { 117 JXL_DASSERT(called_); 118 return histogram_bits_; 119 } 120 121 Status ReclaimAndCharge(BitWriter* JXL_RESTRICT writer, LayerType layer, 122 AuxOut* JXL_RESTRICT aux_out); 123 124 private: 125 friend struct BitWriter; 126 127 // Expands a BitWriter's storage. Must happen before calling Write or 128 // ZeroPadToByte. Must call ReclaimUnused after writing to reclaim the 129 // unused storage so that BitWriter memory use remains tightly bounded. 130 Status Init(BitWriter* JXL_RESTRICT writer); 131 132 Status PrivateReclaim(BitWriter* JXL_RESTRICT writer, 133 size_t* JXL_RESTRICT used_bits, 134 size_t* JXL_RESTRICT unused_bits); 135 136 size_t prev_bits_written_; 137 const size_t max_bits_; 138 size_t histogram_bits_ = 0; 139 bool called_ = false; 140 Allotment* parent_; 141 }; 142 143 size_t bits_written_; 144 PaddedBytes storage_; 145 Allotment* current_allotment_ = nullptr; 146 }; 147 148 } // namespace jxl 149 150 #endif // LIB_JXL_ENC_BIT_WRITER_H_