lib/jxl/toc.cc

   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/toc.h"
   7
   8 #include <stdint.h>
   9
  10 #include "lib/jxl/base/common.h"
  11 #include "lib/jxl/coeff_order.h"
  12 #include "lib/jxl/coeff_order_fwd.h"
  13 #include "lib/jxl/fields.h"
  14
  15 namespace jxl {
  16 size_t MaxBits(const size_t num_sizes) {
  17   const size_t entry_bits = U32Coder::MaxEncodedBits(kTocDist) * num_sizes;
  18   // permutation bit (not its tokens!), padding, entries, padding.
  19   return 1 + kBitsPerByte + entry_bits + kBitsPerByte;
  20 }
  21
  22 Status ReadToc(size_t toc_entries, BitReader* JXL_RESTRICT reader,
  23                std::vector<uint32_t>* JXL_RESTRICT sizes,
  24                std::vector<coeff_order_t>* JXL_RESTRICT permutation) {
  25   if (toc_entries > 65536) {
  26     // Prevent out of memory if invalid JXL codestream causes a bogus amount
  27     // of toc_entries such as 2720436919446 to be computed.
  28     // TODO(lode): verify whether 65536 is a reasonable upper bound
  29     return JXL_FAILURE("too many toc entries");
  30   }
  31
  32   sizes->clear();
  33   sizes->resize(toc_entries);
  34   if (reader->TotalBitsConsumed() >= reader->TotalBytes() * kBitsPerByte) {
  35     return JXL_STATUS(StatusCode::kNotEnoughBytes, "Not enough bytes for TOC");
  36   }
  37   const auto check_bit_budget = [&](size_t num_entries) -> Status {
  38     // U32Coder reads 2 bits to recognize variant and kTocDist cheapest variant
  39     // is Bits(10), this way at least 12 bits are required per toc-entry.
  40     size_t minimal_bit_cost = num_entries * (2 + 10);
  41     size_t bit_budget = reader->TotalBytes() * 8;
  42     size_t expenses = reader->TotalBitsConsumed();
  43     if ((expenses <= bit_budget) &&
  44         (minimal_bit_cost <= bit_budget - expenses)) {
  45       return true;
  46     }
  47     return JXL_STATUS(StatusCode::kNotEnoughBytes, "Not enough bytes for TOC");
  48   };
  49
  50   JXL_DASSERT(toc_entries > 0);
  51   if (reader->ReadFixedBits<1>() == 1) {
  52     JXL_RETURN_IF_ERROR(check_bit_budget(toc_entries));
  53     permutation->resize(toc_entries);
  54     JXL_RETURN_IF_ERROR(DecodePermutation(/*skip=*/0, toc_entries,
  55                                           permutation->data(), reader));
  56   }
  57   JXL_RETURN_IF_ERROR(reader->JumpToByteBoundary());
  58   JXL_RETURN_IF_ERROR(check_bit_budget(toc_entries));
  59   for (size_t i = 0; i < toc_entries; ++i) {
  60     (*sizes)[i] = U32Coder::Read(kTocDist, reader);
  61   }
  62   JXL_RETURN_IF_ERROR(reader->JumpToByteBoundary());
  63   JXL_RETURN_IF_ERROR(check_bit_budget(0));
  64   return true;
  65 }
  66
  67 Status ReadGroupOffsets(size_t toc_entries, BitReader* JXL_RESTRICT reader,
  68                         std::vector<uint64_t>* JXL_RESTRICT offsets,
  69                         std::vector<uint32_t>* JXL_RESTRICT sizes,
  70                         uint64_t* total_size) {
  71   std::vector<coeff_order_t> permutation;
  72   JXL_RETURN_IF_ERROR(ReadToc(toc_entries, reader, sizes, &permutation));
  73
  74   offsets->clear();
  75   offsets->resize(toc_entries);
  76
  77   // Prefix sum starting with 0 and ending with the offset of the last group
  78   uint64_t offset = 0;
  79   for (size_t i = 0; i < toc_entries; ++i) {
  80     if (offset + (*sizes)[i] < offset) {
  81       return JXL_FAILURE("group offset overflow");
  82     }
  83     (*offsets)[i] = offset;
  84     offset += (*sizes)[i];
  85   }
  86   if (total_size) {
  87     *total_size = offset;
  88   }
  89
  90   if (!permutation.empty()) {
  91     std::vector<uint64_t> permuted_offsets;
  92     std::vector<uint32_t> permuted_sizes;
  93     permuted_offsets.reserve(toc_entries);
  94     permuted_sizes.reserve(toc_entries);
  95     for (coeff_order_t index : permutation) {
  96       permuted_offsets.push_back((*offsets)[index]);
  97       permuted_sizes.push_back((*sizes)[index]);
  98     }
  99     std::swap(*offsets, permuted_offsets);
 100     std::swap(*sizes, permuted_sizes);
 101   }
 102
 103   return true;
 104 }
 105 }  // namespace jxl