backport decoding improvements for libjxl

third_party/jpeg-xl/lib/extras/size_constraints.h

@@ -16,22 +16,24 @@ namespace jxl {
 struct SizeConstraints {
   // Upper limit on pixel dimensions/area, enforced by VerifyDimensions
   // (called from decoders). Fuzzers set smaller values to limit memory use.
-  uint32_t dec_max_xsize = 0xFFFFFFFFu;
-  uint32_t dec_max_ysize = 0xFFFFFFFFu;
-  uint64_t dec_max_pixels = 0xFFFFFFFFu;  // Might be up to ~0ull
+  // Default values correspond to JXL level 10.
+  uint32_t dec_max_xsize = 1u << 30;
+  uint32_t dec_max_ysize = 1u << 30;
+  uint64_t dec_max_pixels = static_cast<uint64_t>(1u) << 40;
 };
 
 template <typename T,
           class = typename std::enable_if<std::is_unsigned<T>::value>::type>
 Status VerifyDimensions(const SizeConstraints* constraints, T xs, T ys) {
-  if (!constraints) return true;
+  SizeConstraints limit = {};
+  if (constraints) limit = *constraints;
 
   if (xs == 0 || ys == 0) return JXL_FAILURE("Empty image.");
-  if (xs > constraints->dec_max_xsize) return JXL_FAILURE("Image too wide.");
-  if (ys > constraints->dec_max_ysize) return JXL_FAILURE("Image too tall.");
+  if (xs > limit.dec_max_xsize) return JXL_FAILURE("Image too wide.");
+  if (ys > limit.dec_max_ysize) return JXL_FAILURE("Image too tall.");
 
   const uint64_t num_pixels = static_cast<uint64_t>(xs) * ys;
-  if (num_pixels > constraints->dec_max_pixels) {
+  if (num_pixels > limit.dec_max_pixels) {
     return JXL_FAILURE("Image too big.");
   }
 

third_party/jpeg-xl/lib/jxl/base/float.h

@@ -18,7 +18,7 @@
 namespace jxl {
 
 namespace detail {
-// Based on highway scalar implementation, for testing
+
 static JXL_INLINE float LoadFloat16(uint16_t bits16) {
   const uint32_t sign = bits16 >> 15;
   const uint32_t biased_exp = (bits16 >> 10) & 0x1F;
@@ -32,7 +32,8 @@ static JXL_INLINE float LoadFloat16(uint16_t bits16) {
   }
 
   // Normalized: convert the representation directly (faster than ldexp/tables).
-  const uint32_t biased_exp32 = biased_exp + (127 - 15);
+  const uint32_t biased_exp32 =
+      biased_exp == 0b11111 ? 0b11111111 : biased_exp + (127 - 15);
   const uint32_t mantissa32 = mantissa << (23 - 10);
   const uint32_t bits32 = (sign << 31) | (biased_exp32 << 23) | mantissa32;
 

third_party/jpeg-xl/lib/jxl/cms/color_encoding_cms.h

@@ -511,7 +511,7 @@ struct ColorEncoding {
   // Checks if the color space and transfer function are the same, ignoring
   // rendering intent and ICC bytes
   bool SameColorEncoding(const ColorEncoding& other) const {
-    return SameColorSpace(other) && tf.IsSame(other.tf);
+    return SameColorSpace(other) && tf.IsSame(other.tf) && !cmyk && !other.cmyk;
   }
 
   // Returns true if all fields have been initialized (possibly to kUnknown).

third_party/jpeg-xl/lib/jxl/color_encoding_internal.h

@@ -150,6 +150,7 @@ struct ColorEncoding : public Fields {
   Status SetICC(IccBytes&& icc, const JxlCmsInterface* cms) {
     JXL_ENSURE(cms != nullptr);
     JXL_ENSURE(!icc.empty());
+    storage_.have_fields = true;
     want_icc_ = storage_.SetFieldsFromICC(std::move(icc), *cms);
     return want_icc_;
   }

third_party/jpeg-xl/lib/jxl/dec_ans.cc

@@ -126,9 +126,9 @@ Status ReadHistogram(int precision_bits, std::vector<int32_t>* counts,
       input->Refill();  // for PeekFixedBits + Advance
       int idx = input->PeekFixedBits<7>();
       input->Consume(huff[idx][0]);
-      logcounts[i] = huff[idx][1];
+      logcounts[i] = int(huff[idx][1]) - 1;
       // The RLE symbol.
-      if (logcounts[i] == ANS_LOG_TAB_SIZE + 1) {
+      if (logcounts[i] == ANS_LOG_TAB_SIZE) {
         int rle_length = DecodeVarLenUint8(input);
         same[i] = rle_length + 5;
         i += rle_length + 3;
@@ -142,7 +142,7 @@ Status ReadHistogram(int precision_bits, std::vector<int32_t>* counts,
     // Invalid input, e.g. due to invalid usage of RLE.
     if (omit_pos < 0) return JXL_FAILURE("Invalid histogram.");
     if (static_cast<size_t>(omit_pos) + 1 < logcounts.size() &&
-        logcounts[omit_pos + 1] == ANS_TAB_SIZE + 1) {
+        logcounts[omit_pos + 1] == ANS_LOG_TAB_SIZE) {
       return JXL_FAILURE("Invalid histogram.");
     }
     int prev = 0;
@@ -158,18 +158,17 @@ Status ReadHistogram(int precision_bits, std::vector<int32_t>* counts,
         (*counts)[i] = prev;
         numsame--;
       } else {
-        unsigned int code = logcounts[i];
+        int code = logcounts[i];
         // omit_pos may not be negative at this point (checked before).
-        if (i == static_cast<size_t>(omit_pos)) {
+        if (i == static_cast<size_t>(omit_pos) || code < 0) {
           continue;
-        } else if (code == 0) {
-          continue;
-        } else if (code == 1) {
-          (*counts)[i] = 1;
+        } else if (shift == 0 || code == 0) {
+          // `shift = 0` means `bitcount = 0`
+          (*counts)[i] = 1 << code;
         } else {
-          int bitcount = GetPopulationCountPrecision(code - 1, shift);
-          (*counts)[i] = (1u << (code - 1)) +
-                         (input->ReadBits(bitcount) << (code - 1 - bitcount));
+          int bitcount = GetPopulationCountPrecision(code, shift);
+          (*counts)[i] = (1 << code) +
+                         (input->ReadBits(bitcount) << (code - bitcount));
         }
       }
       total_count += (*counts)[i];

third_party/jpeg-xl/lib/jxl/dec_ans.h

@@ -379,6 +379,8 @@ class ANSSymbolReader {
   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): Check if we could deal with copy tail as well.
+    if (num_to_copy_ != 0) 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_];

third_party/jpeg-xl/lib/jxl/dec_cache.cc

@@ -106,9 +106,7 @@ Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,
                                            PipelineOptions options) {
   JxlMemoryManager* memory_manager = this->memory_manager();
   size_t num_c = 3 + frame_header.nonserialized_metadata->m.num_extra_channels;
-  bool render_noise =
-      (options.render_noise && (frame_header.flags & FrameHeader::kNoise) != 0);
-  size_t num_tmp_c = render_noise ? 3 : 0;
+  size_t num_tmp_c = options.render_noise ? 3 : 0;
 
   if (frame_header.CanBeReferenced()) {
     // Necessary so that SetInputSizes() can allocate output buffers as needed.
@@ -169,7 +167,8 @@ Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,
          ec++) {
       if (frame_header.extra_channel_upsampling[ec] != 1) {
         JXL_RETURN_IF_ERROR(builder.AddStage(GetUpsamplingStage(
-            frame_header.nonserialized_metadata->transform_data, 3 + ec,
+            memory_manager, frame_header.nonserialized_metadata->transform_data,
+            3 + ec,
             CeilLog2Nonzero(frame_header.extra_channel_upsampling[ec]))));
       }
     }
@@ -191,11 +190,11 @@ Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,
         (late_ec_upsample ? frame_header.extra_channel_upsampling.size() : 0);
     for (size_t c = 0; c < nb_channels; c++) {
       JXL_RETURN_IF_ERROR(builder.AddStage(GetUpsamplingStage(
-          frame_header.nonserialized_metadata->transform_data, c,
-          CeilLog2Nonzero(frame_header.upsampling))));
+          memory_manager, frame_header.nonserialized_metadata->transform_data,
+          c, CeilLog2Nonzero(frame_header.upsampling))));
     }
   }
-  if (render_noise) {
+  if (options.render_noise) {
     JXL_RETURN_IF_ERROR(builder.AddStage(GetConvolveNoiseStage(num_c)));
     JXL_RETURN_IF_ERROR(builder.AddStage(GetAddNoiseStage(
         shared->image_features.noise_params, shared->cmap.base(), num_c)));
@@ -330,6 +329,11 @@ Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,
         }
       }
       linear = false;
+    } else {
+      auto cms_stage = GetCmsStage(output_encoding_info, false);
+      if (cms_stage) {
+        JXL_RETURN_IF_ERROR(builder.AddStage(std::move(cms_stage)));
+      }
     }
     (void)linear;
 

third_party/jpeg-xl/lib/jxl/dec_cache.h

@@ -171,7 +171,8 @@ struct PassesDecoderState {
 
     used_acs = 0;
 
-    upsampler8x = GetUpsamplingStage(shared->metadata->transform_data, 0, 3);
+    upsampler8x = GetUpsamplingStage(memory_manager,
+                                     shared->metadata->transform_data, 0, 3);
     if (frame_header.loop_filter.epf_iters > 0) {
       JXL_ASSIGN_OR_RETURN(
           sigma,

third_party/jpeg-xl/lib/jxl/dec_frame.cc

@@ -539,7 +539,10 @@ Status FrameDecoder::ProcessACGroup(size_t ac_group_id,
   }
   decoded_passes_per_ac_group_[ac_group_id] += num_passes;
 
-  if ((frame_header_.flags & FrameHeader::kNoise) != 0) {
+  const bool render_noise =
+      ((frame_header_.flags & FrameHeader::kNoise) != 0) &&
+      (frame_header_.dc_level == 0);
+  if (render_noise) {
     PrepareNoiseInput(*dec_state_, frame_dim_, frame_header_, ac_group_id,
                       thread);
   }
@@ -654,12 +657,15 @@ Status FrameDecoder::ProcessSections(const SectionInfo* sections, size_t num,
                                   "DecodeDCGroup"));
   }
 
+  const bool render_noise =
+      ((frame_header_.flags & FrameHeader::kNoise) != 0) &&
+      (frame_header_.dc_level == 0);
   if (!HasDcGroupToDecode() && !finalized_dc_) {
     PassesDecoderState::PipelineOptions pipeline_options;
     pipeline_options.use_slow_render_pipeline = use_slow_rendering_pipeline_;
     pipeline_options.coalescing = coalescing_;
     pipeline_options.render_spotcolors = render_spotcolors_;
-    pipeline_options.render_noise = true;
+    pipeline_options.render_noise = render_noise;
     JXL_RETURN_IF_ERROR(dec_state_->PreparePipeline(
         frame_header_, &frame_header_.nonserialized_metadata->m, decoded_,
         pipeline_options));

third_party/jpeg-xl/lib/jxl/dec_frame.h

@@ -271,6 +271,8 @@ class FrameDecoder {
       JXL_RETURN_IF_ERROR(dec_state_->render_pipeline->PrepareForThreads(
           storage_size, use_group_ids));
     }
+    JXL_RETURN_IF_ERROR(
+        dec_state_->upsampler8x->PrepareForThreads(num_threads));
     return true;
   }
 

third_party/jpeg-xl/lib/jxl/dec_huffman.cc

@@ -239,20 +239,4 @@ bool HuffmanDecodingData::ReadFromBitStream(size_t alphabet_size,
   return (table_size > 0);
 }
 
-// Decodes the next Huffman coded symbol from the bit-stream.
-uint16_t HuffmanDecodingData::ReadSymbol(BitReader* br) const {
-  size_t n_bits;
-  const HuffmanCode* table = table_.data();
-  table += br->PeekBits(kHuffmanTableBits);
-  n_bits = table->bits;
-  if (n_bits > kHuffmanTableBits) {
-    br->Consume(kHuffmanTableBits);
-    n_bits -= kHuffmanTableBits;
-    table += table->value;
-    table += br->PeekBits(n_bits);
-  }
-  br->Consume(table->bits);
-  return table->value;
-}
-
 }  // namespace jxl

third_party/jpeg-xl/lib/jxl/dec_huffman.h

@@ -19,10 +19,24 @@ static constexpr size_t kHuffmanTableBits = 8u;
 struct HuffmanDecodingData {
   // Decodes the Huffman code lengths from the bit-stream and fills in the
   // pre-allocated table with the corresponding 2-level Huffman decoding table.
-  // Returns false if the Huffman code lengths can not de decoded.
+  // Returns false if the Huffman code lengths can not be decoded.
   bool ReadFromBitStream(size_t alphabet_size, BitReader* br);
 
-  uint16_t ReadSymbol(BitReader* br) const;
+  // Decodes the next Huffman coded symbol from the bit-stream.
+  JXL_INLINE uint16_t ReadSymbol(BitReader* br) const {
+    size_t n_bits;
+    const HuffmanCode* table = table_.data();
+    table += br->PeekBits(kHuffmanTableBits);
+    n_bits = table->bits;
+    if (n_bits > kHuffmanTableBits) {
+      br->Consume(kHuffmanTableBits);
+      n_bits -= kHuffmanTableBits;
+      table += table->value;
+      table += br->PeekBits(n_bits);
+    }
+    br->Consume(table->bits);
+    return table->value;
+  }
 
   std::vector<HuffmanCode> table_;
 };

third_party/jpeg-xl/lib/jxl/dec_modular.cc

@@ -127,6 +127,14 @@ Status int_to_float(const pixel_type* const JXL_RESTRICT row_in,
     }
     int exp = (f >> mant_bits);
     int mantissa = (f & ((1 << mant_bits) - 1));
+    if (exp == (1 << exp_bits) - 1) {
+      // NaN or infinity
+      f = (signbit ? 0x80000000 : 0);
+      f |= 0b11111111 << 23;
+      f |= mantissa << mant_shift;
+      memcpy(&row_out[x], &f, 4);
+      continue;
+    }
     mantissa <<= mant_shift;
     // Try to normalize only if there is space for maneuver.
     if (exp == 0 && exp_bits < 8) {

third_party/jpeg-xl/lib/jxl/dec_xyb.cc

@@ -237,7 +237,7 @@ Status OutputEncodingInfo::SetFromMetadata(const CodecMetadata& metadata) {
   bool orig_grey = orig_color_encoding.IsGray();
   return SetColorEncoding(!xyb_encoded || orig_ok
                               ? orig_color_encoding
-                              : ColorEncoding::LinearSRGB(orig_grey));
+                              : ColorEncoding::SRGB(orig_grey));
 }
 
 Status OutputEncodingInfo::MaybeSetColorEncoding(

third_party/jpeg-xl/lib/jxl/decode.cc

@@ -2659,12 +2659,27 @@ JxlDecoderStatus JxlDecoderSetOutputColorProfile(
   if (dec->post_headers) {
     return JXL_API_ERROR("too late to set the color encoding");
   }
+  auto& output_encoding = dec->passes_state->output_encoding_info;
+  auto& orig_encoding = dec->image_metadata.color_encoding;
+  jxl::ColorEncoding c_out;
+  bool same_encoding = false;
+  if (color_encoding) {
+    JXL_API_RETURN_IF_ERROR(c_out.FromExternal(*color_encoding));
+    same_encoding = c_out.SameColorEncoding(output_encoding.color_encoding);
+  }
   if ((!dec->passes_state->output_encoding_info.cms_set) &&
-      (icc_data != nullptr)) {
+      (icc_data != nullptr ||
+      (!dec->image_metadata.xyb_encoded && !same_encoding))) {
     return JXL_API_ERROR(
         "must set color management system via JxlDecoderSetCms");
   }
-  auto& output_encoding = dec->passes_state->output_encoding_info;
+  if (!orig_encoding.HaveFields() &&
+      dec->passes_state->output_encoding_info.cms_set) {
+    std::vector<uint8_t> tmp_icc = orig_encoding.ICC();
+    JXL_API_RETURN_IF_ERROR(orig_encoding.SetICC(
+        std::move(tmp_icc), &output_encoding.color_management_system));
+    output_encoding.orig_color_encoding = orig_encoding;
+  }
   if (color_encoding) {
     if (dec->image_metadata.color_encoding.IsGray() &&
         color_encoding->color_space != JXL_COLOR_SPACE_GRAY &&
@@ -2674,13 +2689,9 @@ JxlDecoderStatus JxlDecoderSetOutputColorProfile(
     if (color_encoding->color_space == JXL_COLOR_SPACE_UNKNOWN) {
       return JXL_API_ERROR("Unknown output colorspace");
     }
-    jxl::ColorEncoding c_out;
-    JXL_API_RETURN_IF_ERROR(c_out.FromExternal(*color_encoding));
     JXL_API_RETURN_IF_ERROR(!c_out.ICC().empty());
-    if (!c_out.SameColorEncoding(output_encoding.color_encoding)) {
-      JXL_API_RETURN_IF_ERROR(output_encoding.MaybeSetColorEncoding(c_out));
-      dec->image_metadata.color_encoding = output_encoding.color_encoding;
-    }
+    JXL_API_RETURN_IF_ERROR(output_encoding.MaybeSetColorEncoding(c_out));
+    dec->image_metadata.color_encoding = output_encoding.color_encoding;
     return JXL_DEC_SUCCESS;
   }
   // icc_data != nullptr

third_party/jpeg-xl/lib/jxl/decode_test.cc

@@ -1096,16 +1096,16 @@ TEST(DecodeTest, IccProfileTestXybEncoded) {
             JxlDecoderGetColorAsEncodedProfile(
                 dec, JXL_COLOR_PROFILE_TARGET_DATA, &pixel_encoding));
   EXPECT_EQ(JXL_PRIMARIES_SRGB, pixel_encoding.primaries);
-  // The API returns LINEAR by default when the colorspace cannot be represented
+  // The API returns SRGB by default when the colorspace cannot be represented
   // by enum values.
-  EXPECT_EQ(JXL_TRANSFER_FUNCTION_LINEAR, pixel_encoding.transfer_function);
+  EXPECT_EQ(JXL_TRANSFER_FUNCTION_SRGB, pixel_encoding.transfer_function);
 
   // Test the same but with integer format.
   EXPECT_EQ(JXL_DEC_SUCCESS,
             JxlDecoderGetColorAsEncodedProfile(
                 dec, JXL_COLOR_PROFILE_TARGET_DATA, &pixel_encoding));
   EXPECT_EQ(JXL_PRIMARIES_SRGB, pixel_encoding.primaries);
-  EXPECT_EQ(JXL_TRANSFER_FUNCTION_LINEAR, pixel_encoding.transfer_function);
+  EXPECT_EQ(JXL_TRANSFER_FUNCTION_SRGB, pixel_encoding.transfer_function);
 
   // Test after setting the preferred color profile to non-linear sRGB:
   // for XYB images with ICC profile, this setting is expected to take effect.
@@ -1800,11 +1800,14 @@ void SetPreferredColorProfileTest(
                                    xsize, ysize, num_channels, params);
   auto all_encodings = jxl::test::AllEncodings();
   // TODO(firsching): understand why XYB does not work together with icc_dst.
+  // TODO(jon): fix XYB output space in general
+  /*
   if (!icc_dst) {
     all_encodings.push_back(
         {jxl::ColorSpace::kXYB, jxl::WhitePoint::kD65, jxl::Primaries::kCustom,
          jxl::TransferFunction::kUnknown, jxl::RenderingIntent::kPerceptual});
   }
+  */
   for (const auto& c1 : all_encodings) {
     jxl::ColorEncoding c_out = jxl::test::ColorEncodingFromDescriptor(c1);
     float intensity_out = intensity_in;

third_party/jpeg-xl/lib/jxl/enc_heuristics.cc

@@ -810,7 +810,7 @@ StatusOr<Image3F> ReconstructImage(
   options.use_slow_render_pipeline = false;
   options.coalescing = false;
   options.render_spotcolors = false;
-  options.render_noise = true;
+  options.render_noise = ((frame_header.flags & FrameHeader::kNoise) != 0);
 
   JXL_RETURN_IF_ERROR(dec_state.PreparePipeline(
       frame_header, &shared.metadata->m, &decoded, options));