Source code
Revision control
Copy as Markdown
Other Tools
// 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.
#include "lib/jxl/enc_bit_writer.h"
#include <jxl/types.h>
#include <cstring> // memcpy
#include "lib/jxl/base/byte_order.h"
#include "lib/jxl/base/common.h"
#include "lib/jxl/base/status.h"
#include "lib/jxl/enc_aux_out.h"
namespace jxl {
Status BitWriter::WithMaxBits(size_t max_bits, LayerType layer, AuxOut* aux_out,
const std::function<Status()>& function,
bool finished_histogram) {
BitWriter::Allotment allotment(max_bits);
JXL_RETURN_IF_ERROR(allotment.Init(this));
const Status result = function();
if (result && finished_histogram) {
JXL_RETURN_IF_ERROR(allotment.FinishedHistogram(this));
}
JXL_RETURN_IF_ERROR(allotment.ReclaimAndCharge(this, layer, aux_out));
return result;
}
BitWriter::Allotment::Allotment(size_t max_bits) : max_bits_(max_bits) {}
Status BitWriter::Allotment::Init(BitWriter* JXL_RESTRICT writer) {
prev_bits_written_ = writer->BitsWritten();
const size_t prev_bytes = writer->storage_.size();
const size_t next_bytes = DivCeil(max_bits_, kBitsPerByte);
JXL_RETURN_IF_ERROR(writer->storage_.resize(prev_bytes + next_bytes));
parent_ = writer->current_allotment_;
writer->current_allotment_ = this;
return true;
}
BitWriter::Allotment::~Allotment() {
if (!called_) {
// Not calling is a bug - unused storage will not be reclaimed.
JXL_DEBUG_ABORT("Did not call Allotment::ReclaimUnused");
}
}
Status BitWriter::Allotment::FinishedHistogram(BitWriter* JXL_RESTRICT writer) {
if (writer == nullptr) return true;
JXL_ENSURE(!called_); // Call before ReclaimUnused
JXL_ENSURE(histogram_bits_ == 0); // Do not call twice
JXL_ENSURE(writer->BitsWritten() >= prev_bits_written_);
if (writer->BitsWritten() >= prev_bits_written_) {
histogram_bits_ = writer->BitsWritten() - prev_bits_written_;
}
return true;
}
Status BitWriter::Allotment::ReclaimAndCharge(BitWriter* JXL_RESTRICT writer,
LayerType layer,
AuxOut* JXL_RESTRICT aux_out) {
size_t used_bits = 0;
size_t unused_bits = 0;
JXL_RETURN_IF_ERROR(PrivateReclaim(writer, &used_bits, &unused_bits));
// This may be a nested call with aux_out == null. Whenever we know that
// aux_out is null, we can call ReclaimUnused directly.
if (aux_out != nullptr) {
aux_out->layer(layer).total_bits += used_bits;
aux_out->layer(layer).histogram_bits += HistogramBits();
}
return true;
}
Status BitWriter::Allotment::PrivateReclaim(BitWriter* JXL_RESTRICT writer,
size_t* JXL_RESTRICT used_bits,
size_t* JXL_RESTRICT unused_bits) {
JXL_DASSERT(!called_); // Do not call twice
called_ = true;
if (writer == nullptr) return true;
JXL_DASSERT(writer->BitsWritten() >= prev_bits_written_);
*used_bits = writer->BitsWritten() - prev_bits_written_;
JXL_DASSERT(*used_bits <= max_bits_);
*unused_bits = max_bits_ - *used_bits;
// Reclaim unused bytes whole bytes from writer's allotment.
const size_t unused_bytes = *unused_bits / kBitsPerByte; // truncate
JXL_ENSURE(writer->storage_.size() >= unused_bytes);
JXL_RETURN_IF_ERROR(
writer->storage_.resize(writer->storage_.size() - unused_bytes));
writer->current_allotment_ = parent_;
// Ensure we don't also charge the parent for these bits.
auto* parent = parent_;
while (parent != nullptr) {
parent->prev_bits_written_ += *used_bits;
parent = parent->parent_;
}
return true;
}
Status BitWriter::AppendByteAligned(const Span<const uint8_t>& span) {
if (span.empty()) return true;
JXL_RETURN_IF_ERROR(storage_.resize(storage_.size() + span.size() +
1)); // extra zero padding
// Concatenate by copying bytes because both source and destination are bytes.
JXL_ENSURE(BitsWritten() % kBitsPerByte == 0);
size_t pos = BitsWritten() / kBitsPerByte;
memcpy(storage_.data() + pos, span.data(), span.size());
pos += span.size();
JXL_ENSURE(pos < storage_.size());
storage_[pos++] = 0; // for next Write
bits_written_ += span.size() * kBitsPerByte;
return true;
}
Status BitWriter::AppendUnaligned(const BitWriter& other) {
return WithMaxBits(other.BitsWritten(), LayerType::Header, nullptr, [&] {
size_t full_bytes = other.BitsWritten() / kBitsPerByte;
size_t remaining_bits = other.BitsWritten() % kBitsPerByte;
for (size_t i = 0; i < full_bytes; ++i) {
Write(8, other.storage_[i]);
}
if (remaining_bits > 0) {
Write(remaining_bits,
other.storage_[full_bytes] & ((1u << remaining_bits) - 1));
}
return true;
});
}
// TODO(lode): avoid code duplication
Status BitWriter::AppendByteAligned(
const std::vector<std::unique_ptr<BitWriter>>& others) {
// Total size to add so we can preallocate
size_t other_bytes = 0;
for (const auto& writer : others) {
JXL_ENSURE(writer->BitsWritten() % kBitsPerByte == 0);
other_bytes += DivCeil(writer->BitsWritten(), kBitsPerByte);
}
if (other_bytes == 0) {
// No bytes to append: this happens for example when creating per-group
// storage for groups, but not writing anything in them for e.g. lossless
// images with no alpha. Do nothing.
return true;
}
JXL_RETURN_IF_ERROR(storage_.resize(storage_.size() + other_bytes +
1)); // extra zero padding
// Concatenate by copying bytes because both source and destination are bytes.
JXL_ENSURE(BitsWritten() % kBitsPerByte == 0);
size_t pos = DivCeil(BitsWritten(), kBitsPerByte);
for (const auto& writer : others) {
const Span<const uint8_t> span = writer->GetSpan();
memcpy(storage_.data() + pos, span.data(), span.size());
pos += span.size();
}
JXL_ENSURE(pos < storage_.size());
storage_[pos++] = 0; // for next Write
bits_written_ += other_bytes * kBitsPerByte;
return true;
}
// Example: let's assume that 3 bits (Rs below) have been written already:
// BYTE+0 BYTE+1 BYTE+2
// 0000 0RRR ???? ???? ???? ????
//
// Now, we could write up to 5 bits by just shifting them left by 3 bits and
// OR'ing to BYTE-0.
//
// For n > 5 bits, we write the lowest 5 bits as above, then write the next
// lowest bits into BYTE+1 starting from its lower bits and so on.
void BitWriter::Write(size_t n_bits, uint64_t bits) {
JXL_DASSERT((bits >> n_bits) == 0);
JXL_DASSERT(n_bits <= kMaxBitsPerCall);
size_t bytes_written = bits_written_ / kBitsPerByte;
uint8_t* p = &storage_[bytes_written];
const size_t bits_in_first_byte = bits_written_ % kBitsPerByte;
bits <<= bits_in_first_byte;
#if JXL_BYTE_ORDER_LITTLE
uint64_t v = *p;
// Last (partial) or next byte to write must be zero-initialized!
// PaddedBytes initializes the first, and Write/Append maintain this.
JXL_DASSERT(v >> bits_in_first_byte == 0);
v |= bits;
memcpy(p, &v, sizeof(v)); // Write bytes: possibly more than n_bits/8
#else
*p++ |= static_cast<uint8_t>(bits & 0xFF);
for (size_t bits_left_to_write = n_bits + bits_in_first_byte;
bits_left_to_write >= 9; bits_left_to_write -= 8) {
bits >>= 8;
*p++ = static_cast<uint8_t>(bits & 0xFF);
}
*p = 0;
#endif
bits_written_ += n_bits;
}
} // namespace jxl