enc_bit_writer.cc (8077B)
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 #include "lib/jxl/enc_bit_writer.h" 7 8 #include <jxl/types.h> 9 #include <string.h> // memcpy 10 11 #include "lib/jxl/base/byte_order.h" 12 #include "lib/jxl/base/printf_macros.h" 13 #include "lib/jxl/dec_bit_reader.h" 14 #include "lib/jxl/enc_aux_out.h" 15 16 namespace jxl { 17 18 BitWriter::Allotment::Allotment(BitWriter* JXL_RESTRICT writer, size_t max_bits) 19 : max_bits_(max_bits) { 20 if (writer == nullptr) return; 21 prev_bits_written_ = writer->BitsWritten(); 22 const size_t prev_bytes = writer->storage_.size(); 23 const size_t next_bytes = DivCeil(max_bits, kBitsPerByte); 24 writer->storage_.resize(prev_bytes + next_bytes); 25 parent_ = writer->current_allotment_; 26 writer->current_allotment_ = this; 27 } 28 29 BitWriter::Allotment::~Allotment() { 30 if (!called_) { 31 // Not calling is a bug - unused storage will not be reclaimed. 32 JXL_UNREACHABLE("Did not call Allotment::ReclaimUnused"); 33 } 34 } 35 36 void BitWriter::Allotment::FinishedHistogram(BitWriter* JXL_RESTRICT writer) { 37 if (writer == nullptr) return; 38 JXL_ASSERT(!called_); // Call before ReclaimUnused 39 JXL_ASSERT(histogram_bits_ == 0); // Do not call twice 40 JXL_ASSERT(writer->BitsWritten() >= prev_bits_written_); 41 histogram_bits_ = writer->BitsWritten() - prev_bits_written_; 42 } 43 44 void BitWriter::Allotment::ReclaimAndCharge(BitWriter* JXL_RESTRICT writer, 45 size_t layer, 46 AuxOut* JXL_RESTRICT aux_out) { 47 size_t used_bits = 0; 48 size_t unused_bits = 0; 49 PrivateReclaim(writer, &used_bits, &unused_bits); 50 51 #if JXL_FALSE 52 printf("Layer %s bits: max %" PRIuS " used %" PRIuS " unused %" PRIuS "\n", 53 LayerName(layer), MaxBits(), used_bits, unused_bits); 54 #endif 55 56 // This may be a nested call with aux_out == null. Whenever we know that 57 // aux_out is null, we can call ReclaimUnused directly. 58 if (aux_out != nullptr) { 59 aux_out->layers[layer].total_bits += used_bits; 60 aux_out->layers[layer].histogram_bits += HistogramBits(); 61 } 62 } 63 64 void BitWriter::Allotment::PrivateReclaim(BitWriter* JXL_RESTRICT writer, 65 size_t* JXL_RESTRICT used_bits, 66 size_t* JXL_RESTRICT unused_bits) { 67 JXL_ASSERT(!called_); // Do not call twice 68 called_ = true; 69 if (writer == nullptr) return; 70 71 JXL_ASSERT(writer->BitsWritten() >= prev_bits_written_); 72 *used_bits = writer->BitsWritten() - prev_bits_written_; 73 JXL_ASSERT(*used_bits <= max_bits_); 74 *unused_bits = max_bits_ - *used_bits; 75 76 // Reclaim unused bytes whole bytes from writer's allotment. 77 const size_t unused_bytes = *unused_bits / kBitsPerByte; // truncate 78 JXL_ASSERT(writer->storage_.size() >= unused_bytes); 79 writer->storage_.resize(writer->storage_.size() - unused_bytes); 80 writer->current_allotment_ = parent_; 81 // Ensure we don't also charge the parent for these bits. 82 auto* parent = parent_; 83 while (parent != nullptr) { 84 parent->prev_bits_written_ += *used_bits; 85 parent = parent->parent_; 86 } 87 } 88 89 void BitWriter::AppendByteAligned(const Span<const uint8_t>& span) { 90 if (span.empty()) return; 91 storage_.resize(storage_.size() + span.size() + 1); // extra zero padding 92 93 // Concatenate by copying bytes because both source and destination are bytes. 94 JXL_ASSERT(BitsWritten() % kBitsPerByte == 0); 95 size_t pos = BitsWritten() / kBitsPerByte; 96 memcpy(storage_.data() + pos, span.data(), span.size()); 97 pos += span.size(); 98 storage_[pos++] = 0; // for next Write 99 JXL_ASSERT(pos <= storage_.size()); 100 bits_written_ += span.size() * kBitsPerByte; 101 } 102 103 void BitWriter::AppendByteAligned(const BitWriter& other) { 104 JXL_ASSERT(other.BitsWritten() % kBitsPerByte == 0); 105 JXL_ASSERT(other.BitsWritten() / kBitsPerByte != 0); 106 107 AppendByteAligned(other.GetSpan()); 108 } 109 110 void BitWriter::AppendUnaligned(const BitWriter& other) { 111 Allotment allotment(this, other.BitsWritten()); 112 size_t full_bytes = other.BitsWritten() / kBitsPerByte; 113 size_t remaining_bits = other.BitsWritten() % kBitsPerByte; 114 for (size_t i = 0; i < full_bytes; ++i) { 115 Write(8, other.storage_[i]); 116 } 117 if (remaining_bits > 0) { 118 Write(remaining_bits, 119 other.storage_[full_bytes] & ((1u << remaining_bits) - 1)); 120 } 121 allotment.ReclaimAndCharge(this, 0, nullptr); 122 } 123 124 void BitWriter::AppendByteAligned(const std::vector<BitWriter>& others) { 125 // Total size to add so we can preallocate 126 size_t other_bytes = 0; 127 for (const BitWriter& writer : others) { 128 JXL_ASSERT(writer.BitsWritten() % kBitsPerByte == 0); 129 other_bytes += writer.BitsWritten() / kBitsPerByte; 130 } 131 if (other_bytes == 0) { 132 // No bytes to append: this happens for example when creating per-group 133 // storage for groups, but not writing anything in them for e.g. lossless 134 // images with no alpha. Do nothing. 135 return; 136 } 137 storage_.resize(storage_.size() + other_bytes + 1); // extra zero padding 138 139 // Concatenate by copying bytes because both source and destination are bytes. 140 JXL_ASSERT(BitsWritten() % kBitsPerByte == 0); 141 size_t pos = BitsWritten() / kBitsPerByte; 142 for (const BitWriter& writer : others) { 143 const Span<const uint8_t> span = writer.GetSpan(); 144 if (!span.empty()) { 145 memcpy(storage_.data() + pos, span.data(), span.size()); 146 pos += span.size(); 147 } 148 } 149 storage_[pos++] = 0; // for next Write 150 JXL_ASSERT(pos <= storage_.size()); 151 bits_written_ += other_bytes * kBitsPerByte; 152 } 153 154 // TODO(lode): avoid code duplication 155 void BitWriter::AppendByteAligned( 156 const std::vector<std::unique_ptr<BitWriter>>& others) { 157 // Total size to add so we can preallocate 158 size_t other_bytes = 0; 159 for (const auto& writer : others) { 160 JXL_ASSERT(writer->BitsWritten() % kBitsPerByte == 0); 161 other_bytes += writer->BitsWritten() / kBitsPerByte; 162 } 163 if (other_bytes == 0) { 164 // No bytes to append: this happens for example when creating per-group 165 // storage for groups, but not writing anything in them for e.g. lossless 166 // images with no alpha. Do nothing. 167 return; 168 } 169 storage_.resize(storage_.size() + other_bytes + 1); // extra zero padding 170 171 // Concatenate by copying bytes because both source and destination are bytes. 172 JXL_ASSERT(BitsWritten() % kBitsPerByte == 0); 173 size_t pos = BitsWritten() / kBitsPerByte; 174 for (const auto& writer : others) { 175 const Span<const uint8_t> span = writer->GetSpan(); 176 memcpy(storage_.data() + pos, span.data(), span.size()); 177 pos += span.size(); 178 } 179 storage_[pos++] = 0; // for next Write 180 JXL_ASSERT(pos <= storage_.size()); 181 bits_written_ += other_bytes * kBitsPerByte; 182 } 183 184 // Example: let's assume that 3 bits (Rs below) have been written already: 185 // BYTE+0 BYTE+1 BYTE+2 186 // 0000 0RRR ???? ???? ???? ???? 187 // 188 // Now, we could write up to 5 bits by just shifting them left by 3 bits and 189 // OR'ing to BYTE-0. 190 // 191 // For n > 5 bits, we write the lowest 5 bits as above, then write the next 192 // lowest bits into BYTE+1 starting from its lower bits and so on. 193 void BitWriter::Write(size_t n_bits, uint64_t bits) { 194 JXL_DASSERT((bits >> n_bits) == 0); 195 JXL_DASSERT(n_bits <= kMaxBitsPerCall); 196 uint8_t* p = &storage_[bits_written_ / kBitsPerByte]; 197 const size_t bits_in_first_byte = bits_written_ % kBitsPerByte; 198 bits <<= bits_in_first_byte; 199 #if JXL_BYTE_ORDER_LITTLE 200 uint64_t v = *p; 201 // Last (partial) or next byte to write must be zero-initialized! 202 // PaddedBytes initializes the first, and Write/Append maintain this. 203 JXL_DASSERT(v >> bits_in_first_byte == 0); 204 v |= bits; 205 memcpy(p, &v, sizeof(v)); // Write bytes: possibly more than n_bits/8 206 #else 207 *p++ |= static_cast<uint8_t>(bits & 0xFF); 208 for (size_t bits_left_to_write = n_bits + bits_in_first_byte; 209 bits_left_to_write >= 9; bits_left_to_write -= 8) { 210 bits >>= 8; 211 *p++ = static_cast<uint8_t>(bits & 0xFF); 212 } 213 *p = 0; 214 #endif 215 bits_written_ += n_bits; 216 } 217 } // namespace jxl