tor-browser

The Tor Browser
git clone https://git.dasho.dev/tor-browser.git
Log | Files | Refs | README | LICENSE

dec_frame.cc (34281B)


      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/dec_frame.h"
      7 
      8 #include <jxl/decode.h>
      9 #include <jxl/memory_manager.h>
     10 
     11 #include <algorithm>
     12 #include <cstddef>
     13 #include <cstdint>
     14 #include <cstdlib>
     15 #include <memory>
     16 #include <utility>
     17 #include <vector>
     18 
     19 #include "lib/jxl/ac_context.h"
     20 #include "lib/jxl/ac_strategy.h"
     21 #include "lib/jxl/base/bits.h"
     22 #include "lib/jxl/base/common.h"
     23 #include "lib/jxl/base/compiler_specific.h"
     24 #include "lib/jxl/base/data_parallel.h"
     25 #include "lib/jxl/base/printf_macros.h"
     26 #include "lib/jxl/base/rect.h"
     27 #include "lib/jxl/base/status.h"
     28 #include "lib/jxl/chroma_from_luma.h"
     29 #include "lib/jxl/coeff_order.h"
     30 #include "lib/jxl/coeff_order_fwd.h"
     31 #include "lib/jxl/common.h"  // kMaxNumPasses
     32 #include "lib/jxl/compressed_dc.h"
     33 #include "lib/jxl/dct_util.h"
     34 #include "lib/jxl/dec_ans.h"
     35 #include "lib/jxl/dec_bit_reader.h"
     36 #include "lib/jxl/dec_cache.h"
     37 #include "lib/jxl/dec_group.h"
     38 #include "lib/jxl/dec_modular.h"
     39 #include "lib/jxl/dec_noise.h"
     40 #include "lib/jxl/dec_patch_dictionary.h"
     41 #include "lib/jxl/entropy_coder.h"
     42 #include "lib/jxl/epf.h"
     43 #include "lib/jxl/fields.h"
     44 #include "lib/jxl/frame_dimensions.h"
     45 #include "lib/jxl/frame_header.h"
     46 #include "lib/jxl/image_bundle.h"
     47 #include "lib/jxl/image_metadata.h"
     48 #include "lib/jxl/image_ops.h"
     49 #include "lib/jxl/jpeg/jpeg_data.h"
     50 #include "lib/jxl/loop_filter.h"
     51 #include "lib/jxl/passes_state.h"
     52 #include "lib/jxl/quant_weights.h"
     53 #include "lib/jxl/quantizer.h"
     54 #include "lib/jxl/render_pipeline/render_pipeline.h"
     55 #include "lib/jxl/splines.h"
     56 #include "lib/jxl/toc.h"
     57 
     58 namespace jxl {
     59 
     60 namespace {
     61 Status DecodeGlobalDCInfo(BitReader* reader, bool is_jpeg,
     62                          PassesDecoderState* state, ThreadPool* pool) {
     63  JXL_RETURN_IF_ERROR(state->shared_storage.quantizer.Decode(reader));
     64 
     65  JXL_RETURN_IF_ERROR(DecodeBlockCtxMap(state->memory_manager(), reader,
     66                                        &state->shared_storage.block_ctx_map));
     67 
     68  JXL_RETURN_IF_ERROR(state->shared_storage.cmap.DecodeDC(reader));
     69 
     70  // Pre-compute info for decoding a group.
     71  if (is_jpeg) {
     72    state->shared_storage.quantizer.ClearDCMul();  // Don't dequant DC
     73  }
     74 
     75  state->shared_storage.ac_strategy.FillInvalid();
     76  return true;
     77 }
     78 }  // namespace
     79 
     80 Status DecodeFrame(PassesDecoderState* dec_state, ThreadPool* JXL_RESTRICT pool,
     81                   const uint8_t* next_in, size_t avail_in,
     82                   FrameHeader* frame_header, ImageBundle* decoded,
     83                   const CodecMetadata& metadata,
     84                   bool use_slow_rendering_pipeline) {
     85  FrameDecoder frame_decoder(dec_state, metadata, pool,
     86                             use_slow_rendering_pipeline);
     87 
     88  BitReader reader(Bytes(next_in, avail_in));
     89  JXL_RETURN_IF_ERROR(frame_decoder.InitFrame(&reader, decoded,
     90                                              /*is_preview=*/false));
     91  JXL_RETURN_IF_ERROR(frame_decoder.InitFrameOutput());
     92  if (frame_header) {
     93    *frame_header = frame_decoder.GetFrameHeader();
     94  }
     95 
     96  JXL_RETURN_IF_ERROR(reader.AllReadsWithinBounds());
     97  size_t header_bytes = reader.TotalBitsConsumed() / kBitsPerByte;
     98  JXL_RETURN_IF_ERROR(reader.Close());
     99 
    100  size_t processed_bytes = header_bytes;
    101  Status close_ok = true;
    102  std::vector<std::unique_ptr<BitReader>> section_readers;
    103  {
    104    std::vector<std::unique_ptr<BitReaderScopedCloser>> section_closers;
    105    std::vector<FrameDecoder::SectionInfo> section_info;
    106    std::vector<FrameDecoder::SectionStatus> section_status;
    107    size_t pos = header_bytes;
    108    size_t index = 0;
    109    for (auto toc_entry : frame_decoder.Toc()) {
    110      JXL_RETURN_IF_ERROR(pos + toc_entry.size <= avail_in);
    111      auto br = make_unique<BitReader>(Bytes(next_in + pos, toc_entry.size));
    112      section_info.emplace_back(
    113          FrameDecoder::SectionInfo{br.get(), toc_entry.id, index++});
    114      section_closers.emplace_back(
    115          make_unique<BitReaderScopedCloser>(*br, close_ok));
    116      section_readers.emplace_back(std::move(br));
    117      pos += toc_entry.size;
    118    }
    119    section_status.resize(section_info.size());
    120    JXL_RETURN_IF_ERROR(frame_decoder.ProcessSections(
    121        section_info.data(), section_info.size(), section_status.data()));
    122    for (size_t i = 0; i < section_status.size(); i++) {
    123      JXL_RETURN_IF_ERROR(section_status[i] == FrameDecoder::kDone);
    124      processed_bytes += frame_decoder.Toc()[i].size;
    125    }
    126  }
    127  JXL_RETURN_IF_ERROR(close_ok);
    128  JXL_RETURN_IF_ERROR(frame_decoder.FinalizeFrame());
    129  decoded->SetDecodedBytes(processed_bytes);
    130  return true;
    131 }
    132 
    133 Status FrameDecoder::InitFrame(BitReader* JXL_RESTRICT br, ImageBundle* decoded,
    134                               bool is_preview) {
    135  decoded_ = decoded;
    136  JXL_ENSURE(is_finalized_);
    137  JxlMemoryManager* memory_manager = decoded_->memory_manager();
    138 
    139  // Reset the dequantization matrices to their default values.
    140  dec_state_->shared_storage.matrices = DequantMatrices();
    141 
    142  frame_header_.nonserialized_is_preview = is_preview;
    143  JXL_ENSURE(frame_header_.nonserialized_metadata != nullptr);
    144  JXL_RETURN_IF_ERROR(ReadFrameHeader(br, &frame_header_));
    145  frame_dim_ = frame_header_.ToFrameDimensions();
    146  JXL_DEBUG_V(2, "FrameHeader: %s", frame_header_.DebugString().c_str());
    147 
    148  const size_t num_passes = frame_header_.passes.num_passes;
    149  const size_t num_groups = frame_dim_.num_groups;
    150 
    151  // If the previous frame was not a kRegularFrame, `decoded` may have different
    152  // dimensions; must reset to avoid errors.
    153  decoded->RemoveColor();
    154  decoded->ClearExtraChannels();
    155 
    156  decoded->duration = frame_header_.animation_frame.duration;
    157 
    158  if (!frame_header_.nonserialized_is_preview &&
    159      (frame_header_.is_last || frame_header_.animation_frame.duration > 0) &&
    160      (frame_header_.frame_type == kRegularFrame ||
    161       frame_header_.frame_type == kSkipProgressive)) {
    162    ++dec_state_->visible_frame_index;
    163    dec_state_->nonvisible_frame_index = 0;
    164  } else {
    165    ++dec_state_->nonvisible_frame_index;
    166  }
    167 
    168  // Read TOC.
    169  const size_t toc_entries =
    170      NumTocEntries(num_groups, frame_dim_.num_dc_groups, num_passes);
    171  std::vector<uint32_t> sizes;
    172  std::vector<coeff_order_t> permutation;
    173  JXL_RETURN_IF_ERROR(
    174      ReadToc(memory_manager, toc_entries, br, &sizes, &permutation));
    175  bool have_permutation = !permutation.empty();
    176  toc_.resize(toc_entries);
    177  section_sizes_sum_ = 0;
    178  for (size_t i = 0; i < toc_entries; ++i) {
    179    toc_[i].size = sizes[i];
    180    size_t index = have_permutation ? permutation[i] : i;
    181    toc_[index].id = i;
    182    if (section_sizes_sum_ + toc_[i].size < section_sizes_sum_) {
    183      return JXL_FAILURE("group offset overflow");
    184    }
    185    section_sizes_sum_ += toc_[i].size;
    186  }
    187 
    188  if (JXL_DEBUG_V_LEVEL >= 3) {
    189    for (size_t i = 0; i < toc_entries; ++i) {
    190      JXL_DEBUG_V(3, "TOC entry %" PRIuS " size %" PRIuS " id %" PRIuS "", i,
    191                  toc_[i].size, toc_[i].id);
    192    }
    193  }
    194 
    195  JXL_ENSURE((br->TotalBitsConsumed() % kBitsPerByte) == 0);
    196  const size_t group_codes_begin = br->TotalBitsConsumed() / kBitsPerByte;
    197  JXL_ENSURE(!toc_.empty());
    198 
    199  // Overflow check.
    200  if (group_codes_begin + section_sizes_sum_ < group_codes_begin) {
    201    return JXL_FAILURE("Invalid group codes");
    202  }
    203 
    204  if (!frame_header_.chroma_subsampling.Is444() &&
    205      !(frame_header_.flags & FrameHeader::kSkipAdaptiveDCSmoothing) &&
    206      frame_header_.encoding == FrameEncoding::kVarDCT) {
    207    return JXL_FAILURE(
    208        "Non-444 chroma subsampling is not allowed when adaptive DC "
    209        "smoothing is enabled");
    210  }
    211  return true;
    212 }
    213 
    214 Status FrameDecoder::InitFrameOutput() {
    215  JXL_RETURN_IF_ERROR(
    216      InitializePassesSharedState(frame_header_, &dec_state_->shared_storage));
    217  JXL_RETURN_IF_ERROR(dec_state_->Init(frame_header_));
    218  modular_frame_decoder_.Init(frame_dim_);
    219 
    220  if (decoded_->IsJPEG()) {
    221    if (frame_header_.encoding == FrameEncoding::kModular) {
    222      return JXL_FAILURE("Cannot output JPEG from Modular");
    223    }
    224    jpeg::JPEGData* jpeg_data = decoded_->jpeg_data.get();
    225    size_t num_components = jpeg_data->components.size();
    226    if (num_components != 1 && num_components != 3) {
    227      return JXL_FAILURE("Invalid number of components");
    228    }
    229    if (frame_header_.nonserialized_metadata->m.xyb_encoded) {
    230      return JXL_FAILURE("Cannot decode to JPEG an XYB image");
    231    }
    232    auto jpeg_c_map = JpegOrder(ColorTransform::kYCbCr, num_components == 1);
    233    decoded_->jpeg_data->width = frame_dim_.xsize;
    234    decoded_->jpeg_data->height = frame_dim_.ysize;
    235    for (size_t c = 0; c < num_components; c++) {
    236      auto& component = jpeg_data->components[jpeg_c_map[c]];
    237      component.width_in_blocks =
    238          frame_dim_.xsize_blocks >> frame_header_.chroma_subsampling.HShift(c);
    239      component.height_in_blocks =
    240          frame_dim_.ysize_blocks >> frame_header_.chroma_subsampling.VShift(c);
    241      component.h_samp_factor =
    242          1 << frame_header_.chroma_subsampling.RawHShift(c);
    243      component.v_samp_factor =
    244          1 << frame_header_.chroma_subsampling.RawVShift(c);
    245      component.coeffs.resize(component.width_in_blocks *
    246                              component.height_in_blocks * jxl::kDCTBlockSize);
    247    }
    248  }
    249 
    250  // Clear the state.
    251  decoded_dc_global_ = false;
    252  decoded_ac_global_ = false;
    253  is_finalized_ = false;
    254  finalized_dc_ = false;
    255  num_sections_done_ = 0;
    256  decoded_dc_groups_.clear();
    257  decoded_dc_groups_.resize(frame_dim_.num_dc_groups);
    258  decoded_passes_per_ac_group_.clear();
    259  decoded_passes_per_ac_group_.resize(frame_dim_.num_groups, 0);
    260  processed_section_.clear();
    261  processed_section_.resize(toc_.size());
    262  allocated_ = false;
    263  return true;
    264 }
    265 
    266 Status FrameDecoder::ProcessDCGlobal(BitReader* br) {
    267  PassesSharedState& shared = dec_state_->shared_storage;
    268  JxlMemoryManager* memory_manager = shared.memory_manager;
    269  if (frame_header_.flags & FrameHeader::kPatches) {
    270    bool uses_extra_channels = false;
    271    JXL_RETURN_IF_ERROR(shared.image_features.patches.Decode(
    272        memory_manager, br, frame_dim_.xsize_padded, frame_dim_.ysize_padded,
    273        shared.metadata->m.num_extra_channels, &uses_extra_channels));
    274    if (uses_extra_channels && frame_header_.upsampling != 1) {
    275      for (size_t ecups : frame_header_.extra_channel_upsampling) {
    276        if (ecups != frame_header_.upsampling) {
    277          return JXL_FAILURE(
    278              "Cannot use extra channels in patches if color channels are "
    279              "subsampled differently from extra channels");
    280        }
    281      }
    282    }
    283  } else {
    284    shared.image_features.patches.Clear();
    285  }
    286  shared.image_features.splines.Clear();
    287  if (frame_header_.flags & FrameHeader::kSplines) {
    288    JXL_RETURN_IF_ERROR(shared.image_features.splines.Decode(
    289        memory_manager, br, frame_dim_.xsize * frame_dim_.ysize));
    290  }
    291  if (frame_header_.flags & FrameHeader::kNoise) {
    292    JXL_RETURN_IF_ERROR(DecodeNoise(br, &shared.image_features.noise_params));
    293  }
    294  JXL_RETURN_IF_ERROR(dec_state_->shared_storage.matrices.DecodeDC(br));
    295 
    296  if (frame_header_.encoding == FrameEncoding::kVarDCT) {
    297    JXL_RETURN_IF_ERROR(
    298        jxl::DecodeGlobalDCInfo(br, decoded_->IsJPEG(), dec_state_, pool_));
    299  }
    300  // Splines' draw cache uses the color correlation map.
    301  if (frame_header_.flags & FrameHeader::kSplines) {
    302    JXL_RETURN_IF_ERROR(shared.image_features.splines.InitializeDrawCache(
    303        frame_dim_.xsize_upsampled, frame_dim_.ysize_upsampled,
    304        dec_state_->shared->cmap.base()));
    305  }
    306  Status dec_status = modular_frame_decoder_.DecodeGlobalInfo(
    307      br, frame_header_, /*allow_truncated_group=*/false);
    308  if (dec_status.IsFatalError()) return dec_status;
    309  if (dec_status) {
    310    decoded_dc_global_ = true;
    311  }
    312  return dec_status;
    313 }
    314 
    315 Status FrameDecoder::ProcessDCGroup(size_t dc_group_id, BitReader* br) {
    316  const size_t gx = dc_group_id % frame_dim_.xsize_dc_groups;
    317  const size_t gy = dc_group_id / frame_dim_.xsize_dc_groups;
    318  const LoopFilter& lf = frame_header_.loop_filter;
    319  if (frame_header_.encoding == FrameEncoding::kVarDCT &&
    320      !(frame_header_.flags & FrameHeader::kUseDcFrame)) {
    321    JXL_RETURN_IF_ERROR(modular_frame_decoder_.DecodeVarDCTDC(
    322        frame_header_, dc_group_id, br, dec_state_));
    323  }
    324  const Rect mrect(gx * frame_dim_.dc_group_dim, gy * frame_dim_.dc_group_dim,
    325                   frame_dim_.dc_group_dim, frame_dim_.dc_group_dim);
    326  JXL_RETURN_IF_ERROR(modular_frame_decoder_.DecodeGroup(
    327      frame_header_, mrect, br, 3, 1000,
    328      ModularStreamId::ModularDC(dc_group_id),
    329      /*zerofill=*/false, nullptr, nullptr,
    330      /*allow_truncated=*/false));
    331  if (frame_header_.encoding == FrameEncoding::kVarDCT) {
    332    JXL_RETURN_IF_ERROR(modular_frame_decoder_.DecodeAcMetadata(
    333        frame_header_, dc_group_id, br, dec_state_));
    334  } else if (lf.epf_iters > 0) {
    335    FillImage(kInvSigmaNum / lf.epf_sigma_for_modular, &dec_state_->sigma);
    336  }
    337  decoded_dc_groups_[dc_group_id] = JXL_TRUE;
    338  return true;
    339 }
    340 
    341 Status FrameDecoder::FinalizeDC() {
    342  // Do Adaptive DC smoothing if enabled. This *must* happen between all the
    343  // ProcessDCGroup and ProcessACGroup.
    344  JxlMemoryManager* memory_manager = dec_state_->memory_manager();
    345  if (frame_header_.encoding == FrameEncoding::kVarDCT &&
    346      !(frame_header_.flags & FrameHeader::kSkipAdaptiveDCSmoothing) &&
    347      !(frame_header_.flags & FrameHeader::kUseDcFrame)) {
    348    JXL_RETURN_IF_ERROR(AdaptiveDCSmoothing(
    349        memory_manager, dec_state_->shared->quantizer.MulDC(),
    350        &dec_state_->shared_storage.dc_storage, pool_));
    351  }
    352 
    353  finalized_dc_ = true;
    354  return true;
    355 }
    356 
    357 Status FrameDecoder::AllocateOutput() {
    358  if (allocated_) return true;
    359  modular_frame_decoder_.MaybeDropFullImage();
    360  decoded_->origin = frame_header_.frame_origin;
    361  JXL_RETURN_IF_ERROR(
    362      dec_state_->InitForAC(frame_header_.passes.num_passes, nullptr));
    363  allocated_ = true;
    364  return true;
    365 }
    366 
    367 Status FrameDecoder::ProcessACGlobal(BitReader* br) {
    368  JXL_ENSURE(finalized_dc_);
    369  JxlMemoryManager* memory_manager = dec_state_->memory_manager();
    370 
    371  // Decode AC group.
    372  if (frame_header_.encoding == FrameEncoding::kVarDCT) {
    373    JXL_RETURN_IF_ERROR(dec_state_->shared_storage.matrices.Decode(
    374        memory_manager, br, &modular_frame_decoder_));
    375    JXL_RETURN_IF_ERROR(dec_state_->shared_storage.matrices.EnsureComputed(
    376        memory_manager, dec_state_->used_acs));
    377 
    378    size_t num_histo_bits =
    379        CeilLog2Nonzero(dec_state_->shared->frame_dim.num_groups);
    380    dec_state_->shared_storage.num_histograms =
    381        1 + br->ReadBits(num_histo_bits);
    382 
    383    JXL_DEBUG_V(3,
    384                "Processing AC global with %d passes and %" PRIuS
    385                " sets of histograms",
    386                frame_header_.passes.num_passes,
    387                dec_state_->shared_storage.num_histograms);
    388 
    389    dec_state_->code.resize(kMaxNumPasses);
    390    dec_state_->context_map.resize(kMaxNumPasses);
    391    // Read coefficient orders and histograms.
    392    size_t max_num_bits_ac = 0;
    393    for (size_t i = 0; i < frame_header_.passes.num_passes; i++) {
    394      uint16_t used_orders = U32Coder::Read(kOrderEnc, br);
    395      JXL_RETURN_IF_ERROR(DecodeCoeffOrders(
    396          memory_manager, used_orders, dec_state_->used_acs,
    397          &dec_state_->shared_storage
    398               .coeff_orders[i * dec_state_->shared_storage.coeff_order_size],
    399          br));
    400      size_t num_contexts =
    401          dec_state_->shared->num_histograms *
    402          dec_state_->shared_storage.block_ctx_map.NumACContexts();
    403      JXL_RETURN_IF_ERROR(DecodeHistograms(memory_manager, br, num_contexts,
    404                                           &dec_state_->code[i],
    405                                           &dec_state_->context_map[i]));
    406      // Add extra values to enable the cheat in hot loop of DecodeACVarBlock.
    407      dec_state_->context_map[i].resize(
    408          num_contexts + kZeroDensityContextLimit - kZeroDensityContextCount);
    409      max_num_bits_ac =
    410          std::max(max_num_bits_ac, dec_state_->code[i].max_num_bits);
    411    }
    412    max_num_bits_ac += CeilLog2Nonzero(frame_header_.passes.num_passes);
    413    // 16-bit buffer for decoding to JPEG are not implemented.
    414    // TODO(veluca): figure out the exact limit - 16 should still work with
    415    // 16-bit buffers, but we are excluding it for safety.
    416    bool use_16_bit = max_num_bits_ac < 16 && !decoded_->IsJPEG();
    417    bool store = frame_header_.passes.num_passes > 1;
    418    size_t xs = store ? kGroupDim * kGroupDim : 0;
    419    size_t ys = store ? frame_dim_.num_groups : 0;
    420    if (use_16_bit) {
    421      JXL_ASSIGN_OR_RETURN(dec_state_->coefficients,
    422                           ACImageT<int16_t>::Make(memory_manager, xs, ys));
    423    } else {
    424      JXL_ASSIGN_OR_RETURN(dec_state_->coefficients,
    425                           ACImageT<int32_t>::Make(memory_manager, xs, ys));
    426    }
    427    if (store) {
    428      dec_state_->coefficients->ZeroFill();
    429    }
    430  }
    431 
    432  // Set JPEG decoding data.
    433  if (decoded_->IsJPEG()) {
    434    decoded_->color_transform = frame_header_.color_transform;
    435    decoded_->chroma_subsampling = frame_header_.chroma_subsampling;
    436    const std::vector<QuantEncoding>& qe =
    437        dec_state_->shared_storage.matrices.encodings();
    438    if (qe.empty() || qe[0].mode != QuantEncoding::Mode::kQuantModeRAW ||
    439        std::abs(qe[0].qraw.qtable_den - 1.f / (8 * 255)) > 1e-8f) {
    440      return JXL_FAILURE(
    441          "Quantization table is not a JPEG quantization table.");
    442    }
    443    jpeg::JPEGData* jpeg_data = decoded_->jpeg_data.get();
    444    size_t num_components = jpeg_data->components.size();
    445    bool is_gray = (num_components == 1);
    446    JXL_ENSURE(frame_header_.color_transform != ColorTransform::kXYB);
    447    auto jpeg_c_map = JpegOrder(frame_header_.color_transform, is_gray);
    448    size_t qt_set = 0;
    449    JXL_ENSURE(num_components <= 3);
    450    JXL_ENSURE(qe[0].qraw.qtable->size() == 3 * 8 * 8);
    451    int* qtable = qe[0].qraw.qtable->data();
    452    for (size_t c = 0; c < num_components; c++) {
    453      // TODO(eustas): why 1-st quant table for gray?
    454      size_t quant_c = is_gray ? 1 : c;
    455      size_t qpos = jpeg_data->components[jpeg_c_map[c]].quant_idx;
    456      JXL_ENSURE(qpos != jpeg_data->quant.size());
    457      qt_set |= 1 << qpos;
    458      for (size_t x = 0; x < 8; x++) {
    459        for (size_t y = 0; y < 8; y++) {
    460          jpeg_data->quant[qpos].values[x * 8 + y] =
    461              qtable[quant_c * 64 + y * 8 + x];
    462        }
    463      }
    464    }
    465    for (size_t i = 0; i < jpeg_data->quant.size(); i++) {
    466      if (qt_set & (1 << i)) continue;
    467      if (i == 0) return JXL_FAILURE("First quant table unused.");
    468      // Unused quant table is set to copy of previous quant table
    469      for (size_t j = 0; j < 64; j++) {
    470        jpeg_data->quant[i].values[j] = jpeg_data->quant[i - 1].values[j];
    471      }
    472    }
    473  }
    474  decoded_ac_global_ = true;
    475  return true;
    476 }
    477 
    478 Status FrameDecoder::ProcessACGroup(size_t ac_group_id,
    479                                    BitReader* JXL_RESTRICT* br,
    480                                    size_t num_passes, size_t thread,
    481                                    bool force_draw, bool dc_only) {
    482  size_t group_dim = frame_dim_.group_dim;
    483  const size_t gx = ac_group_id % frame_dim_.xsize_groups;
    484  const size_t gy = ac_group_id / frame_dim_.xsize_groups;
    485  const size_t x = gx * group_dim;
    486  const size_t y = gy * group_dim;
    487  JxlMemoryManager* memory_manager = dec_state_->memory_manager();
    488  JXL_DEBUG_V(3,
    489              "Processing AC group %" PRIuS "(%" PRIuS ",%" PRIuS
    490              ") group_dim: %" PRIuS " decoded passes: %u new passes: %" PRIuS,
    491              ac_group_id, gx, gy, group_dim,
    492              decoded_passes_per_ac_group_[ac_group_id], num_passes);
    493 
    494  RenderPipelineInput render_pipeline_input =
    495      dec_state_->render_pipeline->GetInputBuffers(ac_group_id, thread);
    496 
    497  bool should_run_pipeline = true;
    498 
    499  if (frame_header_.encoding == FrameEncoding::kVarDCT) {
    500    JXL_RETURN_IF_ERROR(group_dec_caches_[thread].InitOnce(
    501        memory_manager, frame_header_.passes.num_passes, dec_state_->used_acs));
    502    JXL_RETURN_IF_ERROR(DecodeGroup(
    503        frame_header_, br, num_passes, ac_group_id, dec_state_,
    504        &group_dec_caches_[thread], thread, render_pipeline_input,
    505        decoded_->jpeg_data.get(), decoded_passes_per_ac_group_[ac_group_id],
    506        force_draw, dc_only, &should_run_pipeline));
    507  }
    508 
    509  // don't limit to image dimensions here (is done in DecodeGroup)
    510  const Rect mrect(x, y, group_dim, group_dim);
    511  bool modular_ready = false;
    512  size_t pass0 = decoded_passes_per_ac_group_[ac_group_id];
    513  size_t pass1 =
    514      force_draw ? frame_header_.passes.num_passes : pass0 + num_passes;
    515  for (size_t i = pass0; i < pass1; ++i) {
    516    int minShift;
    517    int maxShift;
    518    frame_header_.passes.GetDownsamplingBracket(i, minShift, maxShift);
    519    bool modular_pass_ready = true;
    520    JXL_DEBUG_V(2, "Decoding modular in group %d pass %d",
    521                static_cast<int>(ac_group_id), static_cast<int>(i));
    522    if (i < pass0 + num_passes) {
    523      JXL_DEBUG_V(2, "Bit reader position: %" PRIuS " / %" PRIuS,
    524                  br[i - pass0]->TotalBitsConsumed(),
    525                  br[i - pass0]->TotalBytes() * kBitsPerByte);
    526      JXL_RETURN_IF_ERROR(modular_frame_decoder_.DecodeGroup(
    527          frame_header_, mrect, br[i - pass0], minShift, maxShift,
    528          ModularStreamId::ModularAC(ac_group_id, i),
    529          /*zerofill=*/false, dec_state_, &render_pipeline_input,
    530          /*allow_truncated=*/false, &modular_pass_ready));
    531    } else {
    532      JXL_RETURN_IF_ERROR(modular_frame_decoder_.DecodeGroup(
    533          frame_header_, mrect, nullptr, minShift, maxShift,
    534          ModularStreamId::ModularAC(ac_group_id, i), /*zerofill=*/true,
    535          dec_state_, &render_pipeline_input,
    536          /*allow_truncated=*/false, &modular_pass_ready));
    537    }
    538    if (modular_pass_ready) modular_ready = true;
    539  }
    540  decoded_passes_per_ac_group_[ac_group_id] += num_passes;
    541 
    542  if ((frame_header_.flags & FrameHeader::kNoise) != 0) {
    543    PrepareNoiseInput(*dec_state_, frame_dim_, frame_header_, ac_group_id,
    544                      thread);
    545  }
    546 
    547  if (!modular_frame_decoder_.UsesFullImage() && !decoded_->IsJPEG()) {
    548    if (should_run_pipeline && modular_ready) {
    549      JXL_RETURN_IF_ERROR(render_pipeline_input.Done());
    550    } else if (force_draw) {
    551      return JXL_FAILURE("Modular group decoding failed.");
    552    }
    553  }
    554  return true;
    555 }
    556 
    557 void FrameDecoder::MarkSections(const SectionInfo* sections, size_t num,
    558                                const SectionStatus* section_status) {
    559  num_sections_done_ += num;
    560  for (size_t i = 0; i < num; i++) {
    561    if (section_status[i] != SectionStatus::kDone) {
    562      processed_section_[sections[i].id] = JXL_FALSE;
    563      num_sections_done_--;
    564    }
    565  }
    566 }
    567 
    568 Status FrameDecoder::ProcessSections(const SectionInfo* sections, size_t num,
    569                                     SectionStatus* section_status) {
    570  if (num == 0) return true;  // Nothing to process
    571  std::fill(section_status, section_status + num, SectionStatus::kSkipped);
    572  size_t dc_global_sec = num;
    573  size_t ac_global_sec = num;
    574  std::vector<size_t> dc_group_sec(frame_dim_.num_dc_groups, num);
    575  std::vector<std::vector<size_t>> ac_group_sec(
    576      frame_dim_.num_groups,
    577      std::vector<size_t>(frame_header_.passes.num_passes, num));
    578  // This keeps track of the number of ac passes we want to process during this
    579  // call of ProcessSections.
    580  std::vector<size_t> desired_num_ac_passes(frame_dim_.num_groups);
    581  bool single_section =
    582      frame_dim_.num_groups == 1 && frame_header_.passes.num_passes == 1;
    583  if (single_section) {
    584    JXL_ENSURE(num == 1);
    585    JXL_ENSURE(sections[0].id == 0);
    586    if (processed_section_[0] == JXL_FALSE) {
    587      processed_section_[0] = JXL_TRUE;
    588      ac_group_sec[0].resize(1);
    589      dc_global_sec = ac_global_sec = dc_group_sec[0] = ac_group_sec[0][0] = 0;
    590      desired_num_ac_passes[0] = 1;
    591    } else {
    592      section_status[0] = SectionStatus::kDuplicate;
    593    }
    594  } else {
    595    size_t ac_global_index = frame_dim_.num_dc_groups + 1;
    596    for (size_t i = 0; i < num; i++) {
    597      JXL_ENSURE(sections[i].id < processed_section_.size());
    598      if (processed_section_[sections[i].id]) {
    599        section_status[i] = SectionStatus::kDuplicate;
    600        continue;
    601      }
    602      if (sections[i].id == 0) {
    603        dc_global_sec = i;
    604      } else if (sections[i].id < ac_global_index) {
    605        dc_group_sec[sections[i].id - 1] = i;
    606      } else if (sections[i].id == ac_global_index) {
    607        ac_global_sec = i;
    608      } else {
    609        size_t ac_idx = sections[i].id - ac_global_index - 1;
    610        size_t acg = ac_idx % frame_dim_.num_groups;
    611        size_t acp = ac_idx / frame_dim_.num_groups;
    612        if (acp >= frame_header_.passes.num_passes) {
    613          return JXL_FAILURE("Invalid section ID");
    614        }
    615        ac_group_sec[acg][acp] = i;
    616      }
    617      processed_section_[sections[i].id] = JXL_TRUE;
    618    }
    619    // Count number of new passes per group.
    620    for (size_t g = 0; g < ac_group_sec.size(); g++) {
    621      size_t j = 0;
    622      for (; j + decoded_passes_per_ac_group_[g] <
    623             frame_header_.passes.num_passes;
    624           j++) {
    625        if (ac_group_sec[g][j + decoded_passes_per_ac_group_[g]] == num) {
    626          break;
    627        }
    628      }
    629      desired_num_ac_passes[g] = j;
    630    }
    631  }
    632  if (dc_global_sec != num) {
    633    Status dc_global_status = ProcessDCGlobal(sections[dc_global_sec].br);
    634    if (dc_global_status.IsFatalError()) return dc_global_status;
    635    if (dc_global_status) {
    636      section_status[dc_global_sec] = SectionStatus::kDone;
    637    } else {
    638      section_status[dc_global_sec] = SectionStatus::kPartial;
    639    }
    640  }
    641 
    642  if (decoded_dc_global_) {
    643    const auto process_section = [this, &dc_group_sec, &num, &sections,
    644                                  &section_status](size_t i,
    645                                                   size_t thread) -> Status {
    646      if (dc_group_sec[i] != num) {
    647        JXL_RETURN_IF_ERROR(ProcessDCGroup(i, sections[dc_group_sec[i]].br));
    648        section_status[dc_group_sec[i]] = SectionStatus::kDone;
    649      }
    650      return true;
    651    };
    652    JXL_RETURN_IF_ERROR(RunOnPool(pool_, 0, dc_group_sec.size(),
    653                                  ThreadPool::NoInit, process_section,
    654                                  "DecodeDCGroup"));
    655  }
    656 
    657  if (!HasDcGroupToDecode() && !finalized_dc_) {
    658    PassesDecoderState::PipelineOptions pipeline_options;
    659    pipeline_options.use_slow_render_pipeline = use_slow_rendering_pipeline_;
    660    pipeline_options.coalescing = coalescing_;
    661    pipeline_options.render_spotcolors = render_spotcolors_;
    662    pipeline_options.render_noise = true;
    663    JXL_RETURN_IF_ERROR(dec_state_->PreparePipeline(
    664        frame_header_, &frame_header_.nonserialized_metadata->m, decoded_,
    665        pipeline_options));
    666    JXL_RETURN_IF_ERROR(FinalizeDC());
    667    JXL_RETURN_IF_ERROR(AllocateOutput());
    668    if (progressive_detail_ >= JxlProgressiveDetail::kDC) {
    669      MarkSections(sections, num, section_status);
    670      return true;
    671    }
    672  }
    673 
    674  if (finalized_dc_ && ac_global_sec != num && !decoded_ac_global_) {
    675    JXL_RETURN_IF_ERROR(ProcessACGlobal(sections[ac_global_sec].br));
    676    section_status[ac_global_sec] = SectionStatus::kDone;
    677  }
    678 
    679  if (progressive_detail_ >= JxlProgressiveDetail::kLastPasses) {
    680    // Mark that we only want the next progression pass.
    681    size_t target_complete_passes = NextNumPassesToPause();
    682    for (size_t i = 0; i < ac_group_sec.size(); i++) {
    683      desired_num_ac_passes[i] =
    684          std::min(desired_num_ac_passes[i],
    685                   target_complete_passes - decoded_passes_per_ac_group_[i]);
    686    }
    687  }
    688 
    689  if (decoded_ac_global_) {
    690    // Mark all the AC groups that we received as not complete yet.
    691    for (size_t i = 0; i < ac_group_sec.size(); i++) {
    692      if (desired_num_ac_passes[i] != 0) {
    693        dec_state_->render_pipeline->ClearDone(i);
    694      }
    695    }
    696 
    697    const auto prepare_storage = [this](size_t num_threads) -> Status {
    698      JXL_RETURN_IF_ERROR(
    699          PrepareStorage(num_threads, decoded_passes_per_ac_group_.size()));
    700      return true;
    701    };
    702    const auto process_group = [this, &ac_group_sec, &desired_num_ac_passes,
    703                                &num, &sections, &section_status](
    704                                   size_t g, size_t thread) -> Status {
    705      if (desired_num_ac_passes[g] == 0) {
    706        // no new AC pass, nothing to do
    707        return true;
    708      }
    709      (void)num;
    710      size_t first_pass = decoded_passes_per_ac_group_[g];
    711      BitReader* JXL_RESTRICT readers[kMaxNumPasses];
    712      for (size_t i = 0; i < desired_num_ac_passes[g]; i++) {
    713        JXL_ENSURE(ac_group_sec[g][first_pass + i] != num);
    714        readers[i] = sections[ac_group_sec[g][first_pass + i]].br;
    715      }
    716      JXL_RETURN_IF_ERROR(ProcessACGroup(
    717          g, readers, desired_num_ac_passes[g], GetStorageLocation(thread, g),
    718          /*force_draw=*/false, /*dc_only=*/false));
    719      for (size_t i = 0; i < desired_num_ac_passes[g]; i++) {
    720        section_status[ac_group_sec[g][first_pass + i]] = SectionStatus::kDone;
    721      }
    722      return true;
    723    };
    724    JXL_RETURN_IF_ERROR(RunOnPool(pool_, 0, ac_group_sec.size(),
    725                                  prepare_storage, process_group,
    726                                  "DecodeGroup"));
    727  }
    728 
    729  MarkSections(sections, num, section_status);
    730  return true;
    731 }
    732 
    733 Status FrameDecoder::Flush() {
    734  bool has_blending = frame_header_.blending_info.mode != BlendMode::kReplace ||
    735                      frame_header_.custom_size_or_origin;
    736  for (const auto& blending_info_ec :
    737       frame_header_.extra_channel_blending_info) {
    738    if (blending_info_ec.mode != BlendMode::kReplace) has_blending = true;
    739  }
    740  // No early Flush() if blending is enabled.
    741  if (has_blending && !is_finalized_) {
    742    return false;
    743  }
    744  // No early Flush() - nothing to do - if the frame is a kSkipProgressive
    745  // frame.
    746  if (frame_header_.frame_type == FrameType::kSkipProgressive &&
    747      !is_finalized_) {
    748    return true;
    749  }
    750  if (decoded_->IsJPEG()) {
    751    // Nothing to do.
    752    return true;
    753  }
    754  JXL_RETURN_IF_ERROR(AllocateOutput());
    755 
    756  uint32_t completely_decoded_ac_pass = *std::min_element(
    757      decoded_passes_per_ac_group_.begin(), decoded_passes_per_ac_group_.end());
    758  if (completely_decoded_ac_pass < frame_header_.passes.num_passes) {
    759    // We don't have all AC yet: force a draw of all the missing areas.
    760    // Mark all sections as not complete.
    761    for (size_t i = 0; i < decoded_passes_per_ac_group_.size(); i++) {
    762      if (decoded_passes_per_ac_group_[i] < frame_header_.passes.num_passes) {
    763        dec_state_->render_pipeline->ClearDone(i);
    764      }
    765    }
    766    const auto prepare_storage = [this](const size_t num_threads) -> Status {
    767      JXL_RETURN_IF_ERROR(
    768          PrepareStorage(num_threads, decoded_passes_per_ac_group_.size()));
    769      return true;
    770    };
    771    const auto process_group = [this](const uint32_t g,
    772                                      size_t thread) -> Status {
    773      if (decoded_passes_per_ac_group_[g] == frame_header_.passes.num_passes) {
    774        // This group was drawn already, nothing to do.
    775        return true;
    776      }
    777      BitReader* JXL_RESTRICT readers[kMaxNumPasses] = {};
    778      JXL_RETURN_IF_ERROR(ProcessACGroup(
    779          g, readers, /*num_passes=*/0, GetStorageLocation(thread, g),
    780          /*force_draw=*/true, /*dc_only=*/!decoded_ac_global_));
    781      return true;
    782    };
    783    JXL_RETURN_IF_ERROR(RunOnPool(pool_, 0, decoded_passes_per_ac_group_.size(),
    784                                  prepare_storage, process_group,
    785                                  "ForceDrawGroup"));
    786  }
    787 
    788  // undo global modular transforms and copy int pixel buffers to float ones
    789  JXL_RETURN_IF_ERROR(modular_frame_decoder_.FinalizeDecoding(
    790      frame_header_, dec_state_, pool_, is_finalized_));
    791 
    792  return true;
    793 }
    794 
    795 int FrameDecoder::SavedAs(const FrameHeader& header) {
    796  if (header.frame_type == FrameType::kDCFrame) {
    797    // bits 16, 32, 64, 128 for DC level
    798    return 16 << (header.dc_level - 1);
    799  } else if (header.CanBeReferenced()) {
    800    // bits 1, 2, 4 and 8 for the references
    801    return 1 << header.save_as_reference;
    802  }
    803 
    804  return 0;
    805 }
    806 
    807 bool FrameDecoder::HasEverything() const {
    808  if (!decoded_dc_global_) return false;
    809  if (!decoded_ac_global_) return false;
    810  if (HasDcGroupToDecode()) return false;
    811  for (const auto& nb_passes : decoded_passes_per_ac_group_) {
    812    if (nb_passes < frame_header_.passes.num_passes) return false;
    813  }
    814  return true;
    815 }
    816 
    817 int FrameDecoder::References() const {
    818  if (is_finalized_) {
    819    return 0;
    820  }
    821  if (!HasEverything()) return 0;
    822 
    823  int result = 0;
    824 
    825  // Blending
    826  if (frame_header_.frame_type == FrameType::kRegularFrame ||
    827      frame_header_.frame_type == FrameType::kSkipProgressive) {
    828    bool cropped = frame_header_.custom_size_or_origin;
    829    if (cropped || frame_header_.blending_info.mode != BlendMode::kReplace) {
    830      result |= (1 << frame_header_.blending_info.source);
    831    }
    832    const auto& extra = frame_header_.extra_channel_blending_info;
    833    for (const auto& ecbi : extra) {
    834      if (cropped || ecbi.mode != BlendMode::kReplace) {
    835        result |= (1 << ecbi.source);
    836      }
    837    }
    838  }
    839 
    840  // Patches
    841  if (frame_header_.flags & FrameHeader::kPatches) {
    842    result |= dec_state_->shared->image_features.patches.GetReferences();
    843  }
    844 
    845  // DC Level
    846  if (frame_header_.flags & FrameHeader::kUseDcFrame) {
    847    // Reads from the next dc level
    848    int dc_level = frame_header_.dc_level + 1;
    849    // bits 16, 32, 64, 128 for DC level
    850    result |= (16 << (dc_level - 1));
    851  }
    852 
    853  return result;
    854 }
    855 
    856 Status FrameDecoder::FinalizeFrame() {
    857  if (is_finalized_) {
    858    return JXL_FAILURE("FinalizeFrame called multiple times");
    859  }
    860  is_finalized_ = true;
    861  if (decoded_->IsJPEG()) {
    862    // Nothing to do.
    863    return true;
    864  }
    865 
    866  // undo global modular transforms and copy int pixel buffers to float ones
    867  JXL_RETURN_IF_ERROR(
    868      modular_frame_decoder_.FinalizeDecoding(frame_header_, dec_state_, pool_,
    869                                              /*inplace=*/true));
    870 
    871  if (frame_header_.CanBeReferenced()) {
    872    auto& info = dec_state_->shared_storage
    873                     .reference_frames[frame_header_.save_as_reference];
    874    *info.frame = std::move(dec_state_->frame_storage_for_referencing);
    875    info.ib_is_in_xyb = frame_header_.save_before_color_transform;
    876  }
    877  return true;
    878 }
    879 
    880 }  // namespace jxl