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