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.
#ifndef LIB_JXL_DEC_ANS_H_
#define LIB_JXL_DEC_ANS_H_
// Library to decode the ANS population counts from the bit-stream and build a
// decoding table from them.
#include <jxl/memory_manager.h>
#include <jxl/types.h>
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <vector>
#include "lib/jxl/ans_common.h"
#include "lib/jxl/ans_params.h"
#include "lib/jxl/base/bits.h"
#include "lib/jxl/base/compiler_specific.h"
#include "lib/jxl/base/status.h"
#include "lib/jxl/dec_bit_reader.h"
#include "lib/jxl/dec_huffman.h"
#include "lib/jxl/field_encodings.h"
#include "lib/jxl/memory_manager_internal.h"
namespace jxl {
class ANSSymbolReader;
// Experiments show that best performance is typically achieved for a
// split-exponent of 3 or 4. Trend seems to be that '4' is better
// for large-ish pictures, and '3' better for rather small-ish pictures.
// This is plausible - the more special symbols we have, the better
// statistics we need to get a benefit out of them.
// Our hybrid-encoding scheme has dedicated tokens for the smallest
// (1 << split_exponents) numbers, and for the rest
// encodes (number of bits) + (msb_in_token sub-leading binary digits) +
// (lsb_in_token lowest binary digits) in the token, with the remaining bits
// then being encoded as data.
//
// Example with split_exponent = 4, msb_in_token = 2, lsb_in_token = 0.
//
// Numbers N in [0 .. 15]:
// These get represented as (token=N, bits='').
// Numbers N >= 16:
// If n is such that 2**n <= N < 2**(n+1),
// and m = N - 2**n is the 'mantissa',
// these get represented as:
// (token=split_token +
// ((n - split_exponent) * 4) +
// (m >> (n - msb_in_token)),
// bits=m & (1 << (n - msb_in_token)) - 1)
// Specifically, we would get:
// N = 0 - 15: (token=N, nbits=0, bits='')
// N = 16 (10000): (token=16, nbits=2, bits='00')
// N = 17 (10001): (token=16, nbits=2, bits='01')
// N = 20 (10100): (token=17, nbits=2, bits='00')
// N = 24 (11000): (token=18, nbits=2, bits='00')
// N = 28 (11100): (token=19, nbits=2, bits='00')
// N = 32 (100000): (token=20, nbits=3, bits='000')
// N = 65535: (token=63, nbits=13, bits='1111111111111')
struct HybridUintConfig {
uint32_t split_exponent;
uint32_t split_token;
uint32_t msb_in_token;
uint32_t lsb_in_token;
JXL_INLINE void Encode(uint32_t value, uint32_t* JXL_RESTRICT token,
uint32_t* JXL_RESTRICT nbits,
uint32_t* JXL_RESTRICT bits) const {
if (value < split_token) {
*token = value;
*nbits = 0;
*bits = 0;
} else {
uint32_t n = FloorLog2Nonzero(value);
uint32_t m = value - (1 << n);
*token = split_token +
((n - split_exponent) << (msb_in_token + lsb_in_token)) +
((m >> (n - msb_in_token)) << lsb_in_token) +
(m & ((1 << lsb_in_token) - 1));
*nbits = n - msb_in_token - lsb_in_token;
*bits = (value >> lsb_in_token) & ((1UL << *nbits) - 1);
}
}
explicit HybridUintConfig(uint32_t split_exponent = 4,
uint32_t msb_in_token = 2,
uint32_t lsb_in_token = 0)
: split_exponent(split_exponent),
split_token(1 << split_exponent),
msb_in_token(msb_in_token),
lsb_in_token(lsb_in_token) {
JXL_DASSERT(split_exponent >= msb_in_token + lsb_in_token);
}
};
struct LZ77Params : public Fields {
LZ77Params();
JXL_FIELDS_NAME(LZ77Params)
Status VisitFields(Visitor* JXL_RESTRICT visitor) override;
bool enabled;
// Symbols above min_symbol use a special hybrid uint encoding and
// represent a length, to be added to min_length.
uint32_t min_symbol;
uint32_t min_length;
// Not serialized by VisitFields.
HybridUintConfig length_uint_config{0, 0, 0};
size_t nonserialized_distance_context;
};
static constexpr size_t kWindowSize = 1 << 20;
static constexpr size_t kNumSpecialDistances = 120;
// Table of special distance codes from WebP lossless.
static constexpr int8_t kSpecialDistances[kNumSpecialDistances][2] = {
{0, 1}, {1, 0}, {1, 1}, {-1, 1}, {0, 2}, {2, 0}, {1, 2}, {-1, 2},
{2, 1}, {-2, 1}, {2, 2}, {-2, 2}, {0, 3}, {3, 0}, {1, 3}, {-1, 3},
{3, 1}, {-3, 1}, {2, 3}, {-2, 3}, {3, 2}, {-3, 2}, {0, 4}, {4, 0},
{1, 4}, {-1, 4}, {4, 1}, {-4, 1}, {3, 3}, {-3, 3}, {2, 4}, {-2, 4},
{4, 2}, {-4, 2}, {0, 5}, {3, 4}, {-3, 4}, {4, 3}, {-4, 3}, {5, 0},
{1, 5}, {-1, 5}, {5, 1}, {-5, 1}, {2, 5}, {-2, 5}, {5, 2}, {-5, 2},
{4, 4}, {-4, 4}, {3, 5}, {-3, 5}, {5, 3}, {-5, 3}, {0, 6}, {6, 0},
{1, 6}, {-1, 6}, {6, 1}, {-6, 1}, {2, 6}, {-2, 6}, {6, 2}, {-6, 2},
{4, 5}, {-4, 5}, {5, 4}, {-5, 4}, {3, 6}, {-3, 6}, {6, 3}, {-6, 3},
{0, 7}, {7, 0}, {1, 7}, {-1, 7}, {5, 5}, {-5, 5}, {7, 1}, {-7, 1},
{4, 6}, {-4, 6}, {6, 4}, {-6, 4}, {2, 7}, {-2, 7}, {7, 2}, {-7, 2},
{3, 7}, {-3, 7}, {7, 3}, {-7, 3}, {5, 6}, {-5, 6}, {6, 5}, {-6, 5},
{8, 0}, {4, 7}, {-4, 7}, {7, 4}, {-7, 4}, {8, 1}, {8, 2}, {6, 6},
{-6, 6}, {8, 3}, {5, 7}, {-5, 7}, {7, 5}, {-7, 5}, {8, 4}, {6, 7},
{-6, 7}, {7, 6}, {-7, 6}, {8, 5}, {7, 7}, {-7, 7}, {8, 6}, {8, 7}};
static JXL_INLINE int SpecialDistance(size_t index, int multiplier) {
int dist = kSpecialDistances[index][0] +
static_cast<int>(multiplier) * kSpecialDistances[index][1];
return (dist > 1) ? dist : 1;
}
struct ANSCode {
AlignedMemory alias_tables;
std::vector<HuffmanDecodingData> huffman_data;
std::vector<HybridUintConfig> uint_config;
std::vector<int> degenerate_symbols;
bool use_prefix_code;
uint8_t log_alpha_size; // for ANS.
LZ77Params lz77;
// Maximum number of bits necessary to represent the result of a
// ReadHybridUint call done with this ANSCode.
size_t max_num_bits = 0;
JxlMemoryManager* memory_manager;
void UpdateMaxNumBits(size_t ctx, size_t symbol);
};
class ANSSymbolReader {
public:
// Invalid symbol reader, to be overwritten.
ANSSymbolReader() = default;
static StatusOr<ANSSymbolReader> Create(const ANSCode* code,
BitReader* JXL_RESTRICT br,
size_t distance_multiplier = 0);
JXL_INLINE size_t ReadSymbolANSWithoutRefill(const size_t histo_idx,
BitReader* JXL_RESTRICT br) {
const uint32_t res = state_ & (ANS_TAB_SIZE - 1u);
const AliasTable::Entry* table =
&alias_tables_[histo_idx << log_alpha_size_];
const AliasTable::Symbol symbol =
AliasTable::Lookup(table, res, log_entry_size_, entry_size_minus_1_);
state_ = symbol.freq * (state_ >> ANS_LOG_TAB_SIZE) + symbol.offset;
#if JXL_TRUE
// Branchless version is about equally fast on SKX.
const uint32_t new_state =
(state_ << 16u) | static_cast<uint32_t>(br->PeekFixedBits<16>());
const bool normalize = state_ < (1u << 16u);
state_ = normalize ? new_state : state_;
br->Consume(normalize ? 16 : 0);
#else
if (JXL_UNLIKELY(state_ < (1u << 16u))) {
state_ = (state_ << 16u) | br->PeekFixedBits<16>();
br->Consume(16);
}
#endif
const uint32_t next_res = state_ & (ANS_TAB_SIZE - 1u);
AliasTable::Prefetch(table, next_res, log_entry_size_);
return symbol.value;
}
JXL_INLINE size_t ReadSymbolHuffWithoutRefill(const size_t histo_idx,
BitReader* JXL_RESTRICT br) {
return huffman_data_[histo_idx].ReadSymbol(br);
}
JXL_INLINE size_t ReadSymbolWithoutRefill(const size_t histo_idx,
BitReader* JXL_RESTRICT br) {
// TODO(veluca): hoist if in hotter loops.
if (JXL_UNLIKELY(use_prefix_code_)) {
return ReadSymbolHuffWithoutRefill(histo_idx, br);
}
return ReadSymbolANSWithoutRefill(histo_idx, br);
}
JXL_INLINE size_t ReadSymbol(const size_t histo_idx,
BitReader* JXL_RESTRICT br) {
br->Refill();
return ReadSymbolWithoutRefill(histo_idx, br);
}
#ifdef FUZZING_BUILD_MODE_UNSAFE_FOR_PRODUCTION
bool CheckANSFinalState() const { return true; }
#else
bool CheckANSFinalState() const { return state_ == (ANS_SIGNATURE << 16u); }
#endif
template <typename BitReader>
static JXL_INLINE uint32_t ReadHybridUintConfig(
const HybridUintConfig& config, size_t token, BitReader* br) {
size_t split_token = config.split_token;
size_t msb_in_token = config.msb_in_token;
size_t lsb_in_token = config.lsb_in_token;
size_t split_exponent = config.split_exponent;
// Fast-track version of hybrid integer decoding.
if (token < split_token) return token;
uint32_t nbits = split_exponent - (msb_in_token + lsb_in_token) +
((token - split_token) >> (msb_in_token + lsb_in_token));
// Max amount of bits for ReadBits is 32 and max valid left shift is 29
// bits. However, for speed no error is propagated here, instead limit the
// nbits size. If nbits > 29, the code stream is invalid, but no error is
// returned.
// Note that in most cases we will emit an error if the histogram allows
// representing numbers that would cause invalid shifts, but we need to
// keep this check as when LZ77 is enabled it might make sense to have an
// histogram that could in principle cause invalid shifts.
nbits &= 31u;
uint32_t low = token & ((1 << lsb_in_token) - 1);
token >>= lsb_in_token;
const size_t bits = br->PeekBits(nbits);
br->Consume(nbits);
size_t ret = (((((1 << msb_in_token) | (token & ((1 << msb_in_token) - 1)))
<< nbits) |
bits)
<< lsb_in_token) |
low;
// TODO(eustas): mark BitReader as unhealthy if nbits > 29 or ret does not
// fit uint32_t
return static_cast<uint32_t>(ret);
}
// Takes a *clustered* idx. Can only use if HuffRleOnly() is true.
JXL_INLINE void ReadHybridUintClusteredHuffRleOnly(size_t ctx,
BitReader* JXL_RESTRICT br,
uint32_t* value,
uint32_t* run) {
JXL_DASSERT(IsHuffRleOnly());
br->Refill(); // covers ReadSymbolWithoutRefill + PeekBits
size_t token = ReadSymbolHuffWithoutRefill(ctx, br);
if (JXL_UNLIKELY(token >= lz77_threshold_)) {
*run =
ReadHybridUintConfig(lz77_length_uint_, token - lz77_threshold_, br) +
lz77_min_length_ - 1;
return;
}
*value = ReadHybridUintConfig(configs[ctx], token, br);
}
bool IsHuffRleOnly() const {
if (lz77_window_ == nullptr) return false;
if (!use_prefix_code_) return false;
for (size_t i = 0; i < kHuffmanTableBits; i++) {
if (huffman_data_[lz77_ctx_].table_[i].bits) return false;
if (huffman_data_[lz77_ctx_].table_[i].value != 1) return false;
}
if (configs[lz77_ctx_].split_token > 1) return false;
return true;
}
bool UsesLZ77() { return lz77_window_ != nullptr; }
// Takes a *clustered* idx. Inlined, for use in hot paths.
template <bool uses_lz77>
JXL_INLINE size_t ReadHybridUintClusteredInlined(size_t ctx,
BitReader* JXL_RESTRICT br) {
if (uses_lz77) {
if (JXL_UNLIKELY(num_to_copy_ > 0)) {
size_t ret = lz77_window_[(copy_pos_++) & kWindowMask];
num_to_copy_--;
lz77_window_[(num_decoded_++) & kWindowMask] = ret;
return ret;
}
}
br->Refill(); // covers ReadSymbolWithoutRefill + PeekBits
size_t token = ReadSymbolWithoutRefill(ctx, br);
if (uses_lz77) {
if (JXL_UNLIKELY(token >= lz77_threshold_)) {
num_to_copy_ = ReadHybridUintConfig(lz77_length_uint_,
token - lz77_threshold_, br) +
lz77_min_length_;
br->Refill(); // covers ReadSymbolWithoutRefill + PeekBits
// Distance code.
size_t token = ReadSymbolWithoutRefill(lz77_ctx_, br);
size_t distance = ReadHybridUintConfig(configs[lz77_ctx_], token, br);
if (JXL_LIKELY(distance < num_special_distances_)) {
distance = special_distances_[distance];
} else {
distance = distance + 1 - num_special_distances_;
}
if (JXL_UNLIKELY(distance > num_decoded_)) {
distance = num_decoded_;
}
if (JXL_UNLIKELY(distance > kWindowSize)) {
distance = kWindowSize;
}
copy_pos_ = num_decoded_ - distance;
if (JXL_UNLIKELY(distance == 0)) {
JXL_DASSERT(lz77_window_ != nullptr);
// distance 0 -> num_decoded_ == copy_pos_ == 0
size_t to_fill = std::min<size_t>(num_to_copy_, kWindowSize);
memset(lz77_window_, 0, to_fill * sizeof(lz77_window_[0]));
}
// TODO(eustas): overflow; mark BitReader as unhealthy
if (num_to_copy_ < lz77_min_length_) return 0;
// the code below is the same as doing this:
// return ReadHybridUintClustered<uses_lz77>(ctx, br);
// but gcc doesn't like recursive inlining
size_t ret = lz77_window_[(copy_pos_++) & kWindowMask];
num_to_copy_--;
lz77_window_[(num_decoded_++) & kWindowMask] = ret;
return ret;
}
}
size_t ret = ReadHybridUintConfig(configs[ctx], token, br);
if (uses_lz77 && lz77_window_)
lz77_window_[(num_decoded_++) & kWindowMask] = ret;
return ret;
}
// same but not inlined
template <bool uses_lz77>
size_t ReadHybridUintClustered(size_t ctx, BitReader* JXL_RESTRICT br) {
return ReadHybridUintClusteredInlined<uses_lz77>(ctx, br);
}
// inlined only in the no-lz77 case
template <bool uses_lz77>
JXL_INLINE size_t
ReadHybridUintClusteredMaybeInlined(size_t ctx, BitReader* JXL_RESTRICT br) {
if (uses_lz77) {
return ReadHybridUintClustered<uses_lz77>(ctx, br);
} else {
return ReadHybridUintClusteredInlined<uses_lz77>(ctx, br);
}
}
// inlined, for use in hot paths
template <bool uses_lz77>
JXL_INLINE size_t
ReadHybridUintInlined(size_t ctx, BitReader* JXL_RESTRICT br,
const std::vector<uint8_t>& context_map) {
return ReadHybridUintClustered<uses_lz77>(context_map[ctx], br);
}
// not inlined, for use in non-hot paths
size_t ReadHybridUint(size_t ctx, BitReader* JXL_RESTRICT br,
const std::vector<uint8_t>& context_map) {
return ReadHybridUintClustered</*uses_lz77=*/true>(context_map[ctx], br);
}
// ctx is a *clustered* context!
// This function will modify the ANS state as if `count` symbols have been
// decoded.
bool IsSingleValueAndAdvance(size_t ctx, uint32_t* value, size_t count) {
// TODO(veluca): No optimization for Huffman mode yet.
if (use_prefix_code_) return false;
// TODO(eustas): propagate "degenerate_symbol" to simplify this method.
const uint32_t res = state_ & (ANS_TAB_SIZE - 1u);
const AliasTable::Entry* table = &alias_tables_[ctx << log_alpha_size_];
AliasTable::Symbol symbol =
AliasTable::Lookup(table, res, log_entry_size_, entry_size_minus_1_);
if (symbol.freq != ANS_TAB_SIZE) return false;
if (configs[ctx].split_token <= symbol.value) return false;
if (symbol.value >= lz77_threshold_) return false;
*value = symbol.value;
if (lz77_window_) {
for (size_t i = 0; i < count; i++) {
lz77_window_[(num_decoded_++) & kWindowMask] = symbol.value;
}
}
return true;
}
static constexpr size_t kMaxCheckpointInterval = 512;
struct Checkpoint {
uint32_t state;
uint32_t num_to_copy;
uint32_t copy_pos;
uint32_t num_decoded;
uint32_t lz77_window[kMaxCheckpointInterval];
};
void Save(Checkpoint* checkpoint) {
checkpoint->state = state_;
checkpoint->num_decoded = num_decoded_;
checkpoint->num_to_copy = num_to_copy_;
checkpoint->copy_pos = copy_pos_;
if (lz77_window_) {
size_t win_start = num_decoded_ & kWindowMask;
size_t win_end = (num_decoded_ + kMaxCheckpointInterval) & kWindowMask;
if (win_end > win_start) {
memcpy(checkpoint->lz77_window, lz77_window_ + win_start,
(win_end - win_start) * sizeof(*lz77_window_));
} else {
memcpy(checkpoint->lz77_window, lz77_window_ + win_start,
(kWindowSize - win_start) * sizeof(*lz77_window_));
memcpy(checkpoint->lz77_window + (kWindowSize - win_start),
lz77_window_, win_end * sizeof(*lz77_window_));
}
}
}
void Restore(const Checkpoint& checkpoint) {
state_ = checkpoint.state;
JXL_DASSERT(num_decoded_ <=
checkpoint.num_decoded + kMaxCheckpointInterval);
num_decoded_ = checkpoint.num_decoded;
num_to_copy_ = checkpoint.num_to_copy;
copy_pos_ = checkpoint.copy_pos;
if (lz77_window_) {
size_t win_start = num_decoded_ & kWindowMask;
size_t win_end = (num_decoded_ + kMaxCheckpointInterval) & kWindowMask;
if (win_end > win_start) {
memcpy(lz77_window_ + win_start, checkpoint.lz77_window,
(win_end - win_start) * sizeof(*lz77_window_));
} else {
memcpy(lz77_window_ + win_start, checkpoint.lz77_window,
(kWindowSize - win_start) * sizeof(*lz77_window_));
memcpy(lz77_window_, checkpoint.lz77_window + (kWindowSize - win_start),
win_end * sizeof(*lz77_window_));
}
}
}
private:
ANSSymbolReader(const ANSCode* code, BitReader* JXL_RESTRICT br,
size_t distance_multiplier,
AlignedMemory&& lz77_window_storage);
const AliasTable::Entry* JXL_RESTRICT alias_tables_; // not owned
const HuffmanDecodingData* huffman_data_;
bool use_prefix_code_;
uint32_t state_ = ANS_SIGNATURE << 16u;
const HybridUintConfig* JXL_RESTRICT configs;
uint32_t log_alpha_size_{};
uint32_t log_entry_size_{};
uint32_t entry_size_minus_1_{};
// LZ77 structures and constants.
static constexpr size_t kWindowMask = kWindowSize - 1;
// a std::vector incurs unacceptable decoding speed loss because of
// initialization.
AlignedMemory lz77_window_storage_;
uint32_t* lz77_window_ = nullptr;
uint32_t num_decoded_ = 0;
uint32_t num_to_copy_ = 0;
uint32_t copy_pos_ = 0;
uint32_t lz77_ctx_ = 0;
uint32_t lz77_min_length_ = 0;
uint32_t lz77_threshold_ = 1 << 20; // bigger than any symbol.
HybridUintConfig lz77_length_uint_;
uint32_t special_distances_[kNumSpecialDistances]{};
uint32_t num_special_distances_{};
};
Status DecodeHistograms(JxlMemoryManager* memory_manager, BitReader* br,
size_t num_contexts, ANSCode* code,
std::vector<uint8_t>* context_map,
bool disallow_lz77 = false);
// Exposed for tests.
Status DecodeUintConfigs(size_t log_alpha_size,
std::vector<HybridUintConfig>* uint_config,
BitReader* br);
} // namespace jxl
#endif // LIB_JXL_DEC_ANS_H_