Bug 1896717 - Update libjxl to 715b44238e67e521bba944e2864eb2933296e51c r=saschanaz

Differential Revision: https://phabricator.services.mozilla.com/D210347

media/libjxl/moz.yaml

@@ -10,9 +10,9 @@ origin:
 
   url: https://github.com/libjxl/libjxl
 
-  release: a741085a1dee343f143f5fdca3212ca13d66e401 (2024-04-23T11:57:55Z).
+  release: 715b44238e67e521bba944e2864eb2933296e51c (2024-05-14T15:25:56Z).
 
-  revision: a741085a1dee343f143f5fdca3212ca13d66e401
+  revision: 715b44238e67e521bba944e2864eb2933296e51c
 
   license: Apache-2.0
 

third_party/jpeg-xl/CMakeLists.txt

@@ -433,7 +433,7 @@ if(JPEGXL_ENABLE_MANPAGES)
   find_program(ASCIIDOC a2x)
   if(ASCIIDOC)
     file(STRINGS "${ASCIIDOC}" ASCIIDOC_SHEBANG LIMIT_COUNT 1)
-    if(ASCIIDOC_SHEBANG MATCHES "/sh|/bash" OR MINGW)
+    if(ASCIIDOC_SHEBANG MATCHES "sh$" OR MINGW)
       set(ASCIIDOC_PY_FOUND ON)
       # Run the program directly and set ASCIIDOC as empty.
       set(ASCIIDOC_PY "${ASCIIDOC}")

third_party/jpeg-xl/lib/extras/dec/decode.h

@@ -8,11 +8,9 @@
 
 // Facade for image decoders (PNG, PNM, ...).
 
-#include <stddef.h>
+#include <cstddef>
-#include <stdint.h>
+#include <cstdint>
-
 #include <string>
-#include <vector>
 
 #include "lib/extras/dec/color_hints.h"
 #include "lib/jxl/base/span.h"

third_party/jpeg-xl/lib/extras/dec/gif.cc

@@ -9,8 +9,8 @@
 #include <gif_lib.h>
 #endif
 #include <jxl/codestream_header.h>
-#include <string.h>
 
+#include <cstring>
 #include <memory>
 #include <utility>
 #include <vector>
@@ -18,7 +18,7 @@
 #include "lib/extras/size_constraints.h"
 #include "lib/jxl/base/compiler_specific.h"
 #include "lib/jxl/base/rect.h"
-#include "lib/jxl/sanitizers.h"
+#include "lib/jxl/base/sanitizers.h"
 
 namespace jxl {
 namespace extras {

third_party/jpeg-xl/lib/extras/dec/jpegli.cc

@@ -6,16 +6,16 @@
 #include "lib/extras/dec/jpegli.h"
 
 #include <setjmp.h>
-#include <stdint.h>
 
 #include <algorithm>
-#include <numeric>
+#include <cstdint>
+#include <memory>
 #include <utility>
 #include <vector>
 
 #include "lib/jpegli/decode.h"
+#include "lib/jxl/base/sanitizers.h"
 #include "lib/jxl/base/status.h"
-#include "lib/jxl/sanitizers.h"
 
 namespace jxl {
 namespace extras {
@@ -181,7 +181,9 @@ Status DecodeJpeg(const std::vector<uint8_t>& compressed,
     }
     int nbcomp = cinfo.num_components;
     if (nbcomp != 1 && nbcomp != 3) {
-      return failure("unsupported number of components in JPEG");
+      std::string msg =
+          "unsupported number of components in JPEG: " + std::to_string(nbcomp);
+      return failure(msg.c_str());
     }
     if (dparams.force_rgb) {
       cinfo.out_color_space = JCS_RGB;

third_party/jpeg-xl/lib/extras/dec/jpg.cc

@@ -8,16 +8,16 @@
 #if JPEGXL_ENABLE_JPEG
 #include "lib/jxl/base/include_jpeglib.h"  // NOLINT
 #endif
-#include <stdint.h>
 
 #include <algorithm>
+#include <cstdint>
 #include <numeric>
 #include <utility>
 #include <vector>
 
 #include "lib/extras/size_constraints.h"
+#include "lib/jxl/base/sanitizers.h"
 #include "lib/jxl/base/status.h"
-#include "lib/jxl/sanitizers.h"
 
 namespace jxl {
 namespace extras {

third_party/jpeg-xl/lib/extras/dec/pgx_test.cc

@@ -9,6 +9,7 @@
 
 #include "lib/extras/packed_image_convert.h"
 #include "lib/jxl/image_bundle.h"
+#include "lib/jxl/test_utils.h"
 #include "lib/jxl/testing.h"
 
 namespace jxl {
@@ -26,7 +27,7 @@ TEST(CodecPGXTest, Test8bits) {
   ThreadPool* pool = nullptr;
 
   EXPECT_TRUE(DecodeImagePGX(MakeSpan(pgx.c_str()), ColorHints(), &ppf));
-  CodecInOut io;
+  CodecInOut io{jxl::test::MemoryManager()};
   EXPECT_TRUE(ConvertPackedPixelFileToCodecInOut(ppf, pool, &io));
 
   ScaleImage(255.f, io.Main().color());
@@ -53,7 +54,7 @@ TEST(CodecPGXTest, Test16bits) {
   ThreadPool* pool = nullptr;
 
   EXPECT_TRUE(DecodeImagePGX(MakeSpan(pgx.c_str()), ColorHints(), &ppf));
-  CodecInOut io;
+  CodecInOut io{jxl::test::MemoryManager()};
   EXPECT_TRUE(ConvertPackedPixelFileToCodecInOut(ppf, pool, &io));
 
   ScaleImage(255.f, io.Main().color());

third_party/jpeg-xl/lib/extras/dec/pnm.cc

@@ -15,6 +15,7 @@
 
 #include "lib/extras/size_constraints.h"
 #include "lib/jxl/base/bits.h"
+#include "lib/jxl/base/c_callback_support.h"
 #include "lib/jxl/base/span.h"
 #include "lib/jxl/base/status.h"
 

third_party/jpeg-xl/lib/extras/enc/jpg.cc

@@ -21,8 +21,8 @@
 
 #include "lib/extras/exif.h"
 #include "lib/jxl/base/common.h"
+#include "lib/jxl/base/sanitizers.h"
 #include "lib/jxl/base/status.h"
-#include "lib/jxl/sanitizers.h"
 #if JPEGXL_ENABLE_SJPEG
 #include "sjpeg.h"
 #include "sjpegi.h"

third_party/jpeg-xl/lib/extras/enc/jxl.cc

@@ -202,7 +202,7 @@ bool EncodeImageJXL(const JXLCompressParams& params, const PackedPixelFile& ppf,
         fprintf(stderr,
                 "JPEG bitstream reconstruction data could not be created. "
                 "Possibly there is too much tail data.\n"
-                "Try using --jpeg_store_metadata 0, to losslessly "
+                "Try using --allow_jpeg_reconstruction 0, to losslessly "
                 "recompress the JPEG image data without bitstream "
                 "reconstruction data.\n");
       } else {
@@ -223,7 +223,14 @@ bool EncodeImageJXL(const JXLCompressParams& params, const PackedPixelFile& ppf,
         std::max<uint32_t>(num_alpha_channels, ppf.info.num_extra_channels);
     basic_info.num_color_channels = ppf.info.num_color_channels;
     const bool lossless = (params.distance == 0);
-    basic_info.uses_original_profile = TO_JXL_BOOL(lossless);
+    auto non_perceptual_option = std::find_if(
+        params.options.begin(), params.options.end(), [](JXLOption option) {
+          return option.id ==
+                 JXL_ENC_FRAME_SETTING_DISABLE_PERCEPTUAL_HEURISTICS;
+        });
+    const bool non_perceptual = non_perceptual_option != params.options.end() &&
+                                non_perceptual_option->ival == 1;
+    basic_info.uses_original_profile = TO_JXL_BOOL(lossless || non_perceptual);
     if (params.override_bitdepth != 0) {
       basic_info.bits_per_sample = params.override_bitdepth;
       basic_info.exponent_bits_per_sample =

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

@@ -12,22 +12,19 @@
 #include <jxl/codestream_header.h>
 #include <jxl/encode.h>
 #include <jxl/types.h>
-#include <stddef.h>
-#include <stdint.h>
-#include <stdlib.h>
-#include <string.h>
 
 #include <algorithm>
 #include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
 #include <functional>
 #include <memory>
-#include <mutex>
-#include <set>
 #include <string>
 #include <vector>
 
 #include "lib/jxl/base/byte_order.h"
-#include "lib/jxl/base/c_callback_support.h"
 #include "lib/jxl/base/common.h"
 #include "lib/jxl/base/status.h"
 

third_party/jpeg-xl/lib/extras/packed_image_convert.cc

@@ -7,6 +7,7 @@
 
 #include <jxl/cms.h>
 #include <jxl/color_encoding.h>
+#include <jxl/memory_manager.h>
 #include <jxl/types.h>
 
 #include <cstdint>
@@ -27,6 +28,7 @@ Status ConvertPackedFrameToImageBundle(const JxlBasicInfo& info,
                                        const PackedFrame& frame,
                                        const CodecInOut& io, ThreadPool* pool,
                                        ImageBundle* bundle) {
+  JxlMemoryManager* memory_manager = io.memory_manager;
   JXL_ASSERT(frame.color.pixels() != nullptr);
   size_t frame_bits_per_sample =
       input_bitdepth.type == JXL_BIT_DEPTH_FROM_PIXEL_FORMAT
@@ -65,8 +67,9 @@ Status ConvertPackedFrameToImageBundle(const JxlBasicInfo& info,
   bundle->extra_channels().resize(io.metadata.m.extra_channel_info.size());
   for (size_t i = 0; i < frame.extra_channels.size(); i++) {
     const auto& ppf_ec = frame.extra_channels[i];
-    JXL_ASSIGN_OR_RETURN(bundle->extra_channels()[i],
+    JXL_ASSIGN_OR_RETURN(
-                         ImageF::Create(ppf_ec.xsize, ppf_ec.ysize));
+        bundle->extra_channels()[i],
+        ImageF::Create(memory_manager, ppf_ec.xsize, ppf_ec.ysize));
     JXL_CHECK(BufferToImageF(ppf_ec.format, ppf_ec.xsize, ppf_ec.ysize,
                              ppf_ec.pixels(), ppf_ec.pixels_size, pool,
                              &bundle->extra_channels()[i]));
@@ -76,6 +79,7 @@ Status ConvertPackedFrameToImageBundle(const JxlBasicInfo& info,
 
 Status ConvertPackedPixelFileToCodecInOut(const PackedPixelFile& ppf,
                                           ThreadPool* pool, CodecInOut* io) {
+  JxlMemoryManager* memory_manager = io->memory_manager;
   const bool has_alpha = ppf.info.alpha_bits != 0;
   JXL_ASSERT(!ppf.frames.empty());
   if (has_alpha) {
@@ -179,7 +183,7 @@ Status ConvertPackedPixelFileToCodecInOut(const PackedPixelFile& ppf,
   // Convert the pixels
   io->frames.clear();
   for (const auto& frame : ppf.frames) {
-    ImageBundle bundle(&io->metadata.m);
+    ImageBundle bundle(memory_manager, &io->metadata.m);
     JXL_RETURN_IF_ERROR(ConvertPackedFrameToImageBundle(
         ppf.info, ppf.input_bitdepth, frame, *io, pool, &bundle));
     io->frames.push_back(std::move(bundle));
@@ -234,6 +238,7 @@ Status ConvertCodecInOutToPackedPixelFile(const CodecInOut& io,
                                           const ColorEncoding& c_desired,
                                           ThreadPool* pool,
                                           PackedPixelFile* ppf) {
+  JxlMemoryManager* memory_manager = io.memory_manager;
   const bool has_alpha = io.metadata.m.HasAlpha();
   JXL_ASSERT(!io.frames.empty());
 
@@ -317,7 +322,7 @@ Status ConvertCodecInOutToPackedPixelFile(const CodecInOut& io,
     JXL_ASSIGN_OR_RETURN(ImageBundle ib, frame.Copy());
     const ImageBundle* to_color_transform = &ib;
     ImageMetadata metadata = io.metadata.m;
-    ImageBundle store(&metadata);
+    ImageBundle store(memory_manager, &metadata);
     const ImageBundle* transformed;
     // TODO(firsching): handle the transform here.
     JXL_RETURN_IF_ERROR(TransformIfNeeded(*to_color_transform, c_desired,

third_party/jpeg-xl/lib/extras/tone_mapping_gbench.cc

@@ -3,15 +3,18 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+#include <jxl/memory_manager.h>
+
 #include "benchmark/benchmark.h"
-#include "lib/extras/codec.h"
 #include "lib/extras/tone_mapping.h"
 #include "lib/jxl/image.h"
+#include "tools/no_memory_manager.h"
 
 namespace jxl {
 
 static void BM_ToneMapping(benchmark::State& state) {
-  JXL_ASSIGN_OR_DIE(Image3F color, Image3F::Create(2268, 1512));
+  JxlMemoryManager* memory_manager = jpegxl::tools::NoMemoryManager();
+  JXL_ASSIGN_OR_DIE(Image3F color, Image3F::Create(memory_manager, 2268, 1512));
   FillImage(0.5f, &color);
 
   // Use linear Rec. 2020 so that `ToneMapTo` doesn't have to convert to it and
@@ -25,9 +28,10 @@ static void BM_ToneMapping(benchmark::State& state) {
 
   for (auto _ : state) {
     state.PauseTiming();
-    CodecInOut tone_mapping_input;
+    CodecInOut tone_mapping_input{memory_manager};
-    JXL_ASSIGN_OR_DIE(Image3F color2,
+    JXL_ASSIGN_OR_DIE(
-                      Image3F::Create(color.xsize(), color.ysize()));
+        Image3F color2,
+        Image3F::Create(memory_manager, color.xsize(), color.ysize()));
     CopyImageTo(color, &color2);
     tone_mapping_input.SetFromImage(std::move(color2), linear_rec2020);
     tone_mapping_input.metadata.m.SetIntensityTarget(255);

third_party/jpeg-xl/lib/include/jxl/color_encoding.h

@@ -32,9 +32,9 @@ typedef enum {
   JXL_COLOR_SPACE_UNKNOWN,
 } JxlColorSpace;
 
-/** Built-in whitepoints for color encoding. When decoding, the numerical xy
+/** Built-in white points for color encoding. When decoding, the numerical xy
- * whitepoint value can be read from the @ref JxlColorEncoding white_point field
+ * white point value can be read from the @ref JxlColorEncoding white_point
- * regardless of the enum value. When encoding, enum values except
+ * field regardless of the enum value. When encoding, enum values except
  * ::JXL_WHITE_POINT_CUSTOM override the numerical fields. Some enum values
  * match a subset of CICP (Rec. ITU-T H.273 | ISO/IEC 23091-2:2019(E)), however
  * the white point and RGB primaries are separate enums here.
@@ -78,7 +78,7 @@ typedef enum {
  * of CICP (Rec. ITU-T H.273 | ISO/IEC 23091-2:2019(E)) unless specified
  * otherwise. */
 typedef enum {
-  /** As specified in SMPTE RP 431-2 */
+  /** As specified in ITU-R BT.709-6 */
   JXL_TRANSFER_FUNCTION_709 = 1,
   /** None of the other table entries describe the transfer function. */
   JXL_TRANSFER_FUNCTION_UNKNOWN = 2,
@@ -97,7 +97,7 @@ typedef enum {
   JXL_TRANSFER_FUNCTION_GAMMA = 65535,
 } JxlTransferFunction;
 
-/** Renderig intent for color encoding, as specified in ISO 15076-1:2010 */
+/** Rendering intent for color encoding, as specified in ISO 15076-1:2010 */
 typedef enum {
   /** vendor-specific */
   JXL_RENDERING_INTENT_PERCEPTUAL = 0,
@@ -117,7 +117,7 @@ typedef struct {
   JxlColorSpace color_space;
 
   /** Built-in white point. If this value is ::JXL_WHITE_POINT_CUSTOM, must
-   * use the numerical whitepoint values from white_point_xy.
+   * use the numerical white point values from white_point_xy.
    */
   JxlWhitePoint white_point;
 

third_party/jpeg-xl/lib/include/jxl/decode.h

@@ -721,7 +721,7 @@ typedef enum {
  * It is often possible to use @ref JxlDecoderGetColorAsICCProfile as an
  * alternative anyway. The following scenarios are possible:
  *  - The JPEG XL image has an attached ICC Profile, in that case, the encoded
- *    structured data is not available, this function will return an error
+ *    structured data is not available and this function will return an error
  *    status. @ref JxlDecoderGetColorAsICCProfile should be called instead.
  *  - The JPEG XL image has an encoded structured color profile, and it
  *    represents an RGB or grayscale color space. This function will return it.
@@ -800,8 +800,8 @@ JXL_EXPORT JxlDecoderStatus JxlDecoderGetICCProfileSize(
  *     or the color profile of the decoded pixels.
  * @param icc_profile buffer to copy the ICC profile into
  * @param size size of the icc_profile buffer in bytes
- * @return ::JXL_DEC_SUCCESS if the profile was successfully returned is
+ * @return ::JXL_DEC_SUCCESS if the profile was successfully returned,
- *     available, ::JXL_DEC_NEED_MORE_INPUT if not yet available, @ref
+ *     ::JXL_DEC_NEED_MORE_INPUT if not yet available, @ref
  *     JXL_DEC_ERROR if the profile doesn't exist or the output size is not
  *     large enough.
  */
@@ -869,7 +869,7 @@ JXL_EXPORT JxlDecoderStatus JxlDecoderSetDesiredIntensityTarget(
  *
  * This function must not be called before @ref JxlDecoderSetCms.
  *
- * @param dec decoder orbject
+ * @param dec decoder object
  * @param color_encoding the output color encoding
  * @param icc_data bytes of the icc profile
  * @param icc_size size of the icc profile in bytes
@@ -913,7 +913,7 @@ JXL_EXPORT JxlDecoderStatus JxlDecoderPreviewOutBufferSize(
     const JxlDecoder* dec, const JxlPixelFormat* format, size_t* size);
 
 /**
- * Sets the buffer to write the small resolution preview image
+ * Sets the buffer to write the low-resolution preview image
  * to. The size of the buffer must be at least as large as given by @ref
  * JxlDecoderPreviewOutBufferSize. The buffer follows the format described
  * by @ref JxlPixelFormat. The preview image dimensions are given by the
@@ -962,10 +962,10 @@ JXL_EXPORT JxlDecoderStatus JxlDecoderGetFrameName(const JxlDecoder* dec,
 
 /**
  * Outputs the blend information for the current frame for a specific extra
- * channel. This function can be called when ::JXL_DEC_FRAME occurred for the
+ * channel. This function can be called once the ::JXL_DEC_FRAME event occurred
- * current frame, even when have_animation in the @ref JxlBasicInfo is @ref
+ * for the current frame, even if the `have_animation` field in the @ref
- * JXL_FALSE. This information is only useful if coalescing is disabled;
+ * JxlBasicInfo is @ref JXL_FALSE. This information is only useful if coalescing
- * otherwise the decoder will have performed blending already.
+ * is disabled; otherwise the decoder will have performed blending already.
  *
  * @param dec decoder object
  * @param index the index of the extra channel
@@ -1344,7 +1344,7 @@ JXL_EXPORT JxlDecoderStatus JxlDecoderSetDecompressBoxes(JxlDecoder* dec,
  *    animation allowing the decoder to jump to individual frames more
  *    efficiently.
  *  - "jbrd": JPEG reconstruction box, contains the information required to
- *    byte-for-byte losslessly recontruct a JPEG-1 image. The JPEG DCT
+ *    byte-for-byte losslessly reconstruct a JPEG-1 image. The JPEG DCT
  *    coefficients (pixel content) themselves as well as the ICC profile are
  *    encoded in the JXL codestream (jxlc or jxlp) itself. EXIF, XMP and JUMBF
  *    metadata is encoded in the corresponding boxes. The jbrd box itself
@@ -1366,7 +1366,7 @@ JXL_EXPORT JxlDecoderStatus JxlDecoderSetDecompressBoxes(JxlDecoder* dec,
  * @param decompressed which box type to get: JXL_FALSE to get the raw box type,
  *     which can be "brob", JXL_TRUE, get the underlying box type.
  * @return ::JXL_DEC_SUCCESS if the value is available, ::JXL_DEC_ERROR if
- *     not, for example the JXL file does not use the container format.
+ *     not, for example the JPEG XL file does not use the container format.
  */
 JXL_EXPORT JxlDecoderStatus JxlDecoderGetBoxType(JxlDecoder* dec,
                                                  JxlBoxType type,

third_party/jpeg-xl/lib/include/jxl/encode.h

@@ -388,6 +388,11 @@ typedef enum {
    */
   JXL_ENC_FRAME_SETTING_USE_FULL_IMAGE_HEURISTICS = 38,
 
+  /** Disable perceptual optimizations. 0 = optimizations enabled (default), 1 =
+   * optimizations disabled.
+   */
+  JXL_ENC_FRAME_SETTING_DISABLE_PERCEPTUAL_HEURISTICS = 39,
+
   /** Enum value not to be used as an option. This value is added to force the
    * C compiler to have the enum to take a known size.
    */

third_party/jpeg-xl/lib/jpegli/libjpeg_test_util.cc

@@ -5,8 +5,11 @@
 
 #include "lib/jpegli/libjpeg_test_util.h"
 
+#include <cstring>
+
 #include "lib/jxl/base/include_jpeglib.h"  // NOLINT
-#include "lib/jxl/sanitizers.h"
+#include "lib/jxl/base/sanitizers.h"
+#include "lib/jxl/base/status.h"
 
 namespace jpegli {
 

third_party/jpeg-xl/lib/jpegli/test_utils.cc

@@ -7,6 +7,7 @@
 
 #include <cmath>
 #include <cstdint>
+#include <cstring>
 #include <fstream>
 #include <sstream>
 
@@ -14,8 +15,8 @@
 #include "lib/jpegli/encode.h"
 #include "lib/jxl/base/byte_order.h"
 #include "lib/jxl/base/printf_macros.h"
+#include "lib/jxl/base/sanitizers.h"
 #include "lib/jxl/base/status.h"
-#include "lib/jxl/sanitizers.h"
 
 #if !defined(TEST_DATA_PATH)
 #include "tools/cpp/runfiles/runfiles.h"

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

@@ -5,9 +5,10 @@
 
 #include "lib/jxl/ac_strategy.h"
 
-#include <string.h>
+#include <jxl/memory_manager.h>
 
 #include <algorithm>
+#include <cstring>
 #include <utility>
 
 #include "lib/jxl/base/bits.h"
@@ -83,9 +84,11 @@ constexpr size_t AcStrategy::kMaxCoeffBlocks;
 constexpr size_t AcStrategy::kMaxBlockDim;
 constexpr size_t AcStrategy::kMaxCoeffArea;
 
-StatusOr<AcStrategyImage> AcStrategyImage::Create(size_t xsize, size_t ysize) {
+StatusOr<AcStrategyImage> AcStrategyImage::Create(
+    JxlMemoryManager* memory_manager, size_t xsize, size_t ysize) {
   AcStrategyImage img;
-  JXL_ASSIGN_OR_RETURN(img.layers_, ImageB::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(img.layers_,
+                       ImageB::Create(memory_manager, xsize, ysize));
   img.row_ = img.layers_.Row(0);
   img.stride_ = img.layers_.PixelsPerRow();
   return img;

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

@@ -6,9 +6,10 @@
 #ifndef LIB_JXL_AC_STRATEGY_H_
 #define LIB_JXL_AC_STRATEGY_H_
 
-#include <stddef.h>
+#include <jxl/memory_manager.h>
-#include <stdint.h>
 
+#include <cstddef>
+#include <cstdint>
 #include <hwy/base.h>  // kMaxVectorSize
 
 #include "lib/jxl/base/compiler_specific.h"
@@ -196,7 +197,8 @@ class AcStrategyRow {
 class AcStrategyImage {
  public:
   AcStrategyImage() = default;
-  static StatusOr<AcStrategyImage> Create(size_t xsize, size_t ysize);
+  static StatusOr<AcStrategyImage> Create(JxlMemoryManager* memory_manager,
+                                          size_t xsize, size_t ysize);
 
   AcStrategyImage(AcStrategyImage&&) = default;
   AcStrategyImage& operator=(AcStrategyImage&&) = default;
@@ -252,6 +254,8 @@ class AcStrategyImage {
   // Count the number of blocks of a given type.
   size_t CountBlocks(AcStrategy::Type type) const;
 
+  JxlMemoryManager* memory_manager() const { return layers_.memory_manager(); }
+
  private:
   ImageB layers_;
   uint8_t* JXL_RESTRICT row_;

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

@@ -3,9 +3,10 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-#include <stddef.h>
+#include <jxl/memory_manager.h>
-#include <stdint.h>
 
+#include <cstddef>
+#include <cstdint>
 #include <vector>
 
 #include "lib/jxl/ans_params.h"
@@ -15,6 +16,7 @@
 #include "lib/jxl/dec_bit_reader.h"
 #include "lib/jxl/enc_ans.h"
 #include "lib/jxl/enc_bit_writer.h"
+#include "lib/jxl/test_utils.h"
 #include "lib/jxl/testing.h"
 
 namespace jxl {
@@ -22,10 +24,11 @@ namespace {
 
 void RoundtripTestcase(int n_histograms, int alphabet_size,
                        const std::vector<Token>& input_values) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   constexpr uint16_t kMagic1 = 0x9e33;
   constexpr uint16_t kMagic2 = 0x8b04;
 
-  BitWriter writer;
+  BitWriter writer{memory_manager};
   // Space for magic bytes.
   BitWriter::Allotment allotment_magic1(&writer, 16);
   writer.Write(16, kMagic1);
@@ -36,8 +39,9 @@ void RoundtripTestcase(int n_histograms, int alphabet_size,
   std::vector<std::vector<Token>> input_values_vec;
   input_values_vec.push_back(input_values);
 
-  BuildAndEncodeHistograms(HistogramParams(), n_histograms, input_values_vec,
+  BuildAndEncodeHistograms(memory_manager, HistogramParams(), n_histograms,
-                           &codes, &context_map, &writer, 0, nullptr);
+                           input_values_vec, &codes, &context_map, &writer, 0,
+                           nullptr);
   WriteTokens(input_values_vec[0], codes, context_map, 0, &writer, 0, nullptr);
 
   // Magic bytes + padding
@@ -54,10 +58,11 @@ void RoundtripTestcase(int n_histograms, int alphabet_size,
 
   std::vector<uint8_t> dec_context_map;
   ANSCode decoded_codes;
-  ASSERT_TRUE(
+  ASSERT_TRUE(DecodeHistograms(memory_manager, &br, n_histograms,
-      DecodeHistograms(&br, n_histograms, &decoded_codes, &dec_context_map));
+                               &decoded_codes, &dec_context_map));
   ASSERT_EQ(dec_context_map, context_map);
-  ANSSymbolReader reader(&decoded_codes, &br);
+  JXL_ASSIGN_OR_DIE(ANSSymbolReader reader,
+                    ANSSymbolReader::Create(&decoded_codes, &br));
 
   for (const Token& symbol : input_values) {
     uint32_t read_symbol =
@@ -156,6 +161,7 @@ TEST(ANSTest, RandomUnbalancedStreamRoundtripBig) {
 }
 
 TEST(ANSTest, UintConfigRoundtrip) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   for (size_t log_alpha_size = 5; log_alpha_size <= 8; log_alpha_size++) {
     std::vector<HybridUintConfig> uint_config;
     std::vector<HybridUintConfig> uint_config_dec;
@@ -168,7 +174,7 @@ TEST(ANSTest, UintConfigRoundtrip) {
     }
     uint_config.emplace_back(log_alpha_size, 0, 0);
     uint_config_dec.resize(uint_config.size());
-    BitWriter writer;
+    BitWriter writer{memory_manager};
     BitWriter::Allotment allotment(&writer, 10 * uint_config.size());
     EncodeUintConfigs(uint_config, &writer, log_alpha_size);
     allotment.ReclaimAndCharge(&writer, 0, nullptr);
@@ -185,6 +191,7 @@ TEST(ANSTest, UintConfigRoundtrip) {
 }
 
 void TestCheckpointing(bool ans, bool lz77) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   std::vector<std::vector<Token>> input_values(1);
   for (size_t i = 0; i < 1024; i++) {
     input_values[0].emplace_back(0, i % 4);
@@ -206,11 +213,11 @@ void TestCheckpointing(bool ans, bool lz77) {
                             : HistogramParams::LZ77Method::kNone;
   params.force_huffman = !ans;
 
-  BitWriter writer;
+  BitWriter writer{memory_manager};
   {
     auto input_values_copy = input_values;
-    BuildAndEncodeHistograms(params, 1, input_values_copy, &codes, &context_map,
+    BuildAndEncodeHistograms(memory_manager, params, 1, input_values_copy,
-                             &writer, 0, nullptr);
+                             &codes, &context_map, &writer, 0, nullptr);
     WriteTokens(input_values_copy[0], codes, context_map, 0, &writer, 0,
                 nullptr);
     writer.ZeroPadToByte();
@@ -225,9 +232,11 @@ void TestCheckpointing(bool ans, bool lz77) {
 
     std::vector<uint8_t> dec_context_map;
     ANSCode decoded_codes;
-    ASSERT_TRUE(DecodeHistograms(&br, 1, &decoded_codes, &dec_context_map));
+    ASSERT_TRUE(DecodeHistograms(memory_manager, &br, 1, &decoded_codes,
+                                 &dec_context_map));
     ASSERT_EQ(dec_context_map, context_map);
-    ANSSymbolReader reader(&decoded_codes, &br);
+    JXL_ASSIGN_OR_DIE(ANSSymbolReader reader,
+                      ANSSymbolReader::Create(&decoded_codes, &br));
 
     ANSSymbolReader::Checkpoint checkpoint;
     size_t br_pos = 0;

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

@@ -115,6 +115,11 @@ class RectT {
            y1() <= other.y1();
   }
 
+  bool IsSame(const RectT& other) const {
+    return x0_ == other.x0_ && xsize_ == other.xsize_ && y0_ == other.y0_ &&
+           ysize_ <= other.ysize_;
+  }
+
   // Returns true if this Rect fully resides in the given image. ImageT could be
   // Plane<T> or Image3<T>; however if ImageT is Rect, results are nonsensical.
   template <class ImageT>

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

@@ -10,7 +10,6 @@
 
 #include "lib/jxl/base/compiler_specific.h"
 #include "lib/jxl/base/sanitizer_definitions.h"
-#include "lib/jxl/image.h"
 
 #if JXL_MEMORY_SANITIZER
 #include <algorithm>
@@ -49,17 +48,17 @@ static JXL_INLINE JXL_MAYBE_UNUSED void MemoryIsInitialized(
 }
 
 // Mark all the bytes of an image (including padding) as poisoned bytes.
-template <typename T>
+template <typename Pixels>
-static JXL_INLINE JXL_MAYBE_UNUSED void PoisonImage(const Plane<T>& im) {
+static JXL_INLINE JXL_MAYBE_UNUSED void PoisonImage(const Pixels& im) {
   PoisonMemory(im.bytes(), im.bytes_per_row() * im.ysize());
 }
 
 namespace {
 
 // Print the uninitialized regions of an image.
-template <typename T>
+template <typename Pixels>
 static JXL_INLINE JXL_MAYBE_UNUSED void PrintImageUninitialized(
-    const Plane<T>& im) {
+    const Pixels& im) {
   fprintf(stderr,
           "Uninitialized regions for image of size %" PRIu64 "x%" PRIu64 ":\n",
           static_cast<uint64_t>(im.xsize()), static_cast<uint64_t>(im.ysize()));
@@ -157,9 +156,9 @@ static JXL_INLINE JXL_MAYBE_UNUSED void PrintImageUninitialized(
 
 // Check that all the pixels in the provided rect of the image are initialized
 // (not poisoned). If any of the values is poisoned it will abort.
-template <typename T>
+template <typename Pixels>
 static JXL_INLINE JXL_MAYBE_UNUSED void CheckImageInitialized(
-    const Plane<T>& im, const Rect& r, size_t c, const char* message) {
+    const Pixels& im, const Rect& r, size_t c, const char* message) {
   JXL_ASSERT(r.x0() <= im.xsize());
   JXL_ASSERT(r.x0() + r.xsize() <= im.xsize());
   JXL_ASSERT(r.y0() <= im.ysize());
@@ -190,9 +189,9 @@ static JXL_INLINE JXL_MAYBE_UNUSED void CheckImageInitialized(
   }
 }
 
-template <typename T>
+template <typename Image>
 static JXL_INLINE JXL_MAYBE_UNUSED void CheckImageInitialized(
-    const Image3<T>& im, const Rect& r, const char* message) {
+    const Image& im, const Rect& r, const char* message) {
   for (size_t c = 0; c < 3; c++) {
     std::string str_message(message);
     str_message += " c=" + std::to_string(c);
@@ -220,8 +219,8 @@ static JXL_INLINE JXL_MAYBE_UNUSED void UnpoisonMemory(const void* m,
 static JXL_INLINE JXL_MAYBE_UNUSED void MemoryIsInitialized(const void* m,
                                                             size_t size) {}
 
-template <typename T>
+template <typename Pixels>
-static JXL_INLINE JXL_MAYBE_UNUSED void PoisonImage(const Plane<T>& im) {}
+static JXL_INLINE JXL_MAYBE_UNUSED void PoisonImage(const Pixels& im) {}
 
 #define JXL_CHECK_IMAGE_INITIALIZED(im, r)
 #define JXL_CHECK_PLANE_INITIALIZED(im, r, c)

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

@@ -3,10 +3,9 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-#include <stddef.h>
-#include <stdint.h>
-
 #include <array>
+#include <cstddef>
+#include <cstdint>
 #include <utility>
 #include <vector>
 
@@ -52,12 +51,13 @@ struct Symbol {
 
 // Reading from output gives the same values.
 TEST(BitReaderTest, TestRoundTrip) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   test::ThreadPoolForTests pool(8);
   EXPECT_TRUE(RunOnPool(
       pool.get(), 0, 1000, ThreadPool::NoInit,
-      [](const uint32_t task, size_t /* thread */) {
+      [&memory_manager](const uint32_t task, size_t /* thread */) {
         constexpr size_t kMaxBits = 8000;
-        BitWriter writer;
+        BitWriter writer{memory_manager};
         BitWriter::Allotment allotment(&writer, kMaxBits);
 
         std::vector<Symbol> symbols;
@@ -86,14 +86,15 @@ TEST(BitReaderTest, TestRoundTrip) {
 
 // SkipBits is the same as reading that many bits.
 TEST(BitReaderTest, TestSkip) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   test::ThreadPoolForTests pool(8);
   EXPECT_TRUE(RunOnPool(
       pool.get(), 0, 96, ThreadPool::NoInit,
-      [](const uint32_t task, size_t /* thread */) {
+      [&memory_manager](const uint32_t task, size_t /* thread */) {
         constexpr size_t kSize = 100;
 
         for (size_t skip = 0; skip < 128; ++skip) {
-          BitWriter writer;
+          BitWriter writer{memory_manager};
           BitWriter::Allotment allotment(&writer, kSize * kBitsPerByte);
           // Start with "task" 1-bits.
           for (size_t i = 0; i < task; ++i) {
@@ -142,11 +143,12 @@ TEST(BitReaderTest, TestSkip) {
 
 // Verifies byte order and different groupings of bits.
 TEST(BitReaderTest, TestOrder) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   constexpr size_t kMaxBits = 16;
 
   // u(1) - bits written into LSBs of first byte
   {
-    BitWriter writer;
+    BitWriter writer{memory_manager};
     BitWriter::Allotment allotment(&writer, kMaxBits);
     for (size_t i = 0; i < 5; ++i) {
       writer.Write(1, 1);
@@ -168,7 +170,7 @@ TEST(BitReaderTest, TestOrder) {
 
   // u(8) - get bytes in the same order
   {
-    BitWriter writer;
+    BitWriter writer{memory_manager};
     BitWriter::Allotment allotment(&writer, kMaxBits);
     writer.Write(8, 0xF8);
     writer.Write(8, 0x3F);
@@ -183,7 +185,7 @@ TEST(BitReaderTest, TestOrder) {
 
   // u(16) - little-endian bytes
   {
-    BitWriter writer;
+    BitWriter writer{memory_manager};
     BitWriter::Allotment allotment(&writer, kMaxBits);
     writer.Write(16, 0xF83F);
 
@@ -197,7 +199,7 @@ TEST(BitReaderTest, TestOrder) {
 
   // Non-byte-aligned, mixed sizes
   {
-    BitWriter writer;
+    BitWriter writer{memory_manager};
     BitWriter::Allotment allotment(&writer, kMaxBits);
     writer.Write(1, 1);
     writer.Write(3, 6);

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

@@ -5,6 +5,8 @@
 
 #include "lib/jxl/blending.h"
 
+#include <jxl/memory_manager.h>
+
 #include <cstddef>
 #include <cstring>
 #include <vector>
@@ -38,9 +40,9 @@ bool NeedsBlending(const FrameHeader& frame_header) {
 }
 
 Status PerformBlending(
-    const float* const* bg, const float* const* fg, float* const* out,
+    JxlMemoryManager* memory_manager, const float* const* bg,
-    size_t x0, size_t xsize, const PatchBlending& color_blending,
+    const float* const* fg, float* const* out, size_t x0, size_t xsize,
-    const PatchBlending* ec_blending,
+    const PatchBlending& color_blending, const PatchBlending* ec_blending,
     const std::vector<ExtraChannelInfo>& extra_channel_info) {
   bool has_alpha = false;
   size_t num_ec = extra_channel_info.size();
@@ -50,7 +52,8 @@ Status PerformBlending(
       break;
     }
   }
-  JXL_ASSIGN_OR_RETURN(ImageF tmp, ImageF::Create(xsize, 3 + num_ec));
+  JXL_ASSIGN_OR_RETURN(ImageF tmp,
+                       ImageF::Create(memory_manager, xsize, 3 + num_ec));
   // Blend extra channels first so that we use the pre-blending alpha.
   for (size_t i = 0; i < num_ec; i++) {
     if (ec_blending[i].mode == PatchBlendMode::kAdd) {

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

@@ -6,6 +6,8 @@
 #ifndef LIB_JXL_BLENDING_H_
 #define LIB_JXL_BLENDING_H_
 
+#include <jxl/memory_manager.h>
+
 #include <cstddef>
 #include <vector>
 
@@ -18,9 +20,9 @@ namespace jxl {
 
 bool NeedsBlending(const FrameHeader& frame_header);
 
-Status PerformBlending(const float* const* bg, const float* const* fg,
+Status PerformBlending(JxlMemoryManager* memory_manager, const float* const* bg,
-                       float* const* out, size_t x0, size_t xsize,
+                       const float* const* fg, float* const* out, size_t x0,
-                       const PatchBlending& color_blending,
+                       size_t xsize, const PatchBlending& color_blending,
                        const PatchBlending* ec_blending,
                        const std::vector<ExtraChannelInfo>& extra_channel_info);
 

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

@@ -5,7 +5,12 @@
 
 #include "lib/jxl/box_content_decoder.h"
 
-#include "lib/jxl/sanitizers.h"
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+
+#include "lib/jxl/base/sanitizers.h"
 
 namespace jxl {
 

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

@@ -22,6 +22,8 @@
 
 #include "lib/jxl/butteraugli/butteraugli.h"
 
+#include <jxl/memory_manager.h>
+
 #include <algorithm>
 #include <cmath>
 #include <cstdint>
@@ -409,8 +411,9 @@ Status SeparateMFAndHF(const ButteraugliParams& params, Image3F* mf, ImageF* hf,
   static const double kSigmaHf = 3.22489901262;
   const size_t xsize = mf->xsize();
   const size_t ysize = mf->ysize();
-  JXL_ASSIGN_OR_RETURN(hf[0], ImageF::Create(xsize, ysize));
+  JxlMemoryManager* memory_manager = mf[0].memory_manager();
-  JXL_ASSIGN_OR_RETURN(hf[1], ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(hf[0], ImageF::Create(memory_manager, xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(hf[1], ImageF::Create(memory_manager, xsize, ysize));
   for (int i = 0; i < 3; ++i) {
     if (i == 2) {
       JXL_RETURN_IF_ERROR(
@@ -465,9 +468,10 @@ Status SeparateHFAndUHF(const ButteraugliParams& params, ImageF* hf,
   const HWY_FULL(float) d;
   const size_t xsize = hf[0].xsize();
   const size_t ysize = hf[0].ysize();
+  JxlMemoryManager* memory_manager = hf[0].memory_manager();
   static const double kSigmaUhf = 1.56416327805;
-  JXL_ASSIGN_OR_RETURN(uhf[0], ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(uhf[0], ImageF::Create(memory_manager, xsize, ysize));
-  JXL_ASSIGN_OR_RETURN(uhf[1], ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(uhf[1], ImageF::Create(memory_manager, xsize, ysize));
   for (int i = 0; i < 2; ++i) {
     // Divide hf into hf and uhf.
     for (size_t y = 0; y < ysize; ++y) {
@@ -531,8 +535,11 @@ void DeallocateHFAndUHF(ImageF* hf, ImageF* uhf) {
 Status SeparateFrequencies(size_t xsize, size_t ysize,
                            const ButteraugliParams& params, BlurTemp* blur_temp,
                            const Image3F& xyb, PsychoImage& ps) {
-  JXL_ASSIGN_OR_RETURN(ps.lf, Image3F::Create(xyb.xsize(), xyb.ysize()));
+  JxlMemoryManager* memory_manager = xyb.memory_manager();
-  JXL_ASSIGN_OR_RETURN(ps.mf, Image3F::Create(xyb.xsize(), xyb.ysize()));
+  JXL_ASSIGN_OR_RETURN(
+      ps.lf, Image3F::Create(memory_manager, xyb.xsize(), xyb.ysize()));
+  JXL_ASSIGN_OR_RETURN(
+      ps.mf, Image3F::Create(memory_manager, xyb.xsize(), xyb.ysize()));
   JXL_RETURN_IF_ERROR(SeparateLFAndMF(params, xyb, &ps.lf, &ps.mf, blur_temp));
   JXL_RETURN_IF_ERROR(SeparateMFAndHF(params, &ps.mf, &ps.hf[0], blur_temp));
   JXL_RETURN_IF_ERROR(
@@ -1204,14 +1211,19 @@ Status Mask(const ImageF& mask0, const ImageF& mask1,
             ImageF* BUTTERAUGLI_RESTRICT diff_ac) {
   const size_t xsize = mask0.xsize();
   const size_t ysize = mask0.ysize();
-  JXL_ASSIGN_OR_RETURN(*mask, ImageF::Create(xsize, ysize));
+  JxlMemoryManager* memory_manager = mask0.memory_manager();
+  JXL_ASSIGN_OR_RETURN(*mask, ImageF::Create(memory_manager, xsize, ysize));
   static const float kMul = 6.19424080439;
   static const float kBias = 12.61050594197;
   static const float kRadius = 2.7;
-  JXL_ASSIGN_OR_RETURN(ImageF diff0, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF diff0,
-  JXL_ASSIGN_OR_RETURN(ImageF diff1, ImageF::Create(xsize, ysize));
+                       ImageF::Create(memory_manager, xsize, ysize));
-  JXL_ASSIGN_OR_RETURN(ImageF blurred0, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF diff1,
-  JXL_ASSIGN_OR_RETURN(ImageF blurred1, ImageF::Create(xsize, ysize));
+                       ImageF::Create(memory_manager, xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF blurred0,
+                       ImageF::Create(memory_manager, xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF blurred1,
+                       ImageF::Create(memory_manager, xsize, ysize));
   DiffPrecompute(mask0, kMul, kBias, &diff0);
   DiffPrecompute(mask1, kMul, kBias, &diff1);
   JXL_RETURN_IF_ERROR(Blur(diff0, kRadius, params, blur_temp, &blurred0));
@@ -1236,8 +1248,11 @@ Status MaskPsychoImage(const PsychoImage& pi0, const PsychoImage& pi1,
                        const ButteraugliParams& params, BlurTemp* blur_temp,
                        ImageF* BUTTERAUGLI_RESTRICT mask,
                        ImageF* BUTTERAUGLI_RESTRICT diff_ac) {
-  JXL_ASSIGN_OR_RETURN(ImageF mask0, ImageF::Create(xsize, ysize));
+  JxlMemoryManager* memory_manager = pi0.hf[0].memory_manager();
-  JXL_ASSIGN_OR_RETURN(ImageF mask1, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF mask0,
+                       ImageF::Create(memory_manager, xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF mask1,
+                       ImageF::Create(memory_manager, xsize, ysize));
   CombineChannelsForMasking(&pi0.hf[0], &pi0.uhf[0], &mask0);
   CombineChannelsForMasking(&pi1.hf[0], &pi1.uhf[0], &mask1);
   JXL_RETURN_IF_ERROR(Mask(mask0, mask1, params, blur_temp, mask, diff_ac));
@@ -1521,23 +1536,28 @@ Status ButteraugliDiffmapInPlace(Image3F& image0, Image3F& image1,
   // image0 and image1 are in linear sRGB color space
   const size_t xsize = image0.xsize();
   const size_t ysize = image0.ysize();
+  JxlMemoryManager* memory_manager = image0.memory_manager();
   BlurTemp blur_temp;
   {
     // Convert image0 and image1 to XYB in-place
-    JXL_ASSIGN_OR_RETURN(Image3F temp, Image3F::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(Image3F temp,
+                         Image3F::Create(memory_manager, xsize, ysize));
     JXL_RETURN_IF_ERROR(
         OpsinDynamicsImage(image0, params, &temp, &blur_temp, &image0));
     JXL_RETURN_IF_ERROR(
         OpsinDynamicsImage(image1, params, &temp, &blur_temp, &image1));
   }
   // image0 and image1 are in XYB color space
-  JXL_ASSIGN_OR_RETURN(ImageF block_diff_dc, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF block_diff_dc,
+                       ImageF::Create(memory_manager, xsize, ysize));
   ZeroFillImage(&block_diff_dc);
   {
     // separate out LF components from image0 and image1 and compute the dc
     // diff image from them
-    JXL_ASSIGN_OR_RETURN(Image3F lf0, Image3F::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(Image3F lf0,
-    JXL_ASSIGN_OR_RETURN(Image3F lf1, Image3F::Create(xsize, ysize));
+                         Image3F::Create(memory_manager, xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(Image3F lf1,
+                         Image3F::Create(memory_manager, xsize, ysize));
     JXL_RETURN_IF_ERROR(
         SeparateLFAndMF(params, image0, &lf0, &image0, &blur_temp));
     JXL_RETURN_IF_ERROR(
@@ -1553,11 +1573,13 @@ Status ButteraugliDiffmapInPlace(Image3F& image0, Image3F& image1,
   JXL_RETURN_IF_ERROR(SeparateMFAndHF(params, &image1, &hf1[0], &blur_temp));
   // image0 and image1 are MF-images in XYB color space
 
-  JXL_ASSIGN_OR_RETURN(ImageF block_diff_ac, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF block_diff_ac,
+                       ImageF::Create(memory_manager, xsize, ysize));
   ZeroFillImage(&block_diff_ac);
   // start accumulating ac diff image from MF images
   {
-    JXL_ASSIGN_OR_RETURN(ImageF diffs, ImageF::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(ImageF diffs,
+                         ImageF::Create(memory_manager, xsize, ysize));
     MaltaDiffMapLF(image0.Plane(1), image1.Plane(1), wMfMalta, wMfMalta,
                    norm1Mf, &diffs, &block_diff_ac);
     MaltaDiffMapLF(image0.Plane(0), image1.Plane(0), wMfMaltaX, wMfMaltaX,
@@ -1579,7 +1601,8 @@ Status ButteraugliDiffmapInPlace(Image3F& image0, Image3F& image1,
   // continue accumulating ac diff image from HF and UHF images
   const float hf_asymmetry = params.hf_asymmetry;
   {
-    JXL_ASSIGN_OR_RETURN(ImageF diffs, ImageF::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(ImageF diffs,
+                         ImageF::Create(memory_manager, xsize, ysize));
     MaltaDiffMap(uhf0[1], uhf1[1], wUhfMalta * hf_asymmetry,
                  wUhfMalta / hf_asymmetry, norm1Uhf, &diffs, &block_diff_ac);
     MaltaDiffMap(uhf0[0], uhf1[0], wUhfMaltaX * hf_asymmetry,
@@ -1597,10 +1620,13 @@ Status ButteraugliDiffmapInPlace(Image3F& image0, Image3F& image1,
   }
 
   // compute mask image from HF and UHF X and Y images
-  JXL_ASSIGN_OR_RETURN(ImageF mask, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF mask,
+                       ImageF::Create(memory_manager, xsize, ysize));
   {
-    JXL_ASSIGN_OR_RETURN(ImageF mask0, ImageF::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(ImageF mask0,
-    JXL_ASSIGN_OR_RETURN(ImageF mask1, ImageF::Create(xsize, ysize));
+                         ImageF::Create(memory_manager, xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(ImageF mask1,
+                         ImageF::Create(memory_manager, xsize, ysize));
     CombineChannelsForMasking(&hf0[0], &uhf0[0], &mask0);
     CombineChannelsForMasking(&hf1[0], &uhf1[0], &mask1);
     DeallocateHFAndUHF(&hf1[0], &uhf1[0]);
@@ -1610,7 +1636,7 @@ Status ButteraugliDiffmapInPlace(Image3F& image0, Image3F& image1,
   }
 
   // compute final diffmap from mask image and ac and dc diff images
-  JXL_ASSIGN_OR_RETURN(diffmap, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(diffmap, ImageF::Create(memory_manager, xsize, ysize));
   for (size_t y = 0; y < ysize; ++y) {
     const float* row_dc = block_diff_dc.Row(y);
     const float* row_ac = block_diff_ac.Row(y);
@@ -1694,7 +1720,8 @@ static inline void CheckImage(const ImageF& image, const char* name) {
 static StatusOr<Image3F> SubSample2x(const Image3F& in) {
   size_t xs = (in.xsize() + 1) / 2;
   size_t ys = (in.ysize() + 1) / 2;
-  JXL_ASSIGN_OR_RETURN(Image3F retval, Image3F::Create(xs, ys));
+  JxlMemoryManager* memory_manager = in.memory_manager();
+  JXL_ASSIGN_OR_RETURN(Image3F retval, Image3F::Create(memory_manager, xs, ys));
   for (size_t c = 0; c < 3; ++c) {
     for (size_t y = 0; y < ys; ++y) {
       for (size_t x = 0; x < xs; ++x) {
@@ -1754,16 +1781,19 @@ StatusOr<std::unique_ptr<ButteraugliComparator>> ButteraugliComparator::Make(
     const Image3F& rgb0, const ButteraugliParams& params) {
   size_t xsize = rgb0.xsize();
   size_t ysize = rgb0.ysize();
+  JxlMemoryManager* memory_manager = rgb0.memory_manager();
   std::unique_ptr<ButteraugliComparator> result =
       std::unique_ptr<ButteraugliComparator>(
           new ButteraugliComparator(xsize, ysize, params));
-  JXL_ASSIGN_OR_RETURN(result->temp_, Image3F::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(result->temp_,
+                       Image3F::Create(memory_manager, xsize, ysize));
 
   if (xsize < 8 || ysize < 8) {
     return result;
   }
 
-  JXL_ASSIGN_OR_RETURN(Image3F xyb0, Image3F::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(Image3F xyb0,
+                       Image3F::Create(memory_manager, xsize, ysize));
   JXL_RETURN_IF_ERROR(HWY_DYNAMIC_DISPATCH(OpsinDynamicsImage)(
       rgb0, params, result->Temp(), &result->blur_temp_, &xyb0));
   result->ReleaseTemp();
@@ -1789,11 +1819,13 @@ Status ButteraugliComparator::Mask(ImageF* BUTTERAUGLI_RESTRICT mask) const {
 
 Status ButteraugliComparator::Diffmap(const Image3F& rgb1,
                                       ImageF& result) const {
+  JxlMemoryManager* memory_manager = rgb1.memory_manager();
   if (xsize_ < 8 || ysize_ < 8) {
     ZeroFillImage(&result);
     return true;
   }
-  JXL_ASSIGN_OR_RETURN(Image3F xyb1, Image3F::Create(xsize_, ysize_));
+  JXL_ASSIGN_OR_RETURN(Image3F xyb1,
+                       Image3F::Create(memory_manager, xsize_, ysize_));
   JXL_RETURN_IF_ERROR(HWY_DYNAMIC_DISPATCH(OpsinDynamicsImage)(
       rgb1, params_, Temp(), &blur_temp_, &xyb1));
   ReleaseTemp();
@@ -1802,8 +1834,9 @@ Status ButteraugliComparator::Diffmap(const Image3F& rgb1,
     if (sub_->xsize_ < 8 || sub_->ysize_ < 8) {
       return true;
     }
-    JXL_ASSIGN_OR_RETURN(Image3F sub_xyb,
+    JXL_ASSIGN_OR_RETURN(
-                         Image3F::Create(sub_->xsize_, sub_->ysize_));
+        Image3F sub_xyb,
+        Image3F::Create(memory_manager, sub_->xsize_, sub_->ysize_));
     JXL_ASSIGN_OR_RETURN(Image3F subsampledRgb1, SubSample2x(rgb1));
     JXL_RETURN_IF_ERROR(HWY_DYNAMIC_DISPATCH(OpsinDynamicsImage)(
         subsampledRgb1, params_, sub_->Temp(), &sub_->blur_temp_, &sub_xyb));
@@ -1817,6 +1850,7 @@ Status ButteraugliComparator::Diffmap(const Image3F& rgb1,
 
 Status ButteraugliComparator::DiffmapOpsinDynamicsImage(const Image3F& xyb1,
                                                         ImageF& result) const {
+  JxlMemoryManager* memory_manager = xyb1.memory_manager();
   if (xsize_ < 8 || ysize_ < 8) {
     ZeroFillImage(&result);
     return true;
@@ -1824,7 +1858,7 @@ Status ButteraugliComparator::DiffmapOpsinDynamicsImage(const Image3F& xyb1,
   PsychoImage pi1;
   JXL_RETURN_IF_ERROR(HWY_DYNAMIC_DISPATCH(SeparateFrequencies)(
       xsize_, ysize_, params_, &blur_temp_, xyb1, pi1));
-  JXL_ASSIGN_OR_RETURN(result, ImageF::Create(xsize_, ysize_));
+  JXL_ASSIGN_OR_RETURN(result, ImageF::Create(memory_manager, xsize_, ysize_));
   return DiffmapPsychoImage(pi1, result);
 }
 
@@ -1850,6 +1884,7 @@ void MaltaDiffMapLF(const ImageF& lum0, const ImageF& lum1, const double w_0gt1,
 
 Status ButteraugliComparator::DiffmapPsychoImage(const PsychoImage& pi1,
                                                  ImageF& diffmap) const {
+  JxlMemoryManager* memory_manager = diffmap.memory_manager();
   if (xsize_ < 8 || ysize_ < 8) {
     ZeroFillImage(&diffmap);
     return true;
@@ -1858,8 +1893,10 @@ Status ButteraugliComparator::DiffmapPsychoImage(const PsychoImage& pi1,
   const float hf_asymmetry_ = params_.hf_asymmetry;
   const float xmul_ = params_.xmul;
 
-  JXL_ASSIGN_OR_RETURN(ImageF diffs, ImageF::Create(xsize_, ysize_));
+  JXL_ASSIGN_OR_RETURN(ImageF diffs,
-  JXL_ASSIGN_OR_RETURN(Image3F block_diff_ac, Image3F::Create(xsize_, ysize_));
+                       ImageF::Create(memory_manager, xsize_, ysize_));
+  JXL_ASSIGN_OR_RETURN(Image3F block_diff_ac,
+                       Image3F::Create(memory_manager, xsize_, ysize_));
   ZeroFillImage(&block_diff_ac);
   MaltaDiffMap(pi0_.uhf[1], pi1.uhf[1], wUhfMalta * hf_asymmetry_,
                wUhfMalta / hf_asymmetry_, norm1Uhf, &diffs, &block_diff_ac, 1);
@@ -1877,7 +1914,8 @@ Status ButteraugliComparator::DiffmapPsychoImage(const PsychoImage& pi1,
   MaltaDiffMapLF(pi0_.mf.Plane(0), pi1.mf.Plane(0), wMfMaltaX, wMfMaltaX,
                  norm1MfX, &diffs, &block_diff_ac, 0);
 
-  JXL_ASSIGN_OR_RETURN(Image3F block_diff_dc, Image3F::Create(xsize_, ysize_));
+  JXL_ASSIGN_OR_RETURN(Image3F block_diff_dc,
+                       Image3F::Create(memory_manager, xsize_, ysize_));
   for (size_t c = 0; c < 3; ++c) {
     if (c < 2) {  // No blue channel error accumulated at HF.
       HWY_DYNAMIC_DISPATCH(L2DiffAsymmetric)
@@ -1925,6 +1963,7 @@ bool ButteraugliDiffmapSmall(const Image3F& rgb0, const Image3F& rgb1,
                              const ButteraugliParams& params, ImageF& diffmap) {
   const size_t xsize = rgb0.xsize();
   const size_t ysize = rgb0.ysize();
+  JxlMemoryManager* memory_manager = rgb0.memory_manager();
   // Butteraugli values for small (where xsize or ysize is smaller
   // than 8 pixels) images are non-sensical, but most likely it is
   // less disruptive to try to compute something than just give up.
@@ -1933,8 +1972,10 @@ bool ButteraugliDiffmapSmall(const Image3F& rgb0, const Image3F& rgb1,
   size_t yborder = ysize < kMax ? (kMax - ysize) / 2 : 0;
   size_t xscaled = std::max<size_t>(kMax, xsize);
   size_t yscaled = std::max<size_t>(kMax, ysize);
-  JXL_ASSIGN_OR_RETURN(Image3F scaled0, Image3F::Create(xscaled, yscaled));
+  JXL_ASSIGN_OR_RETURN(Image3F scaled0,
-  JXL_ASSIGN_OR_RETURN(Image3F scaled1, Image3F::Create(xscaled, yscaled));
+                       Image3F::Create(memory_manager, xscaled, yscaled));
+  JXL_ASSIGN_OR_RETURN(Image3F scaled1,
+                       Image3F::Create(memory_manager, xscaled, yscaled));
   for (int i = 0; i < 3; ++i) {
     for (size_t y = 0; y < yscaled; ++y) {
       for (size_t x = 0; x < xscaled; ++x) {
@@ -1947,7 +1988,7 @@ bool ButteraugliDiffmapSmall(const Image3F& rgb0, const Image3F& rgb1,
   }
   ImageF diffmap_scaled;
   const bool ok = ButteraugliDiffmap(scaled0, scaled1, params, diffmap_scaled);
-  JXL_ASSIGN_OR_RETURN(diffmap, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(diffmap, ImageF::Create(memory_manager, xsize, ysize));
   for (size_t y = 0; y < ysize; ++y) {
     for (size_t x = 0; x < xsize; ++x) {
       diffmap.Row(y)[x] = diffmap_scaled.Row(y + yborder)[x + xborder];
@@ -2111,8 +2152,10 @@ void ScoreToRgb(double score, double good_threshold, double bad_threshold,
 StatusOr<Image3F> CreateHeatMapImage(const ImageF& distmap,
                                      double good_threshold,
                                      double bad_threshold) {
-  JXL_ASSIGN_OR_RETURN(Image3F heatmap,
+  JxlMemoryManager* memory_manager = distmap.memory_manager();
-                       Image3F::Create(distmap.xsize(), distmap.ysize()));
+  JXL_ASSIGN_OR_RETURN(
+      Image3F heatmap,
+      Image3F::Create(memory_manager, distmap.xsize(), distmap.ysize()));
   for (size_t y = 0; y < distmap.ysize(); ++y) {
     const float* BUTTERAUGLI_RESTRICT row_distmap = distmap.ConstRow(y);
     float* BUTTERAUGLI_RESTRICT row_h0 = heatmap.PlaneRow(0, y);

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

@@ -8,11 +8,12 @@
 #ifndef LIB_JXL_BUTTERAUGLI_BUTTERAUGLI_H_
 #define LIB_JXL_BUTTERAUGLI_BUTTERAUGLI_H_
 
-#include <stdlib.h>
+#include <jxl/memory_manager.h>
-#include <string.h>
 
 #include <atomic>
 #include <cstddef>
+#include <cstdlib>
+#include <cstring>
 #include <memory>
 
 #include "lib/jxl/base/compiler_specific.h"
@@ -149,9 +150,11 @@ struct PsychoImage {
 // Hold it here and only allocate on demand to reduce memory usage.
 struct BlurTemp {
   Status GetTransposed(const ImageF &in, ImageF **out) {
+    JxlMemoryManager *memory_manager = in.memory_manager();
     if (transposed_temp.xsize() == 0) {
-      JXL_ASSIGN_OR_RETURN(transposed_temp,
+      JXL_ASSIGN_OR_RETURN(
-                           ImageF::Create(in.ysize(), in.xsize()));
+          transposed_temp,
+          ImageF::Create(memory_manager, in.ysize(), in.xsize()));
     }
     *out = &transposed_temp;
     return true;

third_party/jpeg-xl/lib/jxl/butteraugli/butteraugli_test.cc

@@ -5,8 +5,8 @@
 
 #include "lib/jxl/butteraugli/butteraugli.h"
 
+#include <jxl/memory_manager.h>
 #include <jxl/types.h>
-#include <stddef.h>
 
 #include <algorithm>
 #include <cstdint>
@@ -20,6 +20,7 @@
 #include "lib/jxl/image.h"
 #include "lib/jxl/image_ops.h"
 #include "lib/jxl/test_image.h"
+#include "lib/jxl/test_utils.h"
 #include "lib/jxl/testing.h"
 
 namespace jxl {
@@ -30,7 +31,8 @@ using extras::PackedPixelFile;
 using test::TestImage;
 
 Image3F SinglePixelImage(float red, float green, float blue) {
-  JXL_ASSIGN_OR_DIE(Image3F img, Image3F::Create(1, 1));
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
+  JXL_ASSIGN_OR_DIE(Image3F img, Image3F::Create(memory_manager, 1, 1));
   img.PlaneRow(0, 0)[0] = red;
   img.PlaneRow(1, 0)[0] = green;
   img.PlaneRow(2, 0)[0] = blue;
@@ -38,11 +40,13 @@ Image3F SinglePixelImage(float red, float green, float blue) {
 }
 
 Image3F GetColorImage(const PackedPixelFile& ppf) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   JXL_CHECK(!ppf.frames.empty());
   const PackedImage& image = ppf.frames[0].color;
   const JxlPixelFormat& format = image.format;
   const uint8_t* pixels = reinterpret_cast<const uint8_t*>(image.pixels());
-  JXL_ASSIGN_OR_DIE(Image3F color, Image3F::Create(image.xsize, image.ysize));
+  JXL_ASSIGN_OR_DIE(Image3F color,
+                    Image3F::Create(memory_manager, image.xsize, image.ysize));
   for (size_t c = 0; c < format.num_channels; ++c) {
     JXL_CHECK(ConvertFromExternal(pixels, image.pixels_size, image.xsize,
                                   image.ysize, ppf.info.bits_per_sample, format,
@@ -88,12 +92,14 @@ TEST(ButteraugliInPlaceTest, SinglePixel) {
 }
 
 TEST(ButteraugliInPlaceTest, LargeImage) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   const size_t xsize = 1024;
   const size_t ysize = 1024;
   TestImage img;
   img.SetDimensions(xsize, ysize).AddFrame().RandomFill(777);
   Image3F rgb0 = GetColorImage(img.ppf());
-  JXL_ASSIGN_OR_DIE(Image3F rgb1, Image3F::Create(xsize, ysize));
+  JXL_ASSIGN_OR_DIE(Image3F rgb1,
+                    Image3F::Create(memory_manager, xsize, ysize));
   CopyImageTo(rgb0, &rgb1);
   AddUniformNoise(&rgb1, 0.02f, 7777);
   AddEdge(&rgb1, 0.1f, xsize / 2, xsize / 2);

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

@@ -5,18 +5,21 @@
 
 #include "lib/jxl/chroma_from_luma.h"
 
+#include <jxl/memory_manager.h>
+
 #include "lib/jxl/image_ops.h"
 
 namespace jxl {
 
-StatusOr<ColorCorrelationMap> ColorCorrelationMap::Create(size_t xsize,
+StatusOr<ColorCorrelationMap> ColorCorrelationMap::Create(
-                                                          size_t ysize,
+    JxlMemoryManager* memory_manager, size_t xsize, size_t ysize, bool XYB) {
-                                                          bool XYB) {
   ColorCorrelationMap result;
   size_t xblocks = DivCeil(xsize, kColorTileDim);
   size_t yblocks = DivCeil(ysize, kColorTileDim);
-  JXL_ASSIGN_OR_RETURN(result.ytox_map, ImageSB::Create(xblocks, yblocks));
+  JXL_ASSIGN_OR_RETURN(result.ytox_map,
-  JXL_ASSIGN_OR_RETURN(result.ytob_map, ImageSB::Create(xblocks, yblocks));
+                       ImageSB::Create(memory_manager, xblocks, yblocks));
+  JXL_ASSIGN_OR_RETURN(result.ytob_map,
+                       ImageSB::Create(memory_manager, xblocks, yblocks));
   ZeroFillImage(&result.ytox_map);
   ZeroFillImage(&result.ytob_map);
   if (!XYB) {

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

@@ -9,6 +9,8 @@
 // Chroma-from-luma, computed using heuristics to determine the best linear
 // model for the X and B channels from the Y channel.
 
+#include <jxl/memory_manager.h>
+
 #include <cmath>
 #include <cstddef>
 #include <cstdint>
@@ -52,7 +54,8 @@ struct ColorCorrelationMap {
   // xsize/ysize are in pixels
   // set XYB=false to do something close to no-op cmap (needed for now since
   // cmap is mandatory)
-  static StatusOr<ColorCorrelationMap> Create(size_t xsize, size_t ysize,
+  static StatusOr<ColorCorrelationMap> Create(JxlMemoryManager* memory_manager,
+                                              size_t xsize, size_t ysize,
                                               bool XYB = true);
 
   float YtoXRatio(int32_t x_factor) const {

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

@@ -76,7 +76,7 @@ constexpr Matrix3x3 kBradfordInv{{{0.9869929f, -0.1470543f, 0.1599627f},
                                   {0.4323053f, 0.5183603f, 0.0492912f},
                                   {-0.0085287f, 0.0400428f, 0.9684867f}}};
 
-// Adapts whitepoint x, y to D50
+// Adapts white point x, y to D50
 static Status AdaptToXYZD50(float wx, float wy, Matrix3x3& matrix) {
   bool ok = (wx >= 0) && (wx <= 1) && (wy > 0) && (wy <= 1);
   if (!ok) {
@@ -359,7 +359,7 @@ static Status CreateICCChadMatrix(double wx, double wy, Matrix3x3& result) {
   return true;
 }
 
-// Creates RGB to XYZ matrix given RGB primaries and whitepoint in xy.
+// Creates RGB to XYZ matrix given RGB primaries and white point in xy.
 static Status CreateICCRGBMatrix(double rx, double ry, double gx, double gy,
                                  double bx, double by, double wx, double wy,
                                  Matrix3x3& result) {

third_party/jpeg-xl/lib/jxl/cms/transfer_functions_test.cc

@@ -34,7 +34,7 @@ HWY_NOINLINE void TestPqEncodedFromDisplay() {
     const float actual = GetLane(tf_pq.EncodedFromDisplay(d, Set(d, f)));
     const float expected = TF_PQ_Base::EncodedFromDisplay(intensity, f);
     const float abs_err = std::abs(expected - actual);
-    EXPECT_LT(abs_err, 5e-7) << "f = " << f;
+    EXPECT_LT(abs_err, 6e-7) << "f = " << f;
     max_abs_err = std::max(max_abs_err, abs_err);
   }
   printf("max abs err %e\n", static_cast<double>(max_abs_err));

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

@@ -8,16 +8,14 @@
 
 // Holds inputs/outputs for decoding/encoding images.
 
-#include <stddef.h>
+#include <jxl/memory_manager.h>
-#include <stdint.h>
 
-#include <type_traits>
+#include <cstddef>
+#include <cstdint>
 #include <utility>
 #include <vector>
 
-#include "lib/jxl/base/data_parallel.h"
 #include "lib/jxl/base/status.h"
-#include "lib/jxl/frame_header.h"
 #include "lib/jxl/headers.h"
 #include "lib/jxl/image.h"
 #include "lib/jxl/image_bundle.h"
@@ -37,9 +35,11 @@ struct Blobs {
 // to/from decoding/encoding.
 class CodecInOut {
  public:
-  CodecInOut() : preview_frame(&metadata.m) {
+  explicit CodecInOut(JxlMemoryManager* memory_manager)
+      : memory_manager(memory_manager),
+        preview_frame(memory_manager, &metadata.m) {
     frames.reserve(1);
-    frames.emplace_back(&metadata.m);
+    frames.emplace_back(memory_manager, &metadata.m);
   }
 
   // Move-only.
@@ -97,6 +97,8 @@ class CodecInOut {
 
   // Metadata stored into / retrieved from bitstreams.
 
+  JxlMemoryManager* memory_manager;
+
   Blobs blobs;
 
   CodecMetadata metadata;  // applies to preview and all frames

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

@@ -5,20 +5,18 @@
 
 #include "lib/jxl/coeff_order.h"
 
-#include <stdint.h>
+#include <jxl/memory_manager.h>
 
 #include <algorithm>
+#include <cstdint>
 #include <vector>
 
-#include "lib/jxl/ans_params.h"
+#include "lib/jxl/base/status.h"
-#include "lib/jxl/base/span.h"
 #include "lib/jxl/coeff_order_fwd.h"
 #include "lib/jxl/dec_ans.h"
 #include "lib/jxl/dec_bit_reader.h"
-#include "lib/jxl/entropy_coder.h"
 #include "lib/jxl/lehmer_code.h"
 #include "lib/jxl/modular/encoding/encoding.h"
-#include "lib/jxl/modular/modular_image.h"
 
 namespace jxl {
 
@@ -57,13 +55,14 @@ Status ReadPermutation(size_t skip, size_t size, coeff_order_t* order,
 
 }  // namespace
 
-Status DecodePermutation(size_t skip, size_t size, coeff_order_t* order,
+Status DecodePermutation(JxlMemoryManager* memory_manager, size_t skip,
-                         BitReader* br) {
+                         size_t size, coeff_order_t* order, BitReader* br) {
   std::vector<uint8_t> context_map;
   ANSCode code;
-  JXL_RETURN_IF_ERROR(
+  JXL_RETURN_IF_ERROR(DecodeHistograms(memory_manager, br, kPermutationContexts,
-      DecodeHistograms(br, kPermutationContexts, &code, &context_map));
+                                       &code, &context_map));
-  ANSSymbolReader reader(&code, br);
+  JXL_ASSIGN_OR_RETURN(ANSSymbolReader reader,
+                       ANSSymbolReader::Create(&code, br));
   JXL_RETURN_IF_ERROR(
       ReadPermutation(skip, size, order, br, &reader, context_map));
   if (!reader.CheckANSFinalState()) {
@@ -92,18 +91,19 @@ Status DecodeCoeffOrder(AcStrategy acs, coeff_order_t* order, BitReader* br,
 
 }  // namespace
 
-Status DecodeCoeffOrders(uint16_t used_orders, uint32_t used_acs,
+Status DecodeCoeffOrders(JxlMemoryManager* memory_manager, uint16_t used_orders,
-                         coeff_order_t* order, BitReader* br) {
+                         uint32_t used_acs, coeff_order_t* order,
+                         BitReader* br) {
   uint16_t computed = 0;
   std::vector<uint8_t> context_map;
   ANSCode code;
-  std::unique_ptr<ANSSymbolReader> reader;
+  ANSSymbolReader reader;
   std::vector<coeff_order_t> natural_order;
   // Bitstream does not have histograms if no coefficient order is used.
   if (used_orders != 0) {
-    JXL_RETURN_IF_ERROR(
+    JXL_RETURN_IF_ERROR(DecodeHistograms(
-        DecodeHistograms(br, kPermutationContexts, &code, &context_map));
+        memory_manager, br, kPermutationContexts, &code, &context_map));
-    reader = make_unique<ANSSymbolReader>(&code, br);
+    JXL_ASSIGN_OR_RETURN(reader, ANSSymbolReader::Create(&code, br));
   }
   uint32_t acs_mask = 0;
   for (uint8_t o = 0; o < AcStrategy::kNumValidStrategies; ++o) {
@@ -136,12 +136,12 @@ Status DecodeCoeffOrders(uint16_t used_orders, uint32_t used_acs,
     } else {
       for (size_t c = 0; c < 3; c++) {
         coeff_order_t* dest = used ? &order[CoeffOrderOffset(ord, c)] : nullptr;
-        JXL_RETURN_IF_ERROR(DecodeCoeffOrder(acs, dest, br, reader.get(),
+        JXL_RETURN_IF_ERROR(DecodeCoeffOrder(acs, dest, br, &reader,
                                              natural_order, context_map));
       }
     }
   }
-  if (used_orders && !reader->CheckANSFinalState()) {
+  if (used_orders && !reader.CheckANSFinalState()) {
     return JXL_FAILURE("Invalid ANS stream");
   }
   return true;

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

@@ -6,6 +6,8 @@
 #ifndef LIB_JXL_COEFF_ORDER_H_
 #define LIB_JXL_COEFF_ORDER_H_
 
+#include <jxl/memory_manager.h>
+
 #include <array>
 #include <cstddef>
 #include <cstdint>
@@ -53,12 +55,13 @@ constexpr JXL_MAYBE_UNUSED uint32_t kPermutationContexts = 8;
 
 uint32_t CoeffOrderContext(uint32_t val);
 
-Status DecodeCoeffOrders(uint16_t used_orders, uint32_t used_acs,
+Status DecodeCoeffOrders(JxlMemoryManager* memory_manager, uint16_t used_orders,
-                         coeff_order_t* order, BitReader* br);
+                         uint32_t used_acs, coeff_order_t* order,
-
-Status DecodePermutation(size_t skip, size_t size, coeff_order_t* order,
                          BitReader* br);
 
+Status DecodePermutation(JxlMemoryManager* memory_manager, size_t skip,
+                         size_t size, coeff_order_t* order, BitReader* br);
+
 }  // namespace jxl
 
 #endif  // LIB_JXL_COEFF_ORDER_H_

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

@@ -5,6 +5,8 @@
 
 #include "lib/jxl/coeff_order.h"
 
+#include <jxl/memory_manager.h>
+
 #include <algorithm>
 #include <numeric>  // iota
 #include <utility>
@@ -16,6 +18,7 @@
 #include "lib/jxl/coeff_order_fwd.h"
 #include "lib/jxl/dec_bit_reader.h"
 #include "lib/jxl/enc_coeff_order.h"
+#include "lib/jxl/test_utils.h"
 #include "lib/jxl/testing.h"
 
 namespace jxl {
@@ -23,14 +26,15 @@ namespace {
 
 void RoundtripPermutation(coeff_order_t* perm, coeff_order_t* out, size_t len,
                           size_t* size) {
-  BitWriter writer;
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
+  BitWriter writer{memory_manager};
   EncodePermutation(perm, 0, len, &writer, 0, nullptr);
   writer.ZeroPadToByte();
   Status status = true;
   {
     BitReader reader(writer.GetSpan());
     BitReaderScopedCloser closer(&reader, &status);
-    ASSERT_TRUE(DecodePermutation(0, len, out, &reader));
+    ASSERT_TRUE(DecodePermutation(memory_manager, 0, len, out, &reader));
   }
   ASSERT_TRUE(status);
   *size = writer.GetSpan().size();

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

@@ -155,7 +155,7 @@ struct ColorEncoding : public Fields {
   }
 
   // Sets the raw ICC profile bytes, without parsing the ICC, and without
-  // updating the direct fields such as whitepoint, primaries and color
+  // updating the direct fields such as white point, primaries and color
   // space. Functions to get and set fields, such as SetWhitePoint, cannot be
   // used anymore after this and functions such as IsSRGB return false no matter
   // what the contents of the icc profile.

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

@@ -5,7 +5,7 @@
 
 #include <jxl/cms.h>
 #include <jxl/cms_interface.h>
-#include <stdint.h>
+#include <jxl/memory_manager.h>
 
 #include <algorithm>
 #include <cmath>
@@ -61,17 +61,19 @@ struct Globals {
 
  private:
   Globals() {
+    JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
     in_gray = GenerateGray();
     in_color = GenerateColor();
-    JXL_ASSIGN_OR_DIE(out_gray, ImageF::Create(kWidth, 1));
+    JXL_ASSIGN_OR_DIE(out_gray, ImageF::Create(memory_manager, kWidth, 1));
-    JXL_ASSIGN_OR_DIE(out_color, ImageF::Create(kWidth * 3, 1));
+    JXL_ASSIGN_OR_DIE(out_color, ImageF::Create(memory_manager, kWidth * 3, 1));
 
     c_native = ColorEncoding::LinearSRGB(/*is_gray=*/false);
     c_gray = ColorEncoding::LinearSRGB(/*is_gray=*/true);
   }
 
   static ImageF GenerateGray() {
-    JXL_ASSIGN_OR_DIE(ImageF gray, ImageF::Create(kWidth, 1));
+    JXL_ASSIGN_OR_DIE(ImageF gray,
+                      ImageF::Create(jxl::test::MemoryManager(), kWidth, 1));
     float* JXL_RESTRICT row = gray.Row(0);
     // Increasing left to right
     for (uint32_t x = 0; x < kWidth; ++x) {
@@ -81,7 +83,8 @@ struct Globals {
   }
 
   static ImageF GenerateColor() {
-    JXL_ASSIGN_OR_DIE(ImageF image, ImageF::Create(kWidth * 3, 1));
+    JXL_ASSIGN_OR_DIE(ImageF image, ImageF::Create(jxl::test::MemoryManager(),
+                                                   kWidth * 3, 1));
     float* JXL_RESTRICT interleaved = image.Row(0);
     std::fill(interleaved, interleaved + kWidth * 3, 0.0f);
 
@@ -343,6 +346,7 @@ TEST_F(ColorManagementTest, HlgOotf) {
 }
 
 TEST_F(ColorManagementTest, XYBProfile) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   ColorEncoding c_xyb;
   c_xyb.SetColorSpace(ColorSpace::kXYB);
   c_xyb.SetRenderingIntent(RenderingIntent::kPerceptual);
@@ -358,8 +362,9 @@ TEST_F(ColorManagementTest, XYBProfile) {
 
   ImageMetadata metadata;
   metadata.color_encoding = c_native;
-  ImageBundle ib(&metadata);
+  ImageBundle ib(memory_manager, &metadata);
-  JXL_ASSIGN_OR_DIE(Image3F native, Image3F::Create(kNumColors, 1));
+  JXL_ASSIGN_OR_DIE(Image3F native,
+                    Image3F::Create(memory_manager, kNumColors, 1));
   float mul = 1.0f / (kGridDim - 1);
   for (size_t ir = 0, x = 0; ir < kGridDim; ++ir) {
     for (size_t ig = 0; ig < kGridDim; ++ig) {
@@ -372,10 +377,12 @@ TEST_F(ColorManagementTest, XYBProfile) {
   }
   ib.SetFromImage(std::move(native), c_native);
   const Image3F& in = *ib.color();
-  JXL_ASSIGN_OR_DIE(Image3F opsin, Image3F::Create(kNumColors, 1));
+  JXL_ASSIGN_OR_DIE(Image3F opsin,
+                    Image3F::Create(memory_manager, kNumColors, 1));
   JXL_CHECK(ToXYB(ib, nullptr, &opsin, cms, nullptr));
 
-  JXL_ASSIGN_OR_DIE(Image3F opsin2, Image3F::Create(kNumColors, 1));
+  JXL_ASSIGN_OR_DIE(Image3F opsin2,
+                    Image3F::Create(memory_manager, kNumColors, 1));
   CopyImageTo(opsin, &opsin2);
   ScaleXYB(&opsin2);
 
@@ -389,7 +396,8 @@ TEST_F(ColorManagementTest, XYBProfile) {
   float* dst = xform.BufDst(0);
   ASSERT_TRUE(xform.Run(0, src, dst, kNumColors));
 
-  JXL_ASSIGN_OR_DIE(Image3F out, Image3F::Create(kNumColors, 1));
+  JXL_ASSIGN_OR_DIE(Image3F out,
+                    Image3F::Create(memory_manager, kNumColors, 1));
   for (size_t i = 0; i < kNumColors; ++i) {
     for (size_t c = 0; c < 3; ++c) {
       out.PlaneRow(c, 0)[i] = dst[3 * i + c];

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

@@ -5,11 +5,12 @@
 
 #include "lib/jxl/compressed_dc.h"
 
-#include <stdint.h>
+#include <jxl/memory_manager.h>
-#include <stdlib.h>
-#include <string.h>
 
 #include <algorithm>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
 #include <vector>
 
 #undef HWY_TARGET_INCLUDE
@@ -121,7 +122,8 @@ JXL_INLINE void ComputePixel(
   Store(out, d, out_rows[2] + x);
 }
 
-Status AdaptiveDCSmoothing(const float* dc_factors, Image3F* dc,
+Status AdaptiveDCSmoothing(JxlMemoryManager* memory_manager,
+                           const float* dc_factors, Image3F* dc,
                            ThreadPool* pool) {
   const size_t xsize = dc->xsize();
   const size_t ysize = dc->ysize();
@@ -132,7 +134,8 @@ Status AdaptiveDCSmoothing(const float* dc_factors, Image3F* dc,
   // the x and b channels through color correlation.
   JXL_ASSERT(w1 + w2 < 0.25f);
 
-  JXL_ASSIGN_OR_RETURN(Image3F smoothed, Image3F::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(Image3F smoothed,
+                       Image3F::Create(memory_manager, xsize, ysize));
   // Fill in borders that the loop below will not. First and last are unused.
   for (size_t c = 0; c < 3; c++) {
     for (size_t y : {static_cast<size_t>(0), ysize - 1}) {
@@ -289,9 +292,11 @@ namespace jxl {
 
 HWY_EXPORT(DequantDC);
 HWY_EXPORT(AdaptiveDCSmoothing);
-Status AdaptiveDCSmoothing(const float* dc_factors, Image3F* dc,
+Status AdaptiveDCSmoothing(JxlMemoryManager* memory_manager,
+                           const float* dc_factors, Image3F* dc,
                            ThreadPool* pool) {
-  return HWY_DYNAMIC_DISPATCH(AdaptiveDCSmoothing)(dc_factors, dc, pool);
+  return HWY_DYNAMIC_DISPATCH(AdaptiveDCSmoothing)(memory_manager, dc_factors,
+                                                   dc, pool);
 }
 
 void DequantDC(const Rect& r, Image3F* dc, ImageB* quant_dc, const Image& in,

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

@@ -6,6 +6,8 @@
 #ifndef LIB_JXL_COMPRESSED_DC_H_
 #define LIB_JXL_COMPRESSED_DC_H_
 
+#include <jxl/memory_manager.h>
+
 #include "lib/jxl/ac_context.h"
 #include "lib/jxl/base/data_parallel.h"
 #include "lib/jxl/base/rect.h"
@@ -20,7 +22,8 @@
 namespace jxl {
 
 // Smooth DC in already-smooth areas, to counteract banding.
-Status AdaptiveDCSmoothing(const float* dc_factors, Image3F* dc,
+Status AdaptiveDCSmoothing(JxlMemoryManager* memory_manager,
+                           const float* dc_factors, Image3F* dc,
                            ThreadPool* pool);
 
 void DequantDC(const Rect& r, Image3F* dc, ImageB* quant_dc, const Image& in,

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

@@ -5,10 +5,11 @@
 
 #include "lib/jxl/convolve.h"
 
+#include <jxl/memory_manager.h>
 #include <jxl/types.h>
-#include <time.h>
 
 #include <cinttypes>  // PRIx64
+#include <ctime>
 
 #undef HWY_TARGET_INCLUDE
 #define HWY_TARGET_INCLUDE "lib/jxl/convolve_test.cc"
@@ -70,13 +71,16 @@ void TestNeighbors() {
 
 void VerifySymmetric3(const size_t xsize, const size_t ysize, ThreadPool* pool,
                       Rng* rng) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   const Rect rect(0, 0, xsize, ysize);
 
-  JXL_ASSIGN_OR_DIE(ImageF in, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_DIE(ImageF in, ImageF::Create(memory_manager, xsize, ysize));
   GenerateImage(*rng, &in, 0.0f, 1.0f);
 
-  JXL_ASSIGN_OR_DIE(ImageF out_expected, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_DIE(ImageF out_expected,
-  JXL_ASSIGN_OR_DIE(ImageF out_actual, ImageF::Create(xsize, ysize));
+                    ImageF::Create(memory_manager, xsize, ysize));
+  JXL_ASSIGN_OR_DIE(ImageF out_actual,
+                    ImageF::Create(memory_manager, xsize, ysize));
 
   const WeightsSymmetric3& weights = WeightsSymmetric3Lowpass();
   Symmetric3(in, rect, weights, pool, &out_expected);
@@ -100,7 +104,8 @@ std::vector<Rect> GenerateTestRectangles(size_t xsize, size_t ysize) {
 // Ensures Symmetric and Separable give the same result.
 void VerifySymmetric5(const size_t xsize, const size_t ysize, ThreadPool* pool,
                       Rng* rng) {
-  JXL_ASSIGN_OR_DIE(ImageF in, ImageF::Create(xsize, ysize));
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
+  JXL_ASSIGN_OR_DIE(ImageF in, ImageF::Create(memory_manager, xsize, ysize));
   GenerateImage(*rng, &in, 0.0f, 1.0f);
 
   for (const Rect& in_rect : GenerateTestRectangles(xsize, ysize)) {
@@ -109,8 +114,10 @@ void VerifySymmetric5(const size_t xsize, const size_t ysize, ThreadPool* pool,
               in_rect.xsize(), in_rect.ysize(), in_rect.x0(), in_rect.y0());
     {
       Rect out_rect = in_rect;
-      JXL_ASSIGN_OR_DIE(ImageF out_expected, ImageF::Create(xsize, ysize));
+      JXL_ASSIGN_OR_DIE(ImageF out_expected,
-      JXL_ASSIGN_OR_DIE(ImageF out_actual, ImageF::Create(xsize, ysize));
+                        ImageF::Create(memory_manager, xsize, ysize));
+      JXL_ASSIGN_OR_DIE(ImageF out_actual,
+                        ImageF::Create(memory_manager, xsize, ysize));
       FillImage(-1.0f, &out_expected);
       FillImage(-1.0f, &out_actual);
 
@@ -124,10 +131,12 @@ void VerifySymmetric5(const size_t xsize, const size_t ysize, ThreadPool* pool,
     }
     {
       Rect out_rect(0, 0, in_rect.xsize(), in_rect.ysize());
-      JXL_ASSIGN_OR_DIE(ImageF out_expected,
+      JXL_ASSIGN_OR_DIE(
-                        ImageF::Create(out_rect.xsize(), out_rect.ysize()));
+          ImageF out_expected,
-      JXL_ASSIGN_OR_DIE(ImageF out_actual,
+          ImageF::Create(memory_manager, out_rect.xsize(), out_rect.ysize()));
-                        ImageF::Create(out_rect.xsize(), out_rect.ysize()));
+      JXL_ASSIGN_OR_DIE(
+          ImageF out_actual,
+          ImageF::Create(memory_manager, out_rect.xsize(), out_rect.ysize()));
 
       SlowSeparable5(in, in_rect, WeightsSeparable5Lowpass(), pool,
                      &out_expected, out_rect);
@@ -142,13 +151,16 @@ void VerifySymmetric5(const size_t xsize, const size_t ysize, ThreadPool* pool,
 
 void VerifySeparable5(const size_t xsize, const size_t ysize, ThreadPool* pool,
                       Rng* rng) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   const Rect rect(0, 0, xsize, ysize);
 
-  JXL_ASSIGN_OR_DIE(ImageF in, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_DIE(ImageF in, ImageF::Create(memory_manager, xsize, ysize));
   GenerateImage(*rng, &in, 0.0f, 1.0f);
 
-  JXL_ASSIGN_OR_DIE(ImageF out_expected, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_DIE(ImageF out_expected,
-  JXL_ASSIGN_OR_DIE(ImageF out_actual, ImageF::Create(xsize, ysize));
+                    ImageF::Create(memory_manager, xsize, ysize));
+  JXL_ASSIGN_OR_DIE(ImageF out_actual,
+                    ImageF::Create(memory_manager, xsize, ysize));
 
   const WeightsSeparable5& weights = WeightsSeparable5Lowpass();
   SlowSeparable5(in, rect, weights, pool, &out_expected, rect);
@@ -198,15 +210,16 @@ void TestConvolve() {
 template <class Conv>
 void BenchmarkConv(const char* caption, const Conv& conv,
                    const hwy::FuncInput unpredictable1) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   const size_t kNumInputs = 1;
   const hwy::FuncInput inputs[kNumInputs] = {unpredictable1};
   hwy::Result results[kNumInputs];
 
   const size_t kDim = 160;  // in+out fit in L2
-  JXL_ASSIGN_OR_DIE(ImageF in, ImageF::Create(kDim, kDim));
+  JXL_ASSIGN_OR_DIE(ImageF in, ImageF::Create(memory_manager, kDim, kDim));
   ZeroFillImage(&in);
   in.Row(kDim / 2)[kDim / 2] = unpredictable1;
-  JXL_ASSIGN_OR_DIE(ImageF out, ImageF::Create(kDim, kDim));
+  JXL_ASSIGN_OR_DIE(ImageF out, ImageF::Create(memory_manager, kDim, kDim));
 
   hwy::Params p;
   p.verbose = false;

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

@@ -6,6 +6,8 @@
 #ifndef LIB_JXL_DCT_UTIL_H_
 #define LIB_JXL_DCT_UTIL_H_
 
+#include <jxl/memory_manager.h>
+
 #include <cstddef>
 #include <cstdint>
 #include <memory>
@@ -53,12 +55,14 @@ class ACImageT final : public ACImage {
  public:
   ACImageT() = default;
 
-  static StatusOr<std::unique_ptr<ACImageT>> Make(size_t xsize, size_t ysize) {
+  static StatusOr<std::unique_ptr<ACImageT>> Make(
+      JxlMemoryManager* memory_manager, size_t xsize, size_t ysize) {
     static_assert(
         std::is_same<T, int16_t>::value || std::is_same<T, int32_t>::value,
         "ACImage must be either 32- or 16- bit");
     std::unique_ptr<ACImageT> result = jxl::make_unique<ACImageT>();
-    JXL_ASSIGN_OR_RETURN(result->img_, Image3<T>::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(result->img_,
+                         Image3<T>::Create(memory_manager, xsize, ysize));
     return result;
   }
 

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

@@ -5,8 +5,9 @@
 
 #include "lib/jxl/dec_ans.h"
 
-#include <stdint.h>
+#include <jxl/memory_manager.h>
 
+#include <cstdint>
 #include <vector>
 
 #include "lib/jxl/ans_common.h"
@@ -182,9 +183,11 @@ Status ReadHistogram(int precision_bits, std::vector<int32_t>* counts,
 
 }  // namespace
 
-Status DecodeANSCodes(const size_t num_histograms,
+Status DecodeANSCodes(JxlMemoryManager* memory_manager,
+                      const size_t num_histograms,
                       const size_t max_alphabet_size, BitReader* in,
                       ANSCode* result) {
+  result->memory_manager = memory_manager;
   result->degenerate_symbols.resize(num_histograms, -1);
   if (result->use_prefix_code) {
     JXL_ASSERT(max_alphabet_size <= 1 << PREFIX_MAX_BITS);
@@ -220,9 +223,9 @@ Status DecodeANSCodes(const size_t num_histograms,
     }
   } else {
     JXL_ASSERT(max_alphabet_size <= ANS_MAX_ALPHABET_SIZE);
-    result->alias_tables =
+    size_t alloc_size = num_histograms * (1 << result->log_alpha_size) *
-        AllocateArray(num_histograms * (1 << result->log_alpha_size) *
+                        sizeof(AliasTable::Entry);
-                      sizeof(AliasTable::Entry));
+    result->alias_tables = AllocateArray(alloc_size);
     AliasTable::Entry* alias_tables =
         reinterpret_cast<AliasTable::Entry*>(result->alias_tables.get());
     for (size_t c = 0; c < num_histograms; ++c) {
@@ -325,7 +328,8 @@ void ANSCode::UpdateMaxNumBits(size_t ctx, size_t symbol) {
   max_num_bits = std::max(max_num_bits, total_bits);
 }
 
-Status DecodeHistograms(BitReader* br, size_t num_contexts, ANSCode* code,
+Status DecodeHistograms(JxlMemoryManager* memory_manager, BitReader* br,
+                        size_t num_contexts, ANSCode* code,
                         std::vector<uint8_t>* context_map, bool disallow_lz77) {
   JXL_RETURN_IF_ERROR(Bundle::Read(br, &code->lz77));
   if (code->lz77.enabled) {
@@ -339,7 +343,8 @@ Status DecodeHistograms(BitReader* br, size_t num_contexts, ANSCode* code,
   size_t num_histograms = 1;
   context_map->resize(num_contexts);
   if (num_contexts > 1) {
-    JXL_RETURN_IF_ERROR(DecodeContextMap(context_map, &num_histograms, br));
+    JXL_RETURN_IF_ERROR(
+        DecodeContextMap(memory_manager, context_map, &num_histograms, br));
   }
   JXL_DEBUG_V(
       4, "Decoded context map of size %" PRIuS " and %" PRIuS " histograms",
@@ -355,9 +360,44 @@ Status DecodeHistograms(BitReader* br, size_t num_contexts, ANSCode* code,
   JXL_RETURN_IF_ERROR(
       DecodeUintConfigs(code->log_alpha_size, &code->uint_config, br));
   const size_t max_alphabet_size = 1 << code->log_alpha_size;
-  JXL_RETURN_IF_ERROR(
+  JXL_RETURN_IF_ERROR(DecodeANSCodes(memory_manager, num_histograms,
-      DecodeANSCodes(num_histograms, max_alphabet_size, br, code));
+                                     max_alphabet_size, br, code));
   return true;
 }
 
+StatusOr<ANSSymbolReader> ANSSymbolReader::Create(const ANSCode* code,
+                                                  BitReader* JXL_RESTRICT br,
+                                                  size_t distance_multiplier) {
+  return ANSSymbolReader(code, br, distance_multiplier);
+}
+
+ANSSymbolReader::ANSSymbolReader(const ANSCode* code,
+                                 BitReader* JXL_RESTRICT br,
+                                 size_t distance_multiplier)
+    : alias_tables_(
+          reinterpret_cast<AliasTable::Entry*>(code->alias_tables.get())),
+      huffman_data_(code->huffman_data.data()),
+      use_prefix_code_(code->use_prefix_code),
+      configs(code->uint_config.data()) {
+  if (!use_prefix_code_) {
+    state_ = static_cast<uint32_t>(br->ReadFixedBits<32>());
+    log_alpha_size_ = code->log_alpha_size;
+    log_entry_size_ = ANS_LOG_TAB_SIZE - code->log_alpha_size;
+    entry_size_minus_1_ = (1 << log_entry_size_) - 1;
+  } else {
+    state_ = (ANS_SIGNATURE << 16u);
+  }
+  if (!code->lz77.enabled) return;
+  lz77_window_storage_ = AllocateArray(kWindowSize * sizeof(uint32_t));
+  lz77_window_ = reinterpret_cast<uint32_t*>(lz77_window_storage_.get());
+  lz77_ctx_ = code->lz77.nonserialized_distance_context;
+  lz77_length_uint_ = code->lz77.length_uint_config;
+  lz77_threshold_ = code->lz77.min_symbol;
+  lz77_min_length_ = code->lz77.min_length;
+  num_special_distances_ = distance_multiplier == 0 ? 0 : kNumSpecialDistances;
+  for (size_t i = 0; i < num_special_distances_; i++) {
+    special_distances_[i] = SpecialDistance(i, distance_multiplier);
+  }
+}
+
 }  // namespace jxl

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

@@ -9,6 +9,7 @@
 // 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>
@@ -152,6 +153,7 @@ struct ANSCode {
   // 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);
 };
 
@@ -159,36 +161,9 @@ class ANSSymbolReader {
  public:
   // Invalid symbol reader, to be overwritten.
   ANSSymbolReader() = default;
-  ANSSymbolReader(const ANSCode* code, BitReader* JXL_RESTRICT br,
+  static StatusOr<ANSSymbolReader> Create(const ANSCode* code,
-                  size_t distance_multiplier = 0)
+                                          BitReader* JXL_RESTRICT br,
-      : alias_tables_(
+                                          size_t distance_multiplier = 0);
-            reinterpret_cast<AliasTable::Entry*>(code->alias_tables.get())),
-        huffman_data_(code->huffman_data.data()),
-        use_prefix_code_(code->use_prefix_code),
-        configs(code->uint_config.data()) {
-    if (!use_prefix_code_) {
-      state_ = static_cast<uint32_t>(br->ReadFixedBits<32>());
-      log_alpha_size_ = code->log_alpha_size;
-      log_entry_size_ = ANS_LOG_TAB_SIZE - code->log_alpha_size;
-      entry_size_minus_1_ = (1 << log_entry_size_) - 1;
-    } else {
-      state_ = (ANS_SIGNATURE << 16u);
-    }
-    if (!code->lz77.enabled) return;
-    // a std::vector incurs unacceptable decoding speed loss because of
-    // initialization.
-    lz77_window_storage_ = AllocateArray(kWindowSize * sizeof(uint32_t));
-    lz77_window_ = reinterpret_cast<uint32_t*>(lz77_window_storage_.get());
-    lz77_ctx_ = code->lz77.nonserialized_distance_context;
-    lz77_length_uint_ = code->lz77.length_uint_config;
-    lz77_threshold_ = code->lz77.min_symbol;
-    lz77_min_length_ = code->lz77.min_length;
-    num_special_distances_ =
-        distance_multiplier == 0 ? 0 : kNumSpecialDistances;
-    for (size_t i = 0; i < num_special_distances_; i++) {
-      special_distances_[i] = SpecialDistance(i, distance_multiplier);
-    }
-  }
 
   JXL_INLINE size_t ReadSymbolANSWithoutRefill(const size_t histo_idx,
                                                BitReader* JXL_RESTRICT br) {
@@ -471,6 +446,9 @@ class ANSSymbolReader {
   }
 
  private:
+  ANSSymbolReader(const ANSCode* code, BitReader* JXL_RESTRICT br,
+                  size_t distance_multiplier);
+
   const AliasTable::Entry* JXL_RESTRICT alias_tables_;  // not owned
   const HuffmanDecodingData* huffman_data_;
   bool use_prefix_code_;
@@ -482,6 +460,8 @@ class ANSSymbolReader {
 
   // LZ77 structures and constants.
   static constexpr size_t kWindowMask = kWindowSize - 1;
+  // a std::vector incurs unacceptable decoding speed loss because of
+  // initialization.
   CacheAlignedUniquePtr lz77_window_storage_;
   uint32_t* lz77_window_ = nullptr;
   uint32_t num_decoded_ = 0;
@@ -495,7 +475,8 @@ class ANSSymbolReader {
   uint32_t num_special_distances_{};
 };
 
-Status DecodeHistograms(BitReader* br, size_t num_contexts, ANSCode* code,
+Status DecodeHistograms(JxlMemoryManager* memory_manager, BitReader* br,
+                        size_t num_contexts, ANSCode* code,
                         std::vector<uint8_t>* context_map,
                         bool disallow_lz77 = false);
 

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

@@ -5,6 +5,8 @@
 
 #include "lib/jxl/dec_cache.h"
 
+#include <jxl/memory_manager.h>
+
 #include "lib/jxl/base/status.h"
 #include "lib/jxl/blending.h"
 #include "lib/jxl/common.h"  // JXL_HIGH_PRECISION
@@ -28,8 +30,10 @@
 namespace jxl {
 
 Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,
+                                           const ImageMetadata* metadata,
                                            ImageBundle* decoded,
                                            PipelineOptions options) {
+  JxlMemoryManager* memory_manager = this->memory_manager();
   size_t num_c = 3 + frame_header.nonserialized_metadata->m.num_extra_channels;
   if (options.render_noise && (frame_header.flags & FrameHeader::kNoise) != 0) {
     num_c += 3;
@@ -37,10 +41,10 @@ Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,
 
   if (frame_header.CanBeReferenced()) {
     // Necessary so that SetInputSizes() can allocate output buffers as needed.
-    frame_storage_for_referencing = ImageBundle(decoded->metadata());
+    frame_storage_for_referencing = ImageBundle(memory_manager, metadata);
   }
 
-  RenderPipeline::Builder builder(num_c);
+  RenderPipeline::Builder builder(memory_manager, num_c);
 
   if (options.use_slow_render_pipeline) {
     builder.UseSimpleImplementation();
@@ -120,7 +124,7 @@ Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,
   }
   if (frame_header.dc_level != 0) {
     builder.AddStage(GetWriteToImage3FStage(
-        &shared_storage.dc_frames[frame_header.dc_level - 1]));
+        memory_manager, &shared_storage.dc_frames[frame_header.dc_level - 1]));
   }
 
   if (frame_header.CanBeReferenced() &&
@@ -131,9 +135,8 @@ Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,
 
   bool has_alpha = false;
   size_t alpha_c = 0;
-  for (size_t i = 0; i < decoded->metadata()->extra_channel_info.size(); i++) {
+  for (size_t i = 0; i < metadata->extra_channel_info.size(); i++) {
-    if (decoded->metadata()->extra_channel_info[i].type ==
+    if (metadata->extra_channel_info[i].type == ExtraChannel::kAlpha) {
-        ExtraChannel::kAlpha) {
       has_alpha = true;
       alpha_c = 3 + i;
       break;
@@ -146,7 +149,7 @@ Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,
     JXL_ASSERT(!frame_header.CanBeReferenced() ||
                frame_header.save_before_color_transform);
     JXL_ASSERT(!options.render_spotcolors ||
-               !decoded->metadata()->Find(ExtraChannel::kSpotColor));
+               !metadata->Find(ExtraChannel::kSpotColor));
     bool is_rgba = (main_output.format.num_channels == 4);
     uint8_t* rgb_output = reinterpret_cast<uint8_t*>(main_output.buffer);
     builder.AddStage(GetFastXYBTosRGB8Stage(rgb_output, main_output.stride,
@@ -186,11 +189,9 @@ Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,
 
     if (options.render_spotcolors &&
         frame_header.nonserialized_metadata->m.Find(ExtraChannel::kSpotColor)) {
-      for (size_t i = 0; i < decoded->metadata()->extra_channel_info.size();
+      for (size_t i = 0; i < metadata->extra_channel_info.size(); i++) {
-           i++) {
         // Don't use Find() because there may be multiple spot color channels.
-        const ExtraChannelInfo& eci =
+        const ExtraChannelInfo& eci = metadata->extra_channel_info[i];
-            decoded->metadata()->extra_channel_info[i];
         if (eci.type == ExtraChannel::kSpotColor) {
           builder.AddStage(GetSpotColorStage(3 + i, eci.spot_color));
         }
@@ -251,9 +252,9 @@ Status PassesDecoderState::PreparePipeline(const FrameHeader& frame_header,
     (void)linear;
 
     if (main_output.callback.IsPresent() || main_output.buffer) {
-      builder.AddStage(GetWriteToOutputStage(main_output, width, height,
+      builder.AddStage(GetWriteToOutputStage(
-                                             has_alpha, unpremul_alpha, alpha_c,
+          main_output, width, height, has_alpha, unpremul_alpha, alpha_c,
-                                             undo_orientation, extra_output));
+          undo_orientation, extra_output, memory_manager));
     } else {
       builder.AddStage(
           GetWriteToImageBundleStage(decoded, output_encoding_info));

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

@@ -7,12 +7,13 @@
 #define LIB_JXL_DEC_CACHE_H_
 
 #include <jxl/decode.h>
+#include <jxl/memory_manager.h>
 #include <jxl/types.h>
-#include <stdint.h>
 
 #include <algorithm>
 #include <atomic>
 #include <cmath>
+#include <cstdint>
 #include <hwy/base.h>  // HWY_ALIGN_MAX
 #include <memory>
 #include <vector>
@@ -86,6 +87,10 @@ struct ImageOutput {
 // Per-frame decoder state. All the images here should be accessed through a
 // group rect (either with block units or pixel units).
 struct PassesDecoderState {
+  explicit PassesDecoderState(JxlMemoryManager* memory_manager)
+      : shared_storage(memory_manager),
+        frame_storage_for_referencing(memory_manager) {}
+
   PassesSharedState shared_storage;
   // Allows avoiding copies for encoder loop.
   const PassesSharedState* JXL_RESTRICT shared = &shared_storage;
@@ -144,7 +149,10 @@ struct PassesDecoderState {
     bool render_noise;
   };
 
-  Status PreparePipeline(const FrameHeader& frame_header, ImageBundle* decoded,
+  JxlMemoryManager* memory_manager() const { return shared->memory_manager; }
+
+  Status PreparePipeline(const FrameHeader& frame_header,
+                         const ImageMetadata* metadata, ImageBundle* decoded,
                          PipelineOptions options);
 
   // Information for colour conversions.
@@ -152,6 +160,7 @@ struct PassesDecoderState {
 
   // Initializes decoder-specific structures using information from *shared.
   Status Init(const FrameHeader& frame_header) {
+    JxlMemoryManager* memory_manager = this->memory_manager();
     x_dm_multiplier = std::pow(1 / (1.25f), frame_header.x_qm_scale - 2.0f);
     b_dm_multiplier = std::pow(1 / (1.25f), frame_header.b_qm_scale - 2.0f);
 
@@ -169,7 +178,8 @@ struct PassesDecoderState {
     if (frame_header.loop_filter.epf_iters > 0) {
       JXL_ASSIGN_OR_RETURN(
           sigma,
-          ImageF::Create(shared->frame_dim.xsize_blocks + 2 * kSigmaPadding,
+          ImageF::Create(memory_manager,
+                         shared->frame_dim.xsize_blocks + 2 * kSigmaPadding,
                          shared->frame_dim.ysize_blocks + 2 * kSigmaPadding));
     }
     return true;
@@ -196,16 +206,17 @@ struct PassesDecoderState {
 // Temp images required for decoding a single group. Reduces memory allocations
 // for large images because we only initialize min(#threads, #groups) instances.
 struct GroupDecCache {
-  Status InitOnce(size_t num_passes, size_t used_acs) {
+  Status InitOnce(JxlMemoryManager* memory_manager, size_t num_passes,
+                  size_t used_acs) {
     for (size_t i = 0; i < num_passes; i++) {
       if (num_nzeroes[i].xsize() == 0) {
         // Allocate enough for a whole group - partial groups on the
         // right/bottom border just use a subset. The valid size is passed via
         // Rect.
 
-        JXL_ASSIGN_OR_RETURN(
+        JXL_ASSIGN_OR_RETURN(num_nzeroes[i],
-            num_nzeroes[i],
+                             Image3I::Create(memory_manager, kGroupDimInBlocks,
-            Image3I::Create(kGroupDimInBlocks, kGroupDimInBlocks));
+                                             kGroupDimInBlocks));
       }
     }
     size_t max_block_area = 0;
@@ -235,11 +246,12 @@ struct GroupDecCache {
     return true;
   }
 
-  Status InitDCBufferOnce() {
+  Status InitDCBufferOnce(JxlMemoryManager* memory_manager) {
     if (dc_buffer.xsize() == 0) {
       JXL_ASSIGN_OR_RETURN(
           dc_buffer,
-          ImageF::Create(kGroupDimInBlocks + kRenderPipelineXOffset * 2,
+          ImageF::Create(memory_manager,
+                         kGroupDimInBlocks + kRenderPipelineXOffset * 2,
                          kGroupDimInBlocks + 4));
     }
     return true;

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

@@ -5,14 +5,14 @@
 
 #include "lib/jxl/dec_context_map.h"
 
+#include <jxl/memory_manager.h>
+
 #include <algorithm>
 #include <cstdint>
 #include <vector>
 
-#include "lib/jxl/ans_params.h"
 #include "lib/jxl/base/status.h"
 #include "lib/jxl/dec_ans.h"
-#include "lib/jxl/entropy_coder.h"
 #include "lib/jxl/inverse_mtf-inl.h"
 
 namespace jxl {
@@ -40,7 +40,8 @@ Status VerifyContextMap(const std::vector<uint8_t>& context_map,
 
 }  // namespace
 
-Status DecodeContextMap(std::vector<uint8_t>* context_map, size_t* num_htrees,
+Status DecodeContextMap(JxlMemoryManager* memory_manager,
+                        std::vector<uint8_t>* context_map, size_t* num_htrees,
                         BitReader* input) {
   bool is_simple = static_cast<bool>(input->ReadFixedBits<1>());
   if (is_simple) {
@@ -61,9 +62,10 @@ Status DecodeContextMap(std::vector<uint8_t>* context_map, size_t* num_htrees,
     // in malicious bitstreams by making every context map require its own
     // context map.
     JXL_RETURN_IF_ERROR(
-        DecodeHistograms(input, 1, &code, &sink_ctx_map,
+        DecodeHistograms(memory_manager, input, 1, &code, &sink_ctx_map,
                          /*disallow_lz77=*/context_map->size() <= 2));
-    ANSSymbolReader reader(&code, input);
+    JXL_ASSIGN_OR_RETURN(ANSSymbolReader reader,
+                         ANSSymbolReader::Create(&code, input));
     size_t i = 0;
     uint32_t maxsym = 0;
     while (i < context_map->size()) {

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

@@ -6,9 +6,10 @@
 #ifndef LIB_JXL_DEC_CONTEXT_MAP_H_
 #define LIB_JXL_DEC_CONTEXT_MAP_H_
 
-#include <stddef.h>
+#include <jxl/memory_manager.h>
-#include <stdint.h>
 
+#include <cstddef>
+#include <cstdint>
 #include <vector>
 
 #include "lib/jxl/dec_bit_reader.h"
@@ -22,7 +23,8 @@ constexpr size_t kMaxClusters = 256;
 // context_map->size() must be the number of possible context ids.
 // Sets *num_htrees to the number of different histogram ids in
 // *context_map.
-Status DecodeContextMap(std::vector<uint8_t>* context_map, size_t* num_htrees,
+Status DecodeContextMap(JxlMemoryManager* memory_manager,
+                        std::vector<uint8_t>* context_map, size_t* num_htrees,
                         BitReader* input);
 
 }  // namespace jxl

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

@@ -5,10 +5,11 @@
 
 #include "lib/jxl/dec_external_image.h"
 
+#include <jxl/memory_manager.h>
 #include <jxl/types.h>
-#include <string.h>
 
 #include <algorithm>
+#include <cstring>
 #include <utility>
 #include <vector>
 
@@ -24,7 +25,7 @@
 #include "lib/jxl/base/common.h"
 #include "lib/jxl/base/compiler_specific.h"
 #include "lib/jxl/base/printf_macros.h"
-#include "lib/jxl/sanitizers.h"
+#include "lib/jxl/base/sanitizers.h"
 
 HWY_BEFORE_NAMESPACE();
 namespace jxl {
@@ -111,9 +112,10 @@ Status UndoOrientation(jxl::Orientation undo_orientation, const Plane<T>& image,
                        Plane<T>& out, jxl::ThreadPool* pool) {
   const size_t xsize = image.xsize();
   const size_t ysize = image.ysize();
+  JxlMemoryManager* memory_manager = image.memory_manager();
 
   if (undo_orientation == Orientation::kFlipHorizontal) {
-    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(memory_manager, xsize, ysize));
     JXL_RETURN_IF_ERROR(RunOnPool(
         pool, 0, static_cast<uint32_t>(ysize), ThreadPool::NoInit,
         [&](const uint32_t task, size_t /*thread*/) {
@@ -126,7 +128,7 @@ Status UndoOrientation(jxl::Orientation undo_orientation, const Plane<T>& image,
         },
         "UndoOrientation"));
   } else if (undo_orientation == Orientation::kRotate180) {
-    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(memory_manager, xsize, ysize));
     JXL_RETURN_IF_ERROR(RunOnPool(
         pool, 0, static_cast<uint32_t>(ysize), ThreadPool::NoInit,
         [&](const uint32_t task, size_t /*thread*/) {
@@ -139,7 +141,7 @@ Status UndoOrientation(jxl::Orientation undo_orientation, const Plane<T>& image,
         },
         "UndoOrientation"));
   } else if (undo_orientation == Orientation::kFlipVertical) {
-    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(memory_manager, xsize, ysize));
     JXL_RETURN_IF_ERROR(RunOnPool(
         pool, 0, static_cast<uint32_t>(ysize), ThreadPool::NoInit,
         [&](const uint32_t task, size_t /*thread*/) {
@@ -152,7 +154,7 @@ Status UndoOrientation(jxl::Orientation undo_orientation, const Plane<T>& image,
         },
         "UndoOrientation"));
   } else if (undo_orientation == Orientation::kTranspose) {
-    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(ysize, xsize));
+    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(memory_manager, ysize, xsize));
     JXL_RETURN_IF_ERROR(RunOnPool(
         pool, 0, static_cast<uint32_t>(ysize), ThreadPool::NoInit,
         [&](const uint32_t task, size_t /*thread*/) {
@@ -164,7 +166,7 @@ Status UndoOrientation(jxl::Orientation undo_orientation, const Plane<T>& image,
         },
         "UndoOrientation"));
   } else if (undo_orientation == Orientation::kRotate90) {
-    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(ysize, xsize));
+    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(memory_manager, ysize, xsize));
     JXL_RETURN_IF_ERROR(RunOnPool(
         pool, 0, static_cast<uint32_t>(ysize), ThreadPool::NoInit,
         [&](const uint32_t task, size_t /*thread*/) {
@@ -176,7 +178,7 @@ Status UndoOrientation(jxl::Orientation undo_orientation, const Plane<T>& image,
         },
         "UndoOrientation"));
   } else if (undo_orientation == Orientation::kAntiTranspose) {
-    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(ysize, xsize));
+    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(memory_manager, ysize, xsize));
     JXL_RETURN_IF_ERROR(RunOnPool(
         pool, 0, static_cast<uint32_t>(ysize), ThreadPool::NoInit,
         [&](const uint32_t task, size_t /*thread*/) {
@@ -188,7 +190,7 @@ Status UndoOrientation(jxl::Orientation undo_orientation, const Plane<T>& image,
         },
         "UndoOrientation"));
   } else if (undo_orientation == Orientation::kRotate270) {
-    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(ysize, xsize));
+    JXL_ASSIGN_OR_RETURN(out, Plane<T>::Create(memory_manager, ysize, xsize));
     JXL_RETURN_IF_ERROR(RunOnPool(
         pool, 0, static_cast<uint32_t>(ysize), ThreadPool::NoInit,
         [&](const uint32_t task, size_t /*thread*/) {
@@ -245,6 +247,7 @@ Status ConvertChannelsToExternal(const ImageF* in_channels[],
                                  jxl::Orientation undo_orientation) {
   JXL_DASSERT(num_channels != 0 && num_channels <= kConvertMaxChannels);
   JXL_DASSERT(in_channels[0] != nullptr);
+  JxlMemoryManager* memory_manager = in_channels[0]->memory_manager();
   JXL_CHECK(float_out ? bits_per_sample == 16 || bits_per_sample == 32
                       : bits_per_sample > 0 && bits_per_sample <= 16);
   const bool has_out_image = (out_image != nullptr);
@@ -310,7 +313,7 @@ Status ConvertChannelsToExternal(const ImageF* in_channels[],
   ImageF ones;
   for (size_t c = 0; c < num_channels; ++c) {
     if (!channels[c]) {
-      JXL_ASSIGN_OR_RETURN(ones, ImageF::Create(xsize, 1));
+      JXL_ASSIGN_OR_RETURN(ones, ImageF::Create(memory_manager, xsize, 1));
       FillImage(1.0f, &ones);
       break;
     }
@@ -324,7 +327,7 @@ Status ConvertChannelsToExternal(const ImageF* in_channels[],
           pool, 0, static_cast<uint32_t>(ysize),
           [&](size_t num_threads) {
             StatusOr<Plane<hwy::float16_t>> f16_cache_or =
-                Plane<hwy::float16_t>::Create(xsize,
+                Plane<hwy::float16_t>::Create(memory_manager, xsize,
                                               num_channels * num_threads);
             if (!f16_cache_or.ok()) return false;
             f16_cache = std::move(f16_cache_or).value();
@@ -402,8 +405,8 @@ Status ConvertChannelsToExternal(const ImageF* in_channels[],
     JXL_RETURN_IF_ERROR(RunOnPool(
         pool, 0, static_cast<uint32_t>(ysize),
         [&](size_t num_threads) {
-          StatusOr<Plane<uint32_t>> u32_cache_or =
+          StatusOr<Plane<uint32_t>> u32_cache_or = Plane<uint32_t>::Create(
-              Plane<uint32_t>::Create(xsize, num_channels * num_threads);
+              memory_manager, xsize, num_channels * num_threads);
           if (!u32_cache_or.ok()) return false;
           u32_cache = std::move(u32_cache_or).value();
           return !!InitOutCallback(num_threads);
@@ -457,11 +460,13 @@ Status ConvertToExternal(const jxl::ImageBundle& ib, size_t bits_per_sample,
   size_t color_channels = num_channels <= 2 ? 1 : 3;
 
   const Image3F* color = &ib.color();
+  JxlMemoryManager* memory_manager = color->memory_manager();
   // Undo premultiplied alpha.
   Image3F unpremul;
   if (ib.AlphaIsPremultiplied() && ib.HasAlpha() && unpremul_alpha) {
-    JXL_ASSIGN_OR_RETURN(unpremul,
+    JXL_ASSIGN_OR_RETURN(
-                         Image3F::Create(color->xsize(), color->ysize()));
+        unpremul,
+        Image3F::Create(memory_manager, color->xsize(), color->ysize()));
     CopyImageTo(*color, &unpremul);
     for (size_t y = 0; y < unpremul.ysize(); y++) {
       UnpremultiplyAlpha(unpremul.PlaneRow(0, y), unpremul.PlaneRow(1, y),

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

@@ -3,16 +3,20 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+#include <jxl/memory_manager.h>
+
 #include "benchmark/benchmark.h"
 #include "lib/jxl/dec_external_image.h"
 #include "lib/jxl/image.h"
 #include "lib/jxl/image_ops.h"
+#include "tools/no_memory_manager.h"
 
 namespace jxl {
 namespace {
 
 // Decoder case, interleaves an internal float image.
 void BM_DecExternalImage_ConvertImageRGBA(benchmark::State& state) {
+  JxlMemoryManager* memory_manager = jpegxl::tools::NoMemoryManager();
   const size_t kNumIter = 5;
   size_t xsize = state.range();
   size_t ysize = state.range();
@@ -20,11 +24,12 @@ void BM_DecExternalImage_ConvertImageRGBA(benchmark::State& state) {
 
   ImageMetadata im;
   im.SetAlphaBits(8);
-  ImageBundle ib(&im);
+  ImageBundle ib(memory_manager, &im);
-  JXL_ASSIGN_OR_DIE(Image3F color, Image3F::Create(xsize, ysize));
+  JXL_ASSIGN_OR_DIE(Image3F color,
+                    Image3F::Create(memory_manager, xsize, ysize));
   ZeroFillImage(&color);
   ib.SetFromImage(std::move(color), ColorEncoding::SRGB());
-  JXL_ASSIGN_OR_DIE(ImageF alpha, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_DIE(ImageF alpha, ImageF::Create(memory_manager, xsize, ysize));
   ZeroFillImage(&alpha);
   ib.SetAlpha(std::move(alpha));
 

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

@@ -6,11 +6,12 @@
 #include "lib/jxl/dec_frame.h"
 
 #include <jxl/decode.h>
-#include <stddef.h>
+#include <jxl/memory_manager.h>
-#include <stdint.h>
 
 #include <algorithm>
 #include <atomic>
+#include <cstddef>
+#include <cstdint>
 #include <cstdlib>
 #include <memory>
 #include <utility>
@@ -63,8 +64,8 @@ Status DecodeGlobalDCInfo(BitReader* reader, bool is_jpeg,
                           PassesDecoderState* state, ThreadPool* pool) {
   JXL_RETURN_IF_ERROR(state->shared_storage.quantizer.Decode(reader));
 
-  JXL_RETURN_IF_ERROR(
+  JXL_RETURN_IF_ERROR(DecodeBlockCtxMap(state->memory_manager(), reader,
-      DecodeBlockCtxMap(reader, &state->shared_storage.block_ctx_map));
+                                        &state->shared_storage.block_ctx_map));
 
   JXL_RETURN_IF_ERROR(state->shared_storage.cmap.DecodeDC(reader));
 
@@ -135,6 +136,7 @@ Status FrameDecoder::InitFrame(BitReader* JXL_RESTRICT br, ImageBundle* decoded,
                                bool is_preview) {
   decoded_ = decoded;
   JXL_ASSERT(is_finalized_);
+  JxlMemoryManager* memory_manager = decoded_->memory_manager();
 
   // Reset the dequantization matrices to their default values.
   dec_state_->shared_storage.matrices = DequantMatrices();
@@ -170,7 +172,8 @@ Status FrameDecoder::InitFrame(BitReader* JXL_RESTRICT br, ImageBundle* decoded,
       NumTocEntries(num_groups, frame_dim_.num_dc_groups, num_passes);
   std::vector<uint32_t> sizes;
   std::vector<coeff_order_t> permutation;
-  JXL_RETURN_IF_ERROR(ReadToc(toc_entries, br, &sizes, &permutation));
+  JXL_RETURN_IF_ERROR(
+      ReadToc(memory_manager, toc_entries, br, &sizes, &permutation));
   bool have_permutation = !permutation.empty();
   toc_.resize(toc_entries);
   section_sizes_sum_ = 0;
@@ -264,10 +267,11 @@ Status FrameDecoder::InitFrameOutput() {
 
 Status FrameDecoder::ProcessDCGlobal(BitReader* br) {
   PassesSharedState& shared = dec_state_->shared_storage;
+  JxlMemoryManager* memory_manager = shared.memory_manager;
   if (frame_header_.flags & FrameHeader::kPatches) {
     bool uses_extra_channels = false;
     JXL_RETURN_IF_ERROR(shared.image_features.patches.Decode(
-        br, frame_dim_.xsize_padded, frame_dim_.ysize_padded,
+        memory_manager, br, frame_dim_.xsize_padded, frame_dim_.ysize_padded,
         &uses_extra_channels));
     if (uses_extra_channels && frame_header_.upsampling != 1) {
       for (size_t ecups : frame_header_.extra_channel_upsampling) {
@@ -284,7 +288,7 @@ Status FrameDecoder::ProcessDCGlobal(BitReader* br) {
   shared.image_features.splines.Clear();
   if (frame_header_.flags & FrameHeader::kSplines) {
     JXL_RETURN_IF_ERROR(shared.image_features.splines.Decode(
-        br, frame_dim_.xsize * frame_dim_.ysize));
+        memory_manager, br, frame_dim_.xsize * frame_dim_.ysize));
   }
   if (frame_header_.flags & FrameHeader::kNoise) {
     JXL_RETURN_IF_ERROR(DecodeNoise(br, &shared.image_features.noise_params));
@@ -339,12 +343,13 @@ Status FrameDecoder::ProcessDCGroup(size_t dc_group_id, BitReader* br) {
 Status FrameDecoder::FinalizeDC() {
   // Do Adaptive DC smoothing if enabled. This *must* happen between all the
   // ProcessDCGroup and ProcessACGroup.
+  JxlMemoryManager* memory_manager = dec_state_->memory_manager();
   if (frame_header_.encoding == FrameEncoding::kVarDCT &&
       !(frame_header_.flags & FrameHeader::kSkipAdaptiveDCSmoothing) &&
       !(frame_header_.flags & FrameHeader::kUseDcFrame)) {
-    JXL_RETURN_IF_ERROR(
+    JXL_RETURN_IF_ERROR(AdaptiveDCSmoothing(
-        AdaptiveDCSmoothing(dec_state_->shared->quantizer.MulDC(),
+        memory_manager, dec_state_->shared->quantizer.MulDC(),
-                            &dec_state_->shared_storage.dc_storage, pool_));
+        &dec_state_->shared_storage.dc_storage, pool_));
   }
 
   finalized_dc_ = true;
@@ -363,11 +368,12 @@ Status FrameDecoder::AllocateOutput() {
 
 Status FrameDecoder::ProcessACGlobal(BitReader* br) {
   JXL_CHECK(finalized_dc_);
+  JxlMemoryManager* memory_manager = dec_state_->memory_manager();
 
   // Decode AC group.
   if (frame_header_.encoding == FrameEncoding::kVarDCT) {
     JXL_RETURN_IF_ERROR(dec_state_->shared_storage.matrices.Decode(
-        br, &modular_frame_decoder_));
+        memory_manager, br, &modular_frame_decoder_));
     JXL_RETURN_IF_ERROR(dec_state_->shared_storage.matrices.EnsureComputed(
         dec_state_->used_acs));
 
@@ -389,15 +395,16 @@ Status FrameDecoder::ProcessACGlobal(BitReader* br) {
     for (size_t i = 0; i < frame_header_.passes.num_passes; i++) {
       uint16_t used_orders = U32Coder::Read(kOrderEnc, br);
       JXL_RETURN_IF_ERROR(DecodeCoeffOrders(
-          used_orders, dec_state_->used_acs,
+          memory_manager, used_orders, dec_state_->used_acs,
           &dec_state_->shared_storage
                .coeff_orders[i * dec_state_->shared_storage.coeff_order_size],
           br));
       size_t num_contexts =
           dec_state_->shared->num_histograms *
           dec_state_->shared_storage.block_ctx_map.NumACContexts();
-      JXL_RETURN_IF_ERROR(DecodeHistograms(
+      JXL_RETURN_IF_ERROR(DecodeHistograms(memory_manager, br, num_contexts,
-          br, num_contexts, &dec_state_->code[i], &dec_state_->context_map[i]));
+                                           &dec_state_->code[i],
+                                           &dec_state_->context_map[i]));
       // Add extra values to enable the cheat in hot loop of DecodeACVarBlock.
       dec_state_->context_map[i].resize(
           num_contexts + kZeroDensityContextLimit - kZeroDensityContextCount);
@@ -414,10 +421,10 @@ Status FrameDecoder::ProcessACGlobal(BitReader* br) {
     size_t ys = store ? frame_dim_.num_groups : 0;
     if (use_16_bit) {
       JXL_ASSIGN_OR_RETURN(dec_state_->coefficients,
-                           ACImageT<int16_t>::Make(xs, ys));
+                           ACImageT<int16_t>::Make(memory_manager, xs, ys));
     } else {
       JXL_ASSIGN_OR_RETURN(dec_state_->coefficients,
-                           ACImageT<int32_t>::Make(xs, ys));
+                           ACImageT<int32_t>::Make(memory_manager, xs, ys));
     }
     if (store) {
       dec_state_->coefficients->ZeroFill();
@@ -475,6 +482,7 @@ Status FrameDecoder::ProcessACGroup(size_t ac_group_id,
   const size_t gy = ac_group_id / frame_dim_.xsize_groups;
   const size_t x = gx * group_dim;
   const size_t y = gy * group_dim;
+  JxlMemoryManager* memory_manager = dec_state_->memory_manager();
   JXL_DEBUG_V(3,
               "Processing AC group %" PRIuS "(%" PRIuS ",%" PRIuS
               ") group_dim: %" PRIuS " decoded passes: %u new passes: %" PRIuS,
@@ -488,12 +496,12 @@ Status FrameDecoder::ProcessACGroup(size_t ac_group_id,
 
   if (frame_header_.encoding == FrameEncoding::kVarDCT) {
     JXL_RETURN_IF_ERROR(group_dec_caches_[thread].InitOnce(
-        frame_header_.passes.num_passes, dec_state_->used_acs));
+        memory_manager, frame_header_.passes.num_passes, dec_state_->used_acs));
-    JXL_RETURN_IF_ERROR(DecodeGroup(frame_header_, br, num_passes, ac_group_id,
+    JXL_RETURN_IF_ERROR(DecodeGroup(
-                                    dec_state_, &group_dec_caches_[thread],
+        frame_header_, br, num_passes, ac_group_id, dec_state_,
-                                    thread, render_pipeline_input, decoded_,
+        &group_dec_caches_[thread], thread, render_pipeline_input,
-                                    decoded_passes_per_ac_group_[ac_group_id],
+        decoded_->jpeg_data.get(), decoded_passes_per_ac_group_[ac_group_id],
-                                    force_draw, dc_only, &should_run_pipeline));
+        force_draw, dc_only, &should_run_pipeline));
   }
 
   // don't limit to image dimensions here (is done in DecodeGroup)
@@ -676,8 +684,9 @@ Status FrameDecoder::ProcessSections(const SectionInfo* sections, size_t num,
     pipeline_options.coalescing = coalescing_;
     pipeline_options.render_spotcolors = render_spotcolors_;
     pipeline_options.render_noise = true;
-    JXL_RETURN_IF_ERROR(
+    JXL_RETURN_IF_ERROR(dec_state_->PreparePipeline(
-        dec_state_->PreparePipeline(frame_header_, decoded_, pipeline_options));
+        frame_header_, &frame_header_.nonserialized_metadata->m, decoded_,
+        pipeline_options));
     JXL_RETURN_IF_ERROR(FinalizeDC());
     JXL_RETURN_IF_ERROR(AllocateOutput());
     if (progressive_detail_ >= JxlProgressiveDetail::kDC) {
@@ -892,7 +901,7 @@ Status FrameDecoder::FinalizeFrame() {
   if (frame_header_.CanBeReferenced()) {
     auto& info = dec_state_->shared_storage
                      .reference_frames[frame_header_.save_as_reference];
-    info.frame = std::move(dec_state_->frame_storage_for_referencing);
+    *info.frame = std::move(dec_state_->frame_storage_for_referencing);
     info.ib_is_in_xyb = frame_header_.save_before_color_transform;
   }
   return true;

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

@@ -8,10 +8,10 @@
 
 #include <jxl/decode.h>
 #include <jxl/types.h>
-#include <stdint.h>
 
 #include <algorithm>
 #include <cstddef>
+#include <cstdint>
 #include <limits>
 #include <utility>
 #include <vector>
@@ -49,6 +49,7 @@ class FrameDecoder {
       : dec_state_(dec_state),
         pool_(pool),
         frame_header_(&metadata),
+        modular_frame_decoder_(dec_state_->memory_manager()),
         use_slow_rendering_pipeline_(use_slow_rendering_pipeline) {}
 
   void SetRenderSpotcolors(bool rsc) { render_spotcolors_ = rsc; }

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

@@ -170,7 +170,7 @@ Status DecodeGroupImpl(const FrameHeader& frame_header,
                        PassesDecoderState* JXL_RESTRICT dec_state,
                        size_t thread, size_t group_idx,
                        RenderPipelineInput& render_pipeline_input,
-                       ImageBundle* decoded, DrawMode draw) {
+                       jpeg::JPEGData* jpeg_data, DrawMode draw) {
   // TODO(veluca): investigate cache usage in this function.
   const Rect block_rect =
       dec_state->shared->frame_dim.BlockGroupRect(group_idx);
@@ -209,11 +209,11 @@ Status DecodeGroupImpl(const FrameHeader& frame_header,
   std::array<int, 3> dcoff = {};
 
   // TODO(veluca): all of this should be done only once per image.
-  if (decoded->IsJPEG()) {
+  if (jpeg_data) {
     if (!dec_state->shared->cmap.IsJPEGCompatible()) {
       return JXL_FAILURE("The CfL map is not JPEG-compatible");
     }
-    jpeg_is_gray = (decoded->jpeg_data->components.size() == 1);
+    jpeg_is_gray = (jpeg_data->components.size() == 1);
     jpeg_c_map = JpegOrder(frame_header.color_transform, jpeg_is_gray);
     const std::vector<QuantEncoding>& qe =
         dec_state->shared->matrices.encodings();
@@ -279,8 +279,8 @@ Status DecodeGroupImpl(const FrameHeader& frame_header,
     for (size_t c = 0; c < 3; c++) {
       idct_row[c] = render_pipeline_input.GetBuffer(c).second.Row(
           render_pipeline_input.GetBuffer(c).first, sby[c] * kBlockDim);
-      if (decoded->IsJPEG()) {
+      if (jpeg_data) {
-        auto& component = decoded->jpeg_data->components[jpeg_c_map[c]];
+        auto& component = jpeg_data->components[jpeg_c_map[c]];
         jpeg_row[c] =
             component.coeffs.data() +
             (component.width_in_blocks * (r[c].y0() + sby[c]) + r[c].x0()) *
@@ -344,7 +344,7 @@ Status DecodeGroupImpl(const FrameHeader& frame_header,
           continue;
         }
 
-        if (JXL_UNLIKELY(decoded->IsJPEG())) {
+        if (JXL_UNLIKELY(jpeg_data)) {
           if (acs.Strategy() != AcStrategy::Type::DCT) {
             return JXL_FAILURE(
                 "Can only decode to JPEG if only DCT-8 is used.");
@@ -606,8 +606,10 @@ struct GetBlockFromBitstream : public GetBlock {
       }
       ctx_offset[pass] = cur_histogram * block_ctx_map->NumACContexts();
 
-      decoders[pass] =
+      JXL_ASSIGN_OR_RETURN(
-          ANSSymbolReader(&dec_state->code[pass + first_pass], readers[pass]);
+          decoders[pass],
+          ANSSymbolReader::Create(&dec_state->code[pass + first_pass],
+                                  readers[pass]));
     }
     nzeros_stride = group_dec_cache->num_nzeroes[0].PixelsPerRow();
     for (size_t i = 0; i < num_passes; i++) {
@@ -688,8 +690,9 @@ Status DecodeGroup(const FrameHeader& frame_header,
                    PassesDecoderState* JXL_RESTRICT dec_state,
                    GroupDecCache* JXL_RESTRICT group_dec_cache, size_t thread,
                    RenderPipelineInput& render_pipeline_input,
-                   ImageBundle* JXL_RESTRICT decoded, size_t first_pass,
+                   jpeg::JPEGData* JXL_RESTRICT jpeg_data, size_t first_pass,
                    bool force_draw, bool dc_only, bool* should_run_pipeline) {
+  JxlMemoryManager* memory_manager = dec_state->memory_manager();
   DrawMode draw =
       (num_passes + first_pass == frame_header.passes.num_passes) || force_draw
           ? kDraw
@@ -700,7 +703,7 @@ Status DecodeGroup(const FrameHeader& frame_header,
   }
 
   if (draw == kDraw && num_passes == 0 && first_pass == 0) {
-    JXL_RETURN_IF_ERROR(group_dec_cache->InitDCBufferOnce());
+    JXL_RETURN_IF_ERROR(group_dec_cache->InitDCBufferOnce(memory_manager));
     const YCbCrChromaSubsampling& cs = frame_header.chroma_subsampling;
     for (size_t c : {0, 1, 2}) {
       size_t hs = cs.HShift(c);
@@ -777,7 +780,7 @@ Status DecodeGroup(const FrameHeader& frame_header,
 
   JXL_RETURN_IF_ERROR(HWY_DYNAMIC_DISPATCH(DecodeGroupImpl)(
       frame_header, get_block.get(), group_dec_cache, dec_state, thread,
-      group_idx, render_pipeline_input, decoded, draw));
+      group_idx, render_pipeline_input, jpeg_data, draw));
 
   for (size_t pass = 0; pass < num_passes; pass++) {
     if (!get_block->decoders[pass].CheckANSFinalState()) {
@@ -794,16 +797,18 @@ Status DecodeGroupForRoundtrip(const FrameHeader& frame_header,
                                GroupDecCache* JXL_RESTRICT group_dec_cache,
                                size_t thread,
                                RenderPipelineInput& render_pipeline_input,
-                               ImageBundle* JXL_RESTRICT decoded,
+                               jpeg::JPEGData* JXL_RESTRICT jpeg_data,
                                AuxOut* aux_out) {
+  JxlMemoryManager* memory_manager = dec_state->memory_manager();
   GetBlockFromEncoder get_block(ac, group_idx, frame_header.passes.shift);
   JXL_RETURN_IF_ERROR(group_dec_cache->InitOnce(
+      memory_manager,
       /*num_passes=*/0,
       /*used_acs=*/(1u << AcStrategy::kNumValidStrategies) - 1));
 
   return HWY_DYNAMIC_DISPATCH(DecodeGroupImpl)(
       frame_header, &get_block, group_dec_cache, dec_state, thread, group_idx,
-      render_pipeline_input, decoded, kDraw);
+      render_pipeline_input, jpeg_data, kDraw);
 }
 
 }  // namespace jxl

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

@@ -29,7 +29,7 @@ Status DecodeGroup(const FrameHeader& frame_header,
                    PassesDecoderState* JXL_RESTRICT dec_state,
                    GroupDecCache* JXL_RESTRICT group_dec_cache, size_t thread,
                    RenderPipelineInput& render_pipeline_input,
-                   ImageBundle* JXL_RESTRICT decoded, size_t first_pass,
+                   jpeg::JPEGData* JXL_RESTRICT jpeg_data, size_t first_pass,
                    bool force_draw, bool dc_only, bool* should_run_pipeline);
 
 Status DecodeGroupForRoundtrip(const FrameHeader& frame_header,
@@ -39,7 +39,7 @@ Status DecodeGroupForRoundtrip(const FrameHeader& frame_header,
                                GroupDecCache* JXL_RESTRICT group_dec_cache,
                                size_t thread,
                                RenderPipelineInput& render_pipeline_input,
-                               ImageBundle* JXL_RESTRICT decoded,
+                               jpeg::JPEGData* JXL_RESTRICT jpeg_data,
                                AuxOut* aux_out);
 
 }  // namespace jxl

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

@@ -5,6 +5,8 @@
 
 #include "lib/jxl/dec_modular.h"
 
+#include <jxl/memory_manager.h>
+
 #include <atomic>
 #include <cstdint>
 #include <vector>
@@ -177,6 +179,7 @@ std::string ModularStreamId::DebugString() const {
 Status ModularFrameDecoder::DecodeGlobalInfo(BitReader* reader,
                                              const FrameHeader& frame_header,
                                              bool allow_truncated_group) {
+  JxlMemoryManager* memory_manager = this->memory_manager();
   bool decode_color = frame_header.encoding == FrameEncoding::kModular;
   const auto& metadata = frame_header.nonserialized_metadata->m;
   bool is_gray = metadata.color_encoding.IsGray();
@@ -194,9 +197,10 @@ Status ModularFrameDecoder::DecodeGlobalInfo(BitReader* reader,
           std::min(static_cast<size_t>(1 << 22),
                    1024 + frame_dim.xsize * frame_dim.ysize *
                               (nb_chans + nb_extra) / 16);
-      JXL_RETURN_IF_ERROR(DecodeTree(reader, &tree, tree_size_limit));
       JXL_RETURN_IF_ERROR(
-          DecodeHistograms(reader, (tree.size() + 1) / 2, &code, &context_map));
+          DecodeTree(memory_manager, reader, &tree, tree_size_limit));
+      JXL_RETURN_IF_ERROR(DecodeHistograms(
+          memory_manager, reader, (tree.size() + 1) / 2, &code, &context_map));
     }
   }
   if (!do_color) nb_chans = 0;
@@ -215,7 +219,7 @@ Status ModularFrameDecoder::DecodeGlobalInfo(BitReader* reader,
 
   JXL_ASSIGN_OR_RETURN(
       Image gi,
-      Image::Create(frame_dim.xsize, frame_dim.ysize,
+      Image::Create(memory_manager, frame_dim.xsize, frame_dim.ysize,
                     metadata.bit_depth.bits_per_sample, nb_chans + nb_extra));
 
   all_same_shift = true;
@@ -306,8 +310,8 @@ Status ModularFrameDecoder::DecodeGroup(
               stream.kind == ModularStreamId::kModularAC);
   const size_t xsize = rect.xsize();
   const size_t ysize = rect.ysize();
-  JXL_ASSIGN_OR_RETURN(Image gi,
+  JXL_ASSIGN_OR_RETURN(Image gi, Image::Create(memory_manager_, xsize, ysize,
-                       Image::Create(xsize, ysize, full_image.bitdepth, 0));
+                                               full_image.bitdepth, 0));
   // start at the first bigger-than-groupsize non-metachannel
   size_t c = full_image.nb_meta_channels;
   for (; c < full_image.channel.size(); c++) {
@@ -329,7 +333,8 @@ Status ModularFrameDecoder::DecodeGroup(
         memset(row_out, 0, r.xsize() * sizeof(*row_out));
       }
     } else {
-      JXL_ASSIGN_OR_RETURN(Channel gc, Channel::Create(r.xsize(), r.ysize()));
+      JXL_ASSIGN_OR_RETURN(
+          Channel gc, Channel::Create(memory_manager_, r.xsize(), r.ysize()));
       if (zerofill) ZeroFillImage(&gc.plane);
       gc.hshift = fc.hshift;
       gc.vshift = fc.vshift;
@@ -391,6 +396,7 @@ Status ModularFrameDecoder::DecodeGroup(
 Status ModularFrameDecoder::DecodeVarDCTDC(const FrameHeader& frame_header,
                                            size_t group_id, BitReader* reader,
                                            PassesDecoderState* dec_state) {
+  JxlMemoryManager* memory_manager = dec_state->memory_manager();
   const Rect r = dec_state->shared->frame_dim.DCGroupRect(group_id);
   JXL_DEBUG_V(6, "Decoding VarDCT DC with rect %s", Description(r).c_str());
   // TODO(eustas): investigate if we could reduce the impact of
@@ -399,8 +405,9 @@ Status ModularFrameDecoder::DecodeVarDCTDC(const FrameHeader& frame_header,
   //               3 comes from XybToRgb that cubes the values, and "magic" is
   //               the sum of all other contributions. 2**18 is known to lead
   //               to NaN on input found by fuzzing (see commit message).
-  JXL_ASSIGN_OR_RETURN(
+  JXL_ASSIGN_OR_RETURN(Image image,
-      Image image, Image::Create(r.xsize(), r.ysize(), full_image.bitdepth, 3));
+                       Image::Create(memory_manager, r.xsize(), r.ysize(),
+                                     full_image.bitdepth, 3));
   size_t stream_id = ModularStreamId::VarDCTDC(group_id).ID(frame_dim);
   reader->Refill();
   size_t extra_precision = reader->ReadFixedBits<2>();
@@ -429,6 +436,7 @@ Status ModularFrameDecoder::DecodeVarDCTDC(const FrameHeader& frame_header,
 Status ModularFrameDecoder::DecodeAcMetadata(const FrameHeader& frame_header,
                                              size_t group_id, BitReader* reader,
                                              PassesDecoderState* dec_state) {
+  JxlMemoryManager* memory_manager = dec_state->memory_manager();
   const Rect r = dec_state->shared->frame_dim.DCGroupRect(group_id);
   JXL_DEBUG_V(6, "Decoding AcMetadata with rect %s", Description(r).c_str());
   size_t upper_bound = r.xsize() * r.ysize();
@@ -436,15 +444,19 @@ Status ModularFrameDecoder::DecodeAcMetadata(const FrameHeader& frame_header,
   size_t count = reader->ReadBits(CeilLog2Nonzero(upper_bound)) + 1;
   size_t stream_id = ModularStreamId::ACMetadata(group_id).ID(frame_dim);
   // YToX, YToB, ACS + QF, EPF
-  JXL_ASSIGN_OR_RETURN(
+  JXL_ASSIGN_OR_RETURN(Image image,
-      Image image, Image::Create(r.xsize(), r.ysize(), full_image.bitdepth, 4));
+                       Image::Create(memory_manager, r.xsize(), r.ysize(),
+                                     full_image.bitdepth, 4));
   static_assert(kColorTileDimInBlocks == 8, "Color tile size changed");
   Rect cr(r.x0() >> 3, r.y0() >> 3, (r.xsize() + 7) >> 3, (r.ysize() + 7) >> 3);
-  JXL_ASSIGN_OR_RETURN(image.channel[0],
+  JXL_ASSIGN_OR_RETURN(
-                       Channel::Create(cr.xsize(), cr.ysize(), 3, 3));
+      image.channel[0],
-  JXL_ASSIGN_OR_RETURN(image.channel[1],
+      Channel::Create(memory_manager, cr.xsize(), cr.ysize(), 3, 3));
-                       Channel::Create(cr.xsize(), cr.ysize(), 3, 3));
+  JXL_ASSIGN_OR_RETURN(
-  JXL_ASSIGN_OR_RETURN(image.channel[2], Channel::Create(count, 2, 0, 0));
+      image.channel[1],
+      Channel::Create(memory_manager, cr.xsize(), cr.ysize(), 3, 3));
+  JXL_ASSIGN_OR_RETURN(image.channel[2],
+                       Channel::Create(memory_manager, count, 2, 0, 0));
   ModularOptions options;
   if (!ModularGenericDecompress(
           reader, image, /*header=*/nullptr, stream_id, &options,
@@ -692,7 +704,8 @@ Status ModularFrameDecoder::FinalizeDecoding(const FrameHeader& frame_header,
                                              jxl::ThreadPool* pool,
                                              bool inplace) {
   if (!use_full_image) return true;
-  Image gi;
+  JxlMemoryManager* memory_manager = dec_state->memory_manager();
+  Image gi{memory_manager};
   if (inplace) {
     gi = std::move(full_image);
   } else {
@@ -747,8 +760,8 @@ Status ModularFrameDecoder::FinalizeDecoding(const FrameHeader& frame_header,
 static constexpr const float kAlmostZero = 1e-8f;
 
 Status ModularFrameDecoder::DecodeQuantTable(
-    size_t required_size_x, size_t required_size_y, BitReader* br,
+    JxlMemoryManager* memory_manager, size_t required_size_x,
-    QuantEncoding* encoding, size_t idx,
+    size_t required_size_y, BitReader* br, QuantEncoding* encoding, size_t idx,
     ModularFrameDecoder* modular_frame_decoder) {
   JXL_RETURN_IF_ERROR(F16Coder::Read(br, &encoding->qraw.qtable_den));
   if (encoding->qraw.qtable_den < kAlmostZero) {
@@ -756,8 +769,9 @@ Status ModularFrameDecoder::DecodeQuantTable(
     // be negative.
     return JXL_FAILURE("Invalid qtable_den: value too small");
   }
-  JXL_ASSIGN_OR_RETURN(Image image,
+  JXL_ASSIGN_OR_RETURN(
-                       Image::Create(required_size_x, required_size_y, 8, 3));
+      Image image,
+      Image::Create(memory_manager, required_size_x, required_size_y, 8, 3));
   ModularOptions options;
   if (modular_frame_decoder) {
     JXL_RETURN_IF_ERROR(ModularGenericDecompress(

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

@@ -6,6 +6,8 @@
 #ifndef LIB_JXL_DEC_MODULAR_H_
 #define LIB_JXL_DEC_MODULAR_H_
 
+#include <jxl/memory_manager.h>
+
 #include <cstddef>
 #include <cstdint>
 #include <string>
@@ -91,6 +93,8 @@ struct ModularStreamId {
 
 class ModularFrameDecoder {
  public:
+  explicit ModularFrameDecoder(JxlMemoryManager* memory_manager)
+      : memory_manager_(memory_manager), full_image(memory_manager) {}
   void Init(const FrameDimensions& frame_dim) { this->frame_dim = frame_dim; }
   Status DecodeGlobalInfo(BitReader* reader, const FrameHeader& frame_header,
                           bool allow_truncated_group);
@@ -109,7 +113,8 @@ class ModularFrameDecoder {
   // Decodes a RAW quant table from `br` into the given `encoding`, of size
   // `required_size_x x required_size_y`. If `modular_frame_decoder` is passed,
   // its global tree is used, otherwise no global tree is used.
-  static Status DecodeQuantTable(size_t required_size_x, size_t required_size_y,
+  static Status DecodeQuantTable(JxlMemoryManager* memory_manager,
+                                 size_t required_size_x, size_t required_size_y,
                                  BitReader* br, QuantEncoding* encoding,
                                  size_t idx,
                                  ModularFrameDecoder* modular_frame_decoder);
@@ -122,6 +127,7 @@ class ModularFrameDecoder {
   bool have_dc() const { return have_something; }
   void MaybeDropFullImage();
   bool UsesFullImage() const { return use_full_image; }
+  JxlMemoryManager* memory_manager() const { return memory_manager_; }
 
  private:
   Status ModularImageToDecodedRect(const FrameHeader& frame_header, Image& gi,
@@ -129,7 +135,7 @@ class ModularFrameDecoder {
                                    jxl::ThreadPool* pool,
                                    RenderPipelineInput& render_pipeline_input,
                                    Rect modular_rect) const;
-
+  JxlMemoryManager* memory_manager_;
   Image full_image;
   std::vector<Transform> global_transform;
   FrameDimensions frame_dim;

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

@@ -5,11 +5,12 @@
 
 #include "lib/jxl/dec_patch_dictionary.h"
 
-#include <stdint.h>
+#include <jxl/memory_manager.h>
-#include <stdlib.h>
 #include <sys/types.h>
 
 #include <algorithm>
+#include <cstdint>
+#include <cstdlib>
 #include <utility>
 #include <vector>
 
@@ -25,14 +26,16 @@
 
 namespace jxl {
 
-Status PatchDictionary::Decode(BitReader* br, size_t xsize, size_t ysize,
+Status PatchDictionary::Decode(JxlMemoryManager* memory_manager, BitReader* br,
+                               size_t xsize, size_t ysize,
                                bool* uses_extra_channels) {
   positions_.clear();
   std::vector<uint8_t> context_map;
   ANSCode code;
-  JXL_RETURN_IF_ERROR(
+  JXL_RETURN_IF_ERROR(DecodeHistograms(
-      DecodeHistograms(br, kNumPatchDictionaryContexts, &code, &context_map));
+      memory_manager, br, kNumPatchDictionaryContexts, &code, &context_map));
-  ANSSymbolReader decoder(&code, br);
+  JXL_ASSIGN_OR_RETURN(ANSSymbolReader decoder,
+                       ANSSymbolReader::Create(&code, br));
 
   auto read_num = [&](size_t context) {
     size_t r = decoder.ReadHybridUint(context, br, context_map);
@@ -58,14 +61,14 @@ Status PatchDictionary::Decode(BitReader* br, size_t xsize, size_t ysize,
     PatchReferencePosition ref_pos;
     ref_pos.ref = read_num(kReferenceFrameContext);
     if (ref_pos.ref >= kMaxNumReferenceFrames ||
-        shared_->reference_frames[ref_pos.ref].frame.xsize() == 0) {
+        shared_->reference_frames[ref_pos.ref].frame->xsize() == 0) {
       return JXL_FAILURE("Invalid reference frame ID");
     }
     if (!shared_->reference_frames[ref_pos.ref].ib_is_in_xyb) {
       return JXL_FAILURE(
           "Patches cannot use frames saved post color transforms");
     }
-    const ImageBundle& ib = shared_->reference_frames[ref_pos.ref].frame;
+    const ImageBundle& ib = *shared_->reference_frames[ref_pos.ref].frame;
     ref_pos.x0 = read_num(kPatchReferencePositionContext);
     ref_pos.y0 = read_num(kPatchReferencePositionContext);
     ref_pos.xsize = read_num(kPatchSizeContext) + 1;
@@ -318,6 +321,7 @@ Status PatchDictionary::AddOneRow(float* const* inout, size_t y, size_t x0,
                                   size_t xsize) const {
   size_t num_ec = shared_->metadata->m.num_extra_channels;
   std::vector<const float*> fg_ptrs(3 + num_ec);
+  JxlMemoryManager* memory_manager = shared_->memory_manager;
   for (size_t pos_idx : GetPatchesForRow(y)) {
     const size_t blending_idx = pos_idx * (num_ec + 1);
     const PatchPosition& pos = positions_[pos_idx];
@@ -334,19 +338,20 @@ Status PatchDictionary::AddOneRow(float* const* inout, size_t y, size_t x0,
     size_t patch_x0 = std::max(bx, x0);
     size_t patch_x1 = std::min(bx + patch_xsize, x0 + xsize);
     for (size_t c = 0; c < 3; c++) {
-      fg_ptrs[c] = shared_->reference_frames[ref].frame.color().ConstPlaneRow(
+      fg_ptrs[c] = shared_->reference_frames[ref].frame->color()->ConstPlaneRow(
                        c, ref_pos.y0 + iy) +
                    ref_pos.x0 + x0 - bx;
     }
     for (size_t i = 0; i < num_ec; i++) {
       fg_ptrs[3 + i] =
-          shared_->reference_frames[ref].frame.extra_channels()[i].ConstRow(
+          shared_->reference_frames[ref].frame->extra_channels()[i].ConstRow(
               ref_pos.y0 + iy) +
           ref_pos.x0 + x0 - bx;
     }
     JXL_RETURN_IF_ERROR(PerformBlending(
-        inout, fg_ptrs.data(), inout, patch_x0 - x0, patch_x1 - patch_x0,
+        memory_manager, inout, fg_ptrs.data(), inout, patch_x0 - x0,
-        blendings_[blending_idx], blendings_.data() + blending_idx + 1,
+        patch_x1 - patch_x0, blendings_[blending_idx],
+        blendings_.data() + blending_idx + 1,
         shared_->metadata->m.extra_channel_info));
   }
   return true;

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

@@ -8,6 +8,7 @@
 
 // Chooses reference patches, and avoids encoding them once per occurrence.
 
+#include <jxl/memory_manager.h>
 #include <sys/types.h>
 
 #include <cstddef>
@@ -99,8 +100,8 @@ class PatchDictionary {
 
   bool HasAny() const { return !positions_.empty(); }
 
-  Status Decode(BitReader* br, size_t xsize, size_t ysize,
+  Status Decode(JxlMemoryManager* memory_manager, BitReader* br, size_t xsize,
-                bool* uses_extra_channels);
+                size_t ysize, bool* uses_extra_channels);
 
   void Clear() {
     positions_.clear();

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

@@ -15,6 +15,7 @@
 #include "lib/jxl/base/compiler_specific.h"
 #include "lib/jxl/base/matrix_ops.h"
 #include "lib/jxl/base/rect.h"
+#include "lib/jxl/base/sanitizers.h"
 #include "lib/jxl/base/status.h"
 #include "lib/jxl/cms/jxl_cms_internal.h"
 #include "lib/jxl/cms/opsin_params.h"
@@ -23,7 +24,6 @@
 #include "lib/jxl/image.h"
 #include "lib/jxl/opsin_params.h"
 #include "lib/jxl/quantizer.h"
-#include "lib/jxl/sanitizers.h"
 HWY_BEFORE_NAMESPACE();
 namespace jxl {
 namespace HWY_NAMESPACE {

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

@@ -26,9 +26,7 @@
 #if JPEGXL_ENABLE_BOXES || JPEGXL_ENABLE_TRANSCODE_JPEG
 #include "lib/jxl/box_content_decoder.h"
 #endif
-#include "lib/jxl/dec_external_image.h"
 #include "lib/jxl/dec_frame.h"
-#include "lib/jxl/dec_modular.h"
 #if JPEGXL_ENABLE_TRANSCODE_JPEG
 #include "lib/jxl/decode_to_jpeg.h"
 #endif
@@ -38,10 +36,7 @@
 #include "lib/jxl/headers.h"
 #include "lib/jxl/icc_codec.h"
 #include "lib/jxl/image_bundle.h"
-#include "lib/jxl/loop_filter.h"
 #include "lib/jxl/memory_manager_internal.h"
-#include "lib/jxl/sanitizers.h"
-#include "lib/jxl/toc.h"
 
 namespace {
 
@@ -364,7 +359,7 @@ struct JxlDecoderStruct {
   bool got_transform_data;  // To skip everything before ICC.
   bool got_all_headers;     // Codestream metadata headers.
   bool post_headers;        // Already decoding pixels.
-  jxl::ICCReader icc_reader;
+  std::unique_ptr<jxl::ICCReader> icc_reader;
   jxl::JxlDecoderFrameIndexBox frame_index_box;
   // This means either we actually got the preview image, or determined we
   // cannot get it or there is none.
@@ -687,7 +682,7 @@ void JxlDecoderRewindDecodingState(JxlDecoder* dec) {
   dec->got_transform_data = false;
   dec->got_all_headers = false;
   dec->post_headers = false;
-  dec->icc_reader.Reset();
+  if (dec->icc_reader) dec->icc_reader->Reset();
   dec->got_preview_image = false;
   dec->preview_frame = false;
   dec->file_pos = 0;
@@ -1050,7 +1045,7 @@ JxlDecoderStatus JxlDecoderReadAllHeaders(JxlDecoder* dec) {
 
   if (dec->metadata.m.color_encoding.WantICC()) {
     jxl::Status status =
-        dec->icc_reader.Init(reader.get(), dec->memory_limit_base);
+        dec->icc_reader->Init(reader.get(), dec->memory_limit_base);
     // Always check AllReadsWithinBounds, not all the C++ decoder implementation
     // handles reader out of bounds correctly  yet (e.g. context map). Not
     // checking AllReadsWithinBounds can cause reader->Close() to trigger an
@@ -1063,8 +1058,8 @@ JxlDecoderStatus JxlDecoderReadAllHeaders(JxlDecoder* dec) {
       // Other non-successful status is an error
       return JXL_DEC_ERROR;
     }
-    PaddedBytes decoded_icc;
+    PaddedBytes decoded_icc{&dec->memory_manager};
-    status = dec->icc_reader.Process(reader.get(), &decoded_icc);
+    status = dec->icc_reader->Process(reader.get(), &decoded_icc);
     if (status.code() == StatusCode::kNotEnoughBytes) {
       return dec->RequestMoreInput();
     }
@@ -1087,7 +1082,7 @@ JxlDecoderStatus JxlDecoderReadAllHeaders(JxlDecoder* dec) {
   dec->codestream_bits_ahead = 0;
 
   if (!dec->passes_state) {
-    dec->passes_state.reset(new jxl::PassesDecoderState());
+    dec->passes_state.reset(new jxl::PassesDecoderState(&dec->memory_manager));
   }
 
   JXL_API_RETURN_IF_ERROR(
@@ -1188,6 +1183,10 @@ JxlDecoderStatus JxlDecoderProcessCodestream(JxlDecoder* dec) {
     return JXL_DEC_SUCCESS;
   }
 
+  if (!dec->icc_reader) {
+    dec->icc_reader.reset(new ICCReader(&dec->memory_manager));
+  }
+
   if (!dec->got_all_headers) {
     JxlDecoderStatus status = JxlDecoderReadAllHeaders(dec);
     if (status != JXL_DEC_SUCCESS) return status;
@@ -1233,7 +1232,8 @@ JxlDecoderStatus JxlDecoderProcessCodestream(JxlDecoder* dec) {
       }
 #endif
       if (!dec->ib) {
-        dec->ib.reset(new jxl::ImageBundle(&dec->image_metadata));
+        dec->ib.reset(
+            new jxl::ImageBundle(&dec->memory_manager, &dec->image_metadata));
       }
 #if JPEGXL_ENABLE_TRANSCODE_JPEG
       // If JPEG reconstruction is wanted and possible, set the jpeg_data of
@@ -2345,7 +2345,7 @@ JxlDecoderStatus JxlDecoderFlushImage(JxlDecoder* dec) {
 JXL_EXPORT JxlDecoderStatus JxlDecoderSetCms(JxlDecoder* dec,
                                              const JxlCmsInterface cms) {
   if (!dec->passes_state) {
-    dec->passes_state.reset(new jxl::PassesDecoderState());
+    dec->passes_state.reset(new jxl::PassesDecoderState(&dec->memory_manager));
   }
   dec->passes_state->output_encoding_info.color_management_system = cms;
   dec->passes_state->output_encoding_info.cms_set = true;

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

@@ -24,7 +24,6 @@
 #include <cstdlib>
 #include <cstring>
 #include <ostream>
-#include <set>
 #include <sstream>
 #include <string>
 #include <tuple>
@@ -203,6 +202,7 @@ enum PreviewMode {
 };
 
 void GeneratePreview(PreviewMode preview_mode, ImageBundle* ib) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   if (preview_mode == kSmallPreview) {
     ib->ShrinkTo(ib->xsize() / 7, ib->ysize() / 7);
   } else if (preview_mode == kBigPreview) {
@@ -213,15 +213,17 @@ void GeneratePreview(PreviewMode preview_mode, ImageBundle* ib) {
         }
       }
     };
-    JXL_ASSIGN_OR_DIE(Image3F preview,
+    JXL_ASSIGN_OR_DIE(
-                      Image3F::Create(ib->xsize() * 7, ib->ysize() * 7));
+        Image3F preview,
+        Image3F::Create(memory_manager, ib->xsize() * 7, ib->ysize() * 7));
     for (size_t c = 0; c < 3; ++c) {
       upsample7(ib->color()->Plane(c), &preview.Plane(c));
     }
     std::vector<ImageF> extra_channels;
     for (size_t i = 0; i < ib->extra_channels().size(); ++i) {
-      JXL_ASSIGN_OR_DIE(ImageF ec,
+      JXL_ASSIGN_OR_DIE(
-                        ImageF::Create(ib->xsize() * 7, ib->ysize() * 7));
+          ImageF ec,
+          ImageF::Create(memory_manager, ib->xsize() * 7, ib->ysize() * 7));
       upsample7(ib->extra_channels()[i], &ec);
       extra_channels.emplace_back(std::move(ec));
     }
@@ -257,12 +259,13 @@ struct TestCodestreamParams {
 std::vector<uint8_t> CreateTestJXLCodestream(
     Span<const uint8_t> pixels, size_t xsize, size_t ysize, size_t num_channels,
     const TestCodestreamParams& params) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   // Compress the pixels with JPEG XL.
   bool grayscale = (num_channels <= 2);
   bool have_alpha = ((num_channels & 1) == 0);
   bool include_alpha = have_alpha && params.jpeg_codestream == nullptr;
   size_t bitdepth = params.jpeg_codestream == nullptr ? 16 : 8;
-  CodecInOut io;
+  CodecInOut io{jxl::test::MemoryManager()};
   io.SetSize(xsize, ysize);
   ColorEncoding color_encoding;
   if (params.add_icc_profile) {
@@ -315,8 +318,8 @@ std::vector<uint8_t> CreateTestJXLCodestream(
       Bytes(jpeg_bytes).AppendTo(*params.jpeg_codestream);
       EXPECT_TRUE(jxl::jpeg::DecodeImageJPG(
           jxl::Bytes(jpeg_bytes.data(), jpeg_bytes.size()), &io));
-      EXPECT_TRUE(
+      EXPECT_TRUE(EncodeJPEGData(memory_manager, *io.Main().jpeg_data,
-          EncodeJPEGData(*io.Main().jpeg_data, &jpeg_data, params.cparams));
+                                 &jpeg_data, params.cparams));
       io.metadata.m.xyb_encoded = false;
     } else {
       JXL_ABORT(
@@ -720,7 +723,8 @@ std::vector<uint8_t> GetTestHeader(size_t xsize, size_t ysize,
                                    bool have_container, bool metadata_default,
                                    bool insert_extra_box,
                                    const jxl::IccBytes& icc_profile) {
-  jxl::BitWriter writer;
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
+  jxl::BitWriter writer{memory_manager};
   jxl::BitWriter::Allotment allotment(&writer, 65536);  // Large enough
 
   if (have_container) {
@@ -1352,7 +1356,7 @@ TEST_P(DecodeTestParam, PixelTest) {
     jxl::ColorEncoding color_encoding =
         jxl::ColorEncoding::SRGB(config.grayscale);
 
-    jxl::CodecInOut io;
+    jxl::CodecInOut io{jxl::test::MemoryManager()};
     if (config.include_alpha) io.metadata.m.SetAlphaBits(16);
     io.metadata.m.color_encoding = color_encoding;
     io.SetSize(config.xsize, config.ysize);
@@ -1651,6 +1655,7 @@ TEST(DecodeTest, PixelTestWithICCProfileLossless) {
 }
 
 TEST(DecodeTest, PixelTestWithICCProfileLossy) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   JxlDecoder* dec = JxlDecoderCreate(nullptr);
 
   size_t xsize = 123;
@@ -1680,7 +1685,7 @@ TEST(DecodeTest, PixelTestWithICCProfileLossy) {
   jxl::ColorEncoding color_encoding0;
   EXPECT_TRUE(color_encoding0.SetICC(GetIccTestProfile(), JxlGetDefaultCms()));
   jxl::Span<const uint8_t> span0(pixels.data(), pixels.size());
-  jxl::CodecInOut io0;
+  jxl::CodecInOut io0{memory_manager};
   io0.SetSize(xsize, ysize);
   EXPECT_TRUE(ConvertFromExternal(span0, xsize, ysize, color_encoding0,
                                   /*bits_per_sample=*/16, format_orig,
@@ -1691,7 +1696,7 @@ TEST(DecodeTest, PixelTestWithICCProfileLossy) {
   jxl::Bytes(icc_data).AppendTo(icc);
   EXPECT_TRUE(color_encoding1.SetICC(std::move(icc), JxlGetDefaultCms()));
   jxl::Span<const uint8_t> span1(pixels2.data(), pixels2.size());
-  jxl::CodecInOut io1;
+  jxl::CodecInOut io1{memory_manager};
   io1.SetSize(xsize, ysize);
   EXPECT_TRUE(ConvertFromExternal(span1, xsize, ysize, color_encoding1,
                                   /*bits_per_sample=*/32, format,
@@ -1735,7 +1740,8 @@ double ButteraugliDistance(size_t xsize, size_t ysize,
                            const std::vector<uint8_t>& pixels_out,
                            const jxl::ColorEncoding& color_out,
                            float intensity_out) {
-  jxl::CodecInOut in;
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
+  jxl::CodecInOut in{memory_manager};
   in.metadata.m.color_encoding = color_in;
   in.metadata.m.SetIntensityTarget(intensity_in);
   JxlPixelFormat format_in = {static_cast<uint32_t>(color_in.Channels()),
@@ -1744,7 +1750,7 @@ double ButteraugliDistance(size_t xsize, size_t ysize,
       jxl::Bytes(pixels_in.data(), pixels_in.size()), xsize, ysize, color_in,
       /*bits_per_sample=*/16, format_in,
       /*pool=*/nullptr, &in.Main()));
-  jxl::CodecInOut out;
+  jxl::CodecInOut out{memory_manager};
   out.metadata.m.color_encoding = color_out;
   out.metadata.m.SetIntensityTarget(intensity_out);
   JxlPixelFormat format_out = {static_cast<uint32_t>(color_out.Channels()),
@@ -2058,6 +2064,7 @@ TEST_P(DecodeAllEncodingsWithCMSTest, DecodeWithCMS) {
 // Tests the case of lossy sRGB image without alpha channel, decoded to RGB8
 // and to RGBA8
 TEST(DecodeTest, PixelTestOpaqueSrgbLossy) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   for (unsigned channels = 3; channels <= 4; channels++) {
     JxlDecoder* dec = JxlDecoderCreate(nullptr);
 
@@ -2083,7 +2090,7 @@ TEST(DecodeTest, PixelTestOpaqueSrgbLossy) {
 
     jxl::ColorEncoding color_encoding0 = jxl::ColorEncoding::SRGB(false);
     jxl::Span<const uint8_t> span0(pixels.data(), pixels.size());
-    jxl::CodecInOut io0;
+    jxl::CodecInOut io0{memory_manager};
     io0.SetSize(xsize, ysize);
     EXPECT_TRUE(ConvertFromExternal(span0, xsize, ysize, color_encoding0,
                                     /*bits_per_sample=*/16, format_orig,
@@ -2091,7 +2098,7 @@ TEST(DecodeTest, PixelTestOpaqueSrgbLossy) {
 
     jxl::ColorEncoding color_encoding1 = jxl::ColorEncoding::SRGB(false);
     jxl::Span<const uint8_t> span1(pixels2.data(), pixels2.size());
-    jxl::CodecInOut io1;
+    jxl::CodecInOut io1{memory_manager};
     EXPECT_TRUE(ConvertFromExternal(span1, xsize, ysize, color_encoding1,
                                     /*bits_per_sample=*/8, format,
                                     /*pool=*/nullptr, &io1.Main()));
@@ -2134,7 +2141,7 @@ TEST(DecodeTest, PixelTestOpaqueSrgbLossyNoise) {
 
     jxl::ColorEncoding color_encoding0 = jxl::ColorEncoding::SRGB(false);
     jxl::Span<const uint8_t> span0(pixels.data(), pixels.size());
-    jxl::CodecInOut io0;
+    jxl::CodecInOut io0{jxl::test::MemoryManager()};
     io0.SetSize(xsize, ysize);
     EXPECT_TRUE(ConvertFromExternal(span0, xsize, ysize, color_encoding0,
                                     /*bits_per_sample=*/16, format_orig,
@@ -2142,7 +2149,7 @@ TEST(DecodeTest, PixelTestOpaqueSrgbLossyNoise) {
 
     jxl::ColorEncoding color_encoding1 = jxl::ColorEncoding::SRGB(false);
     jxl::Span<const uint8_t> span1(pixels2.data(), pixels2.size());
-    jxl::CodecInOut io1;
+    jxl::CodecInOut io1{jxl::test::MemoryManager()};
     EXPECT_TRUE(ConvertFromExternal(span1, xsize, ysize, color_encoding1,
                                     /*bits_per_sample=*/8, format,
                                     /*pool=*/nullptr, &io1.Main()));
@@ -2510,6 +2517,7 @@ TEST(DecodeTest, DCNotGettableTest) {
 }
 
 TEST(DecodeTest, PreviewTest) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   size_t xsize = 77;
   size_t ysize = 120;
   std::vector<uint8_t> pixels = jxl::test::GetSomeTestImage(xsize, ysize, 3, 0);
@@ -2540,7 +2548,7 @@ TEST(DecodeTest, PreviewTest) {
               JxlDecoderPreviewOutBufferSize(dec, &format, &buffer_size));
 
     jxl::ColorEncoding c_srgb = jxl::ColorEncoding::SRGB(false);
-    jxl::CodecInOut io0;
+    jxl::CodecInOut io0{memory_manager};
     EXPECT_TRUE(jxl::ConvertFromExternal(
         jxl::Bytes(pixels.data(), pixels.size()), xsize, ysize, c_srgb,
         /*bits_per_sample=*/16, format_orig, /*pool=*/nullptr, &io0.Main()));
@@ -2561,7 +2569,7 @@ TEST(DecodeTest, PreviewTest) {
 
     EXPECT_EQ(JXL_DEC_PREVIEW_IMAGE, JxlDecoderProcessInput(dec));
 
-    jxl::CodecInOut io1;
+    jxl::CodecInOut io1{memory_manager};
     EXPECT_TRUE(
         jxl::ConvertFromExternal(jxl::Bytes(preview.data(), preview.size()),
                                  xsize_preview, ysize_preview, c_srgb,
@@ -2619,6 +2627,7 @@ TEST(DecodeTest, AlignTest) {
 }
 
 TEST(DecodeTest, AnimationTest) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   size_t xsize = 123;
   size_t ysize = 77;
   static const size_t num_frames = 2;
@@ -2627,7 +2636,7 @@ TEST(DecodeTest, AnimationTest) {
   frames[1] = jxl::test::GetSomeTestImage(xsize, ysize, 3, 1);
   JxlPixelFormat format = {3, JXL_TYPE_UINT16, JXL_BIG_ENDIAN, 0};
 
-  jxl::CodecInOut io;
+  jxl::CodecInOut io{memory_manager};
   io.SetSize(xsize, ysize);
   io.metadata.m.SetUintSamples(16);
   io.metadata.m.color_encoding = jxl::ColorEncoding::SRGB(false);
@@ -2642,7 +2651,7 @@ TEST(DecodeTest, AnimationTest) {
   }
 
   for (size_t i = 0; i < num_frames; ++i) {
-    jxl::ImageBundle bundle(&io.metadata.m);
+    jxl::ImageBundle bundle(memory_manager, &io.metadata.m);
 
     EXPECT_TRUE(ConvertFromExternal(
         jxl::Bytes(frames[i].data(), frames[i].size()), xsize, ysize,
@@ -2720,6 +2729,7 @@ TEST(DecodeTest, AnimationTest) {
 }
 
 TEST(DecodeTest, AnimationTestStreaming) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   size_t xsize = 123;
   size_t ysize = 77;
   static const size_t num_frames = 2;
@@ -2728,7 +2738,7 @@ TEST(DecodeTest, AnimationTestStreaming) {
   frames[1] = jxl::test::GetSomeTestImage(xsize, ysize, 3, 1);
   JxlPixelFormat format = {3, JXL_TYPE_UINT16, JXL_BIG_ENDIAN, 0};
 
-  jxl::CodecInOut io;
+  jxl::CodecInOut io{memory_manager};
   io.SetSize(xsize, ysize);
   io.metadata.m.SetUintSamples(16);
   io.metadata.m.color_encoding = jxl::ColorEncoding::SRGB(false);
@@ -2743,7 +2753,7 @@ TEST(DecodeTest, AnimationTestStreaming) {
   }
 
   for (size_t i = 0; i < num_frames; ++i) {
-    jxl::ImageBundle bundle(&io.metadata.m);
+    jxl::ImageBundle bundle(memory_manager, &io.metadata.m);
 
     EXPECT_TRUE(ConvertFromExternal(
         jxl::Bytes(frames[i].data(), frames[i].size()), xsize, ysize,
@@ -2931,6 +2941,7 @@ TEST(DecodeTest, ExtraChannelTest) {
 }
 
 TEST(DecodeTest, SkipCurrentFrameTest) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   size_t xsize = 90;
   size_t ysize = 120;
   constexpr size_t num_frames = 7;
@@ -2940,7 +2951,7 @@ TEST(DecodeTest, SkipCurrentFrameTest) {
   }
   JxlPixelFormat format = {3, JXL_TYPE_UINT16, JXL_BIG_ENDIAN, 0};
 
-  jxl::CodecInOut io;
+  jxl::CodecInOut io{memory_manager};
   io.SetSize(xsize, ysize);
   io.metadata.m.SetUintSamples(16);
   io.metadata.m.color_encoding = jxl::ColorEncoding::SRGB(false);
@@ -2955,7 +2966,7 @@ TEST(DecodeTest, SkipCurrentFrameTest) {
   }
 
   for (size_t i = 0; i < num_frames; ++i) {
-    jxl::ImageBundle bundle(&io.metadata.m);
+    jxl::ImageBundle bundle(memory_manager, &io.metadata.m);
     if (i & 1) {
       // Mark some frames as referenceable, others not.
       bundle.use_for_next_frame = true;
@@ -3043,6 +3054,7 @@ TEST(DecodeTest, SkipCurrentFrameTest) {
 }
 
 TEST(DecodeTest, SkipFrameTest) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   size_t xsize = 90;
   size_t ysize = 120;
   constexpr size_t num_frames = 16;
@@ -3052,7 +3064,7 @@ TEST(DecodeTest, SkipFrameTest) {
   }
   JxlPixelFormat format = {3, JXL_TYPE_UINT16, JXL_BIG_ENDIAN, 0};
 
-  jxl::CodecInOut io;
+  jxl::CodecInOut io{memory_manager};
   io.SetSize(xsize, ysize);
   io.metadata.m.SetUintSamples(16);
   io.metadata.m.color_encoding = jxl::ColorEncoding::SRGB(false);
@@ -3067,7 +3079,7 @@ TEST(DecodeTest, SkipFrameTest) {
   }
 
   for (size_t i = 0; i < num_frames; ++i) {
-    jxl::ImageBundle bundle(&io.metadata.m);
+    jxl::ImageBundle bundle(memory_manager, &io.metadata.m);
     if (i & 1) {
       // Mark some frames as referenceable, others not.
       bundle.use_for_next_frame = true;
@@ -3180,13 +3192,14 @@ TEST(DecodeTest, SkipFrameTest) {
 }
 
 TEST(DecodeTest, SkipFrameWithBlendingTest) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   size_t xsize = 90;
   size_t ysize = 120;
   constexpr size_t num_frames = 16;
   std::vector<uint8_t> frames[num_frames];
   JxlPixelFormat format = {3, JXL_TYPE_UINT16, JXL_BIG_ENDIAN, 0};
 
-  jxl::CodecInOut io;
+  jxl::CodecInOut io{memory_manager};
   io.SetSize(xsize, ysize);
   io.metadata.m.SetUintSamples(16);
   io.metadata.m.color_encoding = jxl::ColorEncoding::SRGB(false);
@@ -3204,7 +3217,7 @@ TEST(DecodeTest, SkipFrameWithBlendingTest) {
       // An internal frame with 0 duration, and use_for_next_frame, this is a
       // frame that is not rendered and not output by the API, but on which the
       // rendered frames depend
-      jxl::ImageBundle bundle_internal(&io.metadata.m);
+      jxl::ImageBundle bundle_internal(memory_manager, &io.metadata.m);
       EXPECT_TRUE(ConvertFromExternal(
           jxl::Bytes(frame_internal.data(), frame_internal.size()), xsize,
           ysize, jxl::ColorEncoding::SRGB(/*is_gray=*/false),
@@ -3219,7 +3232,7 @@ TEST(DecodeTest, SkipFrameWithBlendingTest) {
         jxl::test::GetSomeTestImage(xsize, ysize, 3, i * 2);
     // Actual rendered frame
     frame_durations[i] = 5 + i;
-    jxl::ImageBundle bundle(&io.metadata.m);
+    jxl::ImageBundle bundle(memory_manager, &io.metadata.m);
     EXPECT_TRUE(ConvertFromExternal(jxl::Bytes(frame.data(), frame.size()),
                                     xsize, ysize,
                                     jxl::ColorEncoding::SRGB(/*is_gray=*/false),
@@ -3394,13 +3407,14 @@ TEST(DecodeTest, SkipFrameWithBlendingTest) {
 }
 
 TEST(DecodeTest, SkipFrameWithAlphaBlendingTest) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   size_t xsize = 90;
   size_t ysize = 120;
   constexpr size_t num_frames = 16;
   std::vector<uint8_t> frames[num_frames + 5];
   JxlPixelFormat format = {4, JXL_TYPE_UINT16, JXL_BIG_ENDIAN, 0};
 
-  jxl::CodecInOut io;
+  jxl::CodecInOut io{memory_manager};
   io.SetSize(xsize, ysize);
   io.metadata.m.SetUintSamples(16);
   io.metadata.m.color_encoding = jxl::ColorEncoding::SRGB(false);
@@ -3427,7 +3441,7 @@ TEST(DecodeTest, SkipFrameWithAlphaBlendingTest) {
       // An internal frame with 0 duration, and use_for_next_frame, this is a
       // frame that is not rendered and not output by default by the API, but on
       // which the rendered frames depend
-      jxl::ImageBundle bundle_internal(&io.metadata.m);
+      jxl::ImageBundle bundle_internal(memory_manager, &io.metadata.m);
       EXPECT_TRUE(ConvertFromExternal(
           jxl::Bytes(frame_internal.data(), frame_internal.size()), xsize / 2,
           ysize / 2, jxl::ColorEncoding::SRGB(/*is_gray=*/false),
@@ -3447,7 +3461,7 @@ TEST(DecodeTest, SkipFrameWithAlphaBlendingTest) {
     std::vector<uint8_t> frame =
         jxl::test::GetSomeTestImage(cropxsize, cropysize, 4, i * 2);
     // Actual rendered frame
-    jxl::ImageBundle bundle(&io.metadata.m);
+    jxl::ImageBundle bundle(memory_manager, &io.metadata.m);
     EXPECT_TRUE(ConvertFromExternal(jxl::Bytes(frame.data(), frame.size()),
                                     cropxsize, cropysize,
                                     jxl::ColorEncoding::SRGB(/*is_gray=*/false),
@@ -3684,14 +3698,15 @@ TEST(DecodeTest, SkipFrameWithAlphaBlendingTest) {
 }
 
 TEST(DecodeTest, OrientedCroppedFrameTest) {
-  const auto test = [](bool keep_orientation, uint32_t orientation,
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
-                       uint32_t resampling) {
+  const auto test = [&](bool keep_orientation, uint32_t orientation,
+                        uint32_t resampling) {
     size_t xsize = 90;
     size_t ysize = 120;
     JxlPixelFormat format = {4, JXL_TYPE_UINT16, JXL_BIG_ENDIAN, 0};
     size_t oxsize = (!keep_orientation && orientation > 4 ? ysize : xsize);
     size_t oysize = (!keep_orientation && orientation > 4 ? xsize : ysize);
-    jxl::CodecInOut io;
+    jxl::CodecInOut io{memory_manager};
     io.SetSize(xsize, ysize);
     io.metadata.m.SetUintSamples(16);
     io.metadata.m.color_encoding = jxl::ColorEncoding::SRGB(false);
@@ -3707,7 +3722,7 @@ TEST(DecodeTest, OrientedCroppedFrameTest) {
 
       std::vector<uint8_t> frame =
           jxl::test::GetSomeTestImage(cropxsize, cropysize, 4, i * 2);
-      jxl::ImageBundle bundle(&io.metadata.m);
+      jxl::ImageBundle bundle(memory_manager, &io.metadata.m);
       EXPECT_TRUE(ConvertFromExternal(
           jxl::Bytes(frame.data(), frame.size()), cropxsize, cropysize,
           jxl::ColorEncoding::SRGB(/*is_gray=*/false),
@@ -3841,6 +3856,7 @@ struct StreamPositions {
 
 void AnalyzeCodestream(const std::vector<uint8_t>& data,
                        StreamPositions* streampos) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   // Unbox data to codestream and mark where it is broken up by boxes.
   std::vector<uint8_t> codestream;
   std::vector<std::pair<size_t, size_t>> breakpoints;
@@ -3926,8 +3942,9 @@ void AnalyzeCodestream(const std::vector<uint8_t>& data,
                            frame_header.passes.num_passes);
     std::vector<uint64_t> section_offsets;
     std::vector<uint32_t> section_sizes;
-    ASSERT_TRUE(ReadGroupOffsets(toc_entries, &br, &section_offsets,
+    ASSERT_TRUE(ReadGroupOffsets(memory_manager, toc_entries, &br,
-                                 &section_sizes, &groups_total_size));
+                                 &section_offsets, &section_sizes,
+                                 &groups_total_size));
     EXPECT_EQ(br.TotalBitsConsumed() % jxl::kBitsPerByte, 0);
     size_t sections_start = br.TotalBitsConsumed() / jxl::kBitsPerByte;
     p.toc_end = add_offset(sections_start);
@@ -4778,6 +4795,7 @@ JXL_GTEST_INSTANTIATE_TEST_SUITE_P(DecodeProgressiveTestInstantiation,
                                    DecodeProgressiveTest,
                                    ::testing::Range(0, 8));
 TEST_P(DecodeProgressiveTest, ProgressiveEventTest) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   const int params = GetParam();
   bool single_group = ((params & 1) != 0);
   bool lossless = (((params >> 1) & 1) != 0);
@@ -4804,7 +4822,7 @@ TEST_P(DecodeProgressiveTest, ProgressiveEventTest) {
         jxl::test::GetSomeTestImage(xsize, ysize, num_channels, 0);
     JxlPixelFormat format = {num_channels, JXL_TYPE_UINT16, JXL_BIG_ENDIAN, 0};
     jxl::ColorEncoding color_encoding = jxl::ColorEncoding::SRGB(false);
-    jxl::CodecInOut io;
+    jxl::CodecInOut io{memory_manager};
     EXPECT_TRUE(jxl::ConvertFromExternal(
         jxl::Bytes(pixels.data(), pixels.size()), xsize, ysize, color_encoding,
         /*bits_per_sample=*/16, format,
@@ -4921,7 +4939,7 @@ TEST_P(DecodeProgressiveTest, ProgressiveEventTest) {
       jxl::ButteraugliParams ba;
       std::vector<float> distances(kNumPasses + 1);
       for (int p = 0;; p = next_pass(p)) {
-        jxl::CodecInOut io1;
+        jxl::CodecInOut io1{memory_manager};
         EXPECT_TRUE(jxl::ConvertFromExternal(
             jxl::Bytes(passes[p].data(), passes[p].size()), xsize, ysize,
             color_encoding,
@@ -4995,24 +5013,26 @@ JXL_TRANSCODE_JPEG_TEST(DecodeTest, JPEGReconstructTestCodestream) {
 }
 
 JXL_TRANSCODE_JPEG_TEST(DecodeTest, JPEGReconstructionTest) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   const std::string jpeg_path = "jxl/flower/flower.png.im_q85_420.jpg";
   const std::vector<uint8_t> orig = jxl::test::ReadTestData(jpeg_path);
-  jxl::CodecInOut orig_io;
+  jxl::CodecInOut orig_io{memory_manager};
   ASSERT_TRUE(jxl::jpeg::DecodeImageJPG(jxl::Bytes(orig), &orig_io));
   orig_io.metadata.m.xyb_encoded = false;
-  jxl::BitWriter writer;
+  jxl::BitWriter writer{memory_manager};
   ASSERT_TRUE(WriteCodestreamHeaders(&orig_io.metadata, &writer, nullptr));
   writer.ZeroPadToByte();
   jxl::CompressParams cparams;
   cparams.color_transform = jxl::ColorTransform::kNone;
-  ASSERT_TRUE(jxl::EncodeFrame(cparams, jxl::FrameInfo{}, &orig_io.metadata,
+  ASSERT_TRUE(jxl::EncodeFrame(memory_manager, cparams, jxl::FrameInfo{},
-                               orig_io.Main(), *JxlGetDefaultCms(),
+                               &orig_io.metadata, orig_io.Main(),
+                               *JxlGetDefaultCms(),
                                /*pool=*/nullptr, &writer,
                                /*aux_out=*/nullptr));
 
   std::vector<uint8_t> jpeg_data;
-  ASSERT_TRUE(
+  ASSERT_TRUE(EncodeJPEGData(memory_manager, *orig_io.Main().jpeg_data.get(),
-      EncodeJPEGData(*orig_io.Main().jpeg_data.get(), &jpeg_data, cparams));
+                             &jpeg_data, cparams));
   std::vector<uint8_t> container;
   jxl::Bytes(jxl::kContainerHeader).AppendTo(container);
   jxl::AppendBoxHeader(jxl::MakeBoxType("jbrd"), jpeg_data.size(), false,
@@ -5549,12 +5569,14 @@ JXL_BOXES_TEST(DecodeTest, PartialCodestreamBoxTest) {
 }
 
 TEST(DecodeTest, SpotColorTest) {
-  jxl::CodecInOut io;
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
+  jxl::CodecInOut io{memory_manager};
   size_t xsize = 55;
   size_t ysize = 257;
   io.metadata.m.color_encoding = jxl::ColorEncoding::LinearSRGB();
-  JXL_ASSIGN_OR_DIE(Image3F main, Image3F::Create(xsize, ysize));
+  JXL_ASSIGN_OR_DIE(Image3F main,
-  JXL_ASSIGN_OR_DIE(ImageF spot, ImageF::Create(xsize, ysize));
+                    Image3F::Create(memory_manager, xsize, ysize));
+  JXL_ASSIGN_OR_DIE(ImageF spot, ImageF::Create(memory_manager, xsize, ysize));
   jxl::ZeroFillImage(&main);
   jxl::ZeroFillImage(&spot);
 

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

@@ -5,6 +5,8 @@
 
 #include "lib/jxl/enc_ac_strategy.h"
 
+#include <jxl/memory_manager.h>
+
 #include <algorithm>
 #include <cmath>
 #include <cstdint>
@@ -213,7 +215,9 @@ const uint8_t* TypeMask(const uint8_t& raw_strategy) {
 Status DumpAcStrategy(const AcStrategyImage& ac_strategy, size_t xsize,
                       size_t ysize, const char* tag, AuxOut* aux_out,
                       const CompressParams& cparams) {
-  JXL_ASSIGN_OR_RETURN(Image3F color_acs, Image3F::Create(xsize, ysize));
+  JxlMemoryManager* memory_manager = ac_strategy.memory_manager();
+  JXL_ASSIGN_OR_RETURN(Image3F color_acs,
+                       Image3F::Create(memory_manager, xsize, ysize));
   for (size_t y = 0; y < ysize; y++) {
     float* JXL_RESTRICT rows[3] = {
         color_acs.PlaneRow(0, y),
@@ -451,11 +455,11 @@ float EstimateEntropy(const AcStrategy& acs, float entropy_mul, size_t x,
         }
       }
       static const double kChannelMul[3] = {
-          10.2,
+          pow(10.2, 8.0),
-          1.0,
+          pow(1.0, 8.0),
-          1.03,
+          pow(1.03, 8.0),
       };
-      lossc = Mul(Set(df8, pow(kChannelMul[c], 8.0)), lossc);
+      lossc = Mul(Set(df8, kChannelMul[c]), lossc);
       loss = Add(loss, lossc);
     }
     entropy += config.cost_delta * GetLane(SumOfLanes(df, entropy_v));
@@ -846,7 +850,7 @@ void ProcessRectACS(const CompressParams& cparams, const ACSConfig& config,
     float entropy_mul;
   };
   // These numbers need to be figured out manually and looking at
-  // ringing next to sky etc. Optimization will find larger numbers
+  // ringing next to sky etc. Optimization will find smaller numbers
   // and produce more ringing than is ideal. Larger numbers will
   // help stop ringing.
   const float entropy_mul16X8 = 1.25;

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

@@ -5,12 +5,13 @@
 
 #include "lib/jxl/enc_adaptive_quantization.h"
 
-#include <stddef.h>
+#include <jxl/memory_manager.h>
-#include <stdlib.h>
 
 #include <algorithm>
 #include <atomic>
 #include <cmath>
+#include <cstddef>
+#include <cstdlib>
 #include <string>
 #include <vector>
 
@@ -171,30 +172,28 @@ V GammaModulation(const D d, const size_t x, const size_t y,
   JXL_DASSERT(kBias > jxl::cms::kOpsinAbsorbanceBias[2]);
   auto overall_ratio = Zero(d);
   auto bias = Set(d, kBias);
-  auto half = Set(d, 0.5f);
   for (size_t dy = 0; dy < 8; ++dy) {
     const float* const JXL_RESTRICT row_in_x = rect.ConstRow(xyb_x, y + dy);
     const float* const JXL_RESTRICT row_in_y = rect.ConstRow(xyb_y, y + dy);
     for (size_t dx = 0; dx < 8; dx += Lanes(d)) {
       const auto iny = Add(Load(d, row_in_y + x + dx), bias);
       const auto inx = Load(d, row_in_x + x + dx);
+
       const auto r = Sub(iny, inx);
-      const auto g = Add(iny, inx);
       const auto ratio_r =
           RatioOfDerivativesOfCubicRootToSimpleGamma</*invert=*/true>(d, r);
+      overall_ratio = Add(overall_ratio, ratio_r);
+
+      const auto g = Add(iny, inx);
       const auto ratio_g =
           RatioOfDerivativesOfCubicRootToSimpleGamma</*invert=*/true>(d, g);
-      const auto avg_ratio = Mul(half, Add(ratio_r, ratio_g));
+      overall_ratio = Add(overall_ratio, ratio_g);
-
-      overall_ratio = Add(overall_ratio, avg_ratio);
     }
   }
-  overall_ratio = Mul(SumOfLanes(d, overall_ratio), Set(d, 1.0f / 64));
+  overall_ratio = Mul(SumOfLanes(d, overall_ratio), Set(d, 0.5f / 64));
   // ideally -1.0, but likely optimal correction adds some entropy, so slightly
   // less than that.
-  // ln(2) constant folded in because we want std::log but have FastLog2f.
+  const auto kGam = Set(d, 0.1005613337192697f);
-  static const float v = 0.14507933746197058f;
-  const auto kGam = Set(d, v * 0.693147180559945f);
   return MulAdd(kGam, FastLog2f(d, overall_ratio), out_val);
 }
 
@@ -396,13 +395,15 @@ void FuzzyErosion(const float butteraugli_target, const Rect& from_rect,
 }
 
 struct AdaptiveQuantizationImpl {
-  Status PrepareBuffers(size_t num_threads) {
+  Status PrepareBuffers(JxlMemoryManager* memory_manager, size_t num_threads) {
-    JXL_ASSIGN_OR_RETURN(diff_buffer,
+    JXL_ASSIGN_OR_RETURN(
-                         ImageF::Create(kEncTileDim + 8, num_threads));
+        diff_buffer,
+        ImageF::Create(memory_manager, kEncTileDim + 8, num_threads));
     for (size_t i = pre_erosion.size(); i < num_threads; i++) {
-      JXL_ASSIGN_OR_RETURN(ImageF tmp,
+      JXL_ASSIGN_OR_RETURN(
-                           ImageF::Create(kEncTileDimInBlocks * 2 + 2,
+          ImageF tmp,
-                                          kEncTileDimInBlocks * 2 + 2));
+          ImageF::Create(memory_manager, kEncTileDimInBlocks * 2 + 2,
+                         kEncTileDimInBlocks * 2 + 2));
       pre_erosion.emplace_back(std::move(tmp));
     }
     return true;
@@ -572,7 +573,8 @@ struct AdaptiveQuantizationImpl {
   ImageF diff_buffer;
 };
 
-Status Blur1x1Masking(ThreadPool* pool, ImageF* mask1x1, const Rect& rect) {
+Status Blur1x1Masking(JxlMemoryManager* memory_manager, ThreadPool* pool,
+                      ImageF* mask1x1, const Rect& rect) {
   // Blur the mask1x1 to obtain the masking image.
   // Before blurring it contains an image of absolute value of the
   // Laplacian of the intensity channel.
@@ -598,7 +600,8 @@ Status Blur1x1Masking(ThreadPool* pool, ImageF* mask1x1, const Rect& rect) {
                         {HWY_REP4(normalize_mul * kFilterMask1x1[1])},
                         {HWY_REP4(normalize_mul * kFilterMask1x1[4])},
                         {HWY_REP4(normalize_mul * kFilterMask1x1[3])}};
-  JXL_ASSIGN_OR_RETURN(ImageF temp, ImageF::Create(rect.xsize(), rect.ysize()));
+  JXL_ASSIGN_OR_RETURN(
+      ImageF temp, ImageF::Create(memory_manager, rect.xsize(), rect.ysize()));
   Symmetric5(*mask1x1, rect, weights, pool, &temp);
   *mask1x1 = std::move(temp);
   return true;
@@ -613,15 +616,19 @@ StatusOr<ImageF> AdaptiveQuantizationMap(const float butteraugli_target,
   AdaptiveQuantizationImpl impl;
   const size_t xsize_blocks = rect.xsize() / kBlockDim;
   const size_t ysize_blocks = rect.ysize() / kBlockDim;
-  JXL_ASSIGN_OR_RETURN(impl.aq_map, ImageF::Create(xsize_blocks, ysize_blocks));
+  JxlMemoryManager* memory_manager = xyb.memory_manager();
-  JXL_ASSIGN_OR_RETURN(*mask, ImageF::Create(xsize_blocks, ysize_blocks));
+  JXL_ASSIGN_OR_RETURN(
-  JXL_ASSIGN_OR_RETURN(*mask1x1, ImageF::Create(xyb.xsize(), xyb.ysize()));
+      impl.aq_map, ImageF::Create(memory_manager, xsize_blocks, ysize_blocks));
+  JXL_ASSIGN_OR_RETURN(
+      *mask, ImageF::Create(memory_manager, xsize_blocks, ysize_blocks));
+  JXL_ASSIGN_OR_RETURN(
+      *mask1x1, ImageF::Create(memory_manager, xyb.xsize(), xyb.ysize()));
   JXL_CHECK(RunOnPool(
       pool, 0,
       DivCeil(xsize_blocks, kEncTileDimInBlocks) *
           DivCeil(ysize_blocks, kEncTileDimInBlocks),
       [&](const size_t num_threads) {
-        return !!impl.PrepareBuffers(num_threads);
+        return !!impl.PrepareBuffers(memory_manager, num_threads);
       },
       [&](const uint32_t tid, const size_t thread) {
         size_t n_enc_tiles = DivCeil(xsize_blocks, kEncTileDimInBlocks);
@@ -637,7 +644,7 @@ StatusOr<ImageF> AdaptiveQuantizationMap(const float butteraugli_target,
       },
       "AQ DiffPrecompute"));
 
-  JXL_RETURN_IF_ERROR(Blur1x1Masking(pool, mask1x1, rect));
+  JXL_RETURN_IF_ERROR(Blur1x1Masking(memory_manager, pool, mask1x1, rect));
   return std::move(impl).aq_map;
 }
 
@@ -675,10 +682,12 @@ Status DumpHeatmap(const CompressParams& cparams, const AuxOut* aux_out,
 Status DumpHeatmaps(const CompressParams& cparams, const AuxOut* aux_out,
                     float ba_target, const ImageF& quant_field,
                     const ImageF& tile_heatmap, const ImageF& bt_diffmap) {
+  JxlMemoryManager* memory_manager = quant_field.memory_manager();
   if (JXL_DEBUG_ADAPTIVE_QUANTIZATION) {
     if (!WantDebugOutput(cparams)) return true;
-    JXL_ASSIGN_OR_RETURN(ImageF inv_qmap, ImageF::Create(quant_field.xsize(),
+    JXL_ASSIGN_OR_RETURN(ImageF inv_qmap,
-                                                         quant_field.ysize()));
+                         ImageF::Create(memory_manager, quant_field.xsize(),
+                                        quant_field.ysize()));
     for (size_t y = 0; y < quant_field.ysize(); ++y) {
       const float* JXL_RESTRICT row_q = quant_field.ConstRow(y);
       float* JXL_RESTRICT row_inv_q = inv_qmap.Row(y);
@@ -702,8 +711,9 @@ StatusOr<ImageF> TileDistMap(const ImageF& distmap, int tile_size, int margin,
                              const AcStrategyImage& ac_strategy) {
   const int tile_xsize = (distmap.xsize() + tile_size - 1) / tile_size;
   const int tile_ysize = (distmap.ysize() + tile_size - 1) / tile_size;
+  JxlMemoryManager* memory_manager = distmap.memory_manager();
   JXL_ASSIGN_OR_RETURN(ImageF tile_distmap,
-                       ImageF::Create(tile_xsize, tile_ysize));
+                       ImageF::Create(memory_manager, tile_xsize, tile_ysize));
   size_t distmap_stride = tile_distmap.PixelsPerRow();
   for (int tile_y = 0; tile_y < tile_ysize; ++tile_y) {
     AcStrategyRow ac_strategy_row = ac_strategy.ConstRow(tile_y);
@@ -774,8 +784,9 @@ StatusOr<ImageBundle> RoundtripImage(const FrameHeader& frame_header,
                                      PassesEncoderState* enc_state,
                                      const JxlCmsInterface& cms,
                                      ThreadPool* pool) {
+  JxlMemoryManager* memory_manager = enc_state->memory_manager();
   std::unique_ptr<PassesDecoderState> dec_state =
-      jxl::make_unique<PassesDecoderState>();
+      jxl::make_unique<PassesDecoderState>(memory_manager);
   JXL_CHECK(dec_state->output_encoding_info.SetFromMetadata(
       *enc_state->shared.metadata));
   dec_state->shared = &enc_state->shared;
@@ -787,18 +798,19 @@ StatusOr<ImageBundle> RoundtripImage(const FrameHeader& frame_header,
 
   size_t num_special_frames = enc_state->special_frames.size();
   size_t num_passes = enc_state->progressive_splitter.GetNumPasses();
-  ModularFrameEncoder modular_frame_encoder(frame_header, enc_state->cparams,
+  ModularFrameEncoder modular_frame_encoder(memory_manager, frame_header,
-                                            false);
+                                            enc_state->cparams, false);
   JXL_CHECK(InitializePassesEncoder(frame_header, opsin, Rect(opsin), cms, pool,
                                     enc_state, &modular_frame_encoder,
                                     nullptr));
   JXL_CHECK(dec_state->Init(frame_header));
   JXL_CHECK(dec_state->InitForAC(num_passes, pool));
 
-  ImageBundle decoded(&enc_state->shared.metadata->m);
+  ImageBundle decoded(memory_manager, &enc_state->shared.metadata->m);
   decoded.origin = frame_header.frame_origin;
-  JXL_ASSIGN_OR_RETURN(Image3F tmp,
+  JXL_ASSIGN_OR_RETURN(
-                       Image3F::Create(opsin.xsize(), opsin.ysize()));
+      Image3F tmp,
+      Image3F::Create(memory_manager, opsin.xsize(), opsin.ysize()));
   decoded.SetFromImage(std::move(tmp),
                        dec_state->output_encoding_info.color_encoding);
 
@@ -811,7 +823,8 @@ StatusOr<ImageBundle> RoundtripImage(const FrameHeader& frame_header,
   // Same as frame_header.nonserialized_metadata->m
   const ImageMetadata& metadata = *decoded.metadata();
 
-  JXL_CHECK(dec_state->PreparePipeline(frame_header, &decoded, options));
+  JXL_CHECK(dec_state->PreparePipeline(
+      frame_header, &enc_state->shared.metadata->m, &decoded, options));
 
   hwy::AlignedUniquePtr<GroupDecCache[]> group_dec_caches;
   const auto allocate_storage = [&](const size_t num_threads) -> Status {
@@ -834,7 +847,7 @@ StatusOr<ImageBundle> RoundtripImage(const FrameHeader& frame_header,
         dec_state->render_pipeline->GetInputBuffers(group_index, thread);
     JXL_CHECK(DecodeGroupForRoundtrip(
         frame_header, enc_state->coeffs, group_index, dec_state.get(),
-        &group_dec_caches[thread], thread, input, &decoded, nullptr));
+        &group_dec_caches[thread], thread, input, nullptr, nullptr));
     for (size_t c = 0; c < metadata.num_extra_channels; c++) {
       std::pair<ImageF*, Rect> ri = input.GetBuffer(3 + c);
       FillPlane(0.0f, ri.first, ri.second);
@@ -869,6 +882,7 @@ Status FindBestQuantization(const FrameHeader& frame_header,
     // so in this case we enable it only for very high butteraugli targets.
     return true;
   }
+  JxlMemoryManager* memory_manager = enc_state->memory_manager();
   Quantizer& quantizer = enc_state->shared.quantizer;
   ImageI& raw_quant_field = enc_state->shared.raw_quant_field;
 
@@ -886,7 +900,7 @@ Status FindBestQuantization(const FrameHeader& frame_header,
   ImageF tile_distmap;
   JXL_ASSIGN_OR_RETURN(
       ImageF initial_quant_field,
-      ImageF::Create(quant_field.xsize(), quant_field.ysize()));
+      ImageF::Create(memory_manager, quant_field.xsize(), quant_field.ysize()));
   CopyImageTo(quant_field, &initial_quant_field);
 
   float initial_qf_min;

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

@@ -5,16 +5,14 @@
 
 #include "lib/jxl/enc_ans.h"
 
+#include <jxl/memory_manager.h>
 #include <jxl/types.h>
-#include <stdint.h>
 
 #include <algorithm>
 #include <array>
 #include <cmath>
 #include <cstdint>
 #include <limits>
-#include <numeric>
-#include <type_traits>
 #include <unordered_map>
 #include <utility>
 #include <vector>
@@ -439,6 +437,7 @@ uint32_t ComputeBestMethod(
 // Returns an estimate of the cost of encoding this histogram and the
 // corresponding data.
 size_t BuildAndStoreANSEncodingData(
+    JxlMemoryManager* memory_manager,
     HistogramParams::ANSHistogramStrategy ans_histogram_strategy,
     const ANSHistBin* histogram, size_t alphabet_size, size_t log_alpha_size,
     bool use_prefix_code, ANSEncSymbolInfo* info, BitWriter* writer) {
@@ -454,7 +453,7 @@ size_t BuildAndStoreANSEncodingData(
       std::vector<uint8_t> depths(alphabet_size);
       std::vector<uint16_t> bits(alphabet_size);
       if (writer == nullptr) {
-        BitWriter tmp_writer;
+        BitWriter tmp_writer{memory_manager};
         BitWriter::Allotment allotment(
             &tmp_writer, 8 * alphabet_size + 8);  // safe upper bound
         BuildAndStoreHuffmanTree(histo.data(), alphabet_size, depths.data(),
@@ -715,7 +714,7 @@ class HistogramBuilder {
 
   // NOTE: `layer` is only for clustered_entropy; caller does ReclaimAndCharge.
   size_t BuildAndStoreEntropyCodes(
-      const HistogramParams& params,
+      JxlMemoryManager* memory_manager, const HistogramParams& params,
       const std::vector<std::vector<Token>>& tokens, EntropyEncodingData* codes,
       std::vector<uint8_t>* context_map, BitWriter* writer, size_t layer,
       AuxOut* aux_out) const {
@@ -810,14 +809,15 @@ class HistogramBuilder {
       codes->encoding_info.back().resize(alphabet_size);
       BitWriter* histo_writer = writer;
       if (params.streaming_mode) {
-        codes->encoded_histograms.emplace_back();
+        codes->encoded_histograms.emplace_back(memory_manager);
         histo_writer = &codes->encoded_histograms.back();
       }
       BitWriter::Allotment allotment(histo_writer, 256 + alphabet_size * 24);
       cost += BuildAndStoreANSEncodingData(
-          params.ans_histogram_strategy, clustered_histograms[c].data_.data(),
+          memory_manager, params.ans_histogram_strategy,
-          alphabet_size, log_alpha_size, codes->use_prefix_code,
+          clustered_histograms[c].data_.data(), alphabet_size, log_alpha_size,
-          codes->encoding_info.back().data(), histo_writer);
+          codes->use_prefix_code, codes->encoding_info.back().data(),
+          histo_writer);
       allotment.FinishedHistogram(histo_writer);
       allotment.ReclaimAndCharge(histo_writer, layer, aux_out);
       if (params.streaming_mode) {
@@ -1535,13 +1535,11 @@ void EncodeHistograms(const std::vector<uint8_t>& context_map,
   allotment.ReclaimAndCharge(writer, layer, aux_out);
 }
 
-size_t BuildAndEncodeHistograms(const HistogramParams& params,
+size_t BuildAndEncodeHistograms(
-                                size_t num_contexts,
+    JxlMemoryManager* memory_manager, const HistogramParams& params,
-                                std::vector<std::vector<Token>>& tokens,
+    size_t num_contexts, std::vector<std::vector<Token>>& tokens,
-                                EntropyEncodingData* codes,
+    EntropyEncodingData* codes, std::vector<uint8_t>* context_map,
-                                std::vector<uint8_t>* context_map,
+    BitWriter* writer, size_t layer, AuxOut* aux_out) {
-                                BitWriter* writer, size_t layer,
-                                AuxOut* aux_out) {
   size_t total_bits = 0;
   codes->lz77.nonserialized_distance_context = num_contexts;
   std::vector<std::vector<Token>> tokens_lz77;
@@ -1655,19 +1653,20 @@ size_t BuildAndEncodeHistograms(const HistogramParams& params,
         CreateFlatHistogram(alphabet_size, ANS_TAB_SIZE);
     codes->encoding_info.emplace_back();
     codes->encoding_info.back().resize(alphabet_size);
-    codes->encoded_histograms.emplace_back();
+    codes->encoded_histograms.emplace_back(memory_manager);
     BitWriter* histo_writer = &codes->encoded_histograms.back();
     BitWriter::Allotment allotment(histo_writer, 256 + alphabet_size * 24);
     BuildAndStoreANSEncodingData(
-        params.ans_histogram_strategy, counts.data(), alphabet_size,
+        memory_manager, params.ans_histogram_strategy, counts.data(),
-        log_alpha_size, codes->use_prefix_code,
+        alphabet_size, log_alpha_size, codes->use_prefix_code,
         codes->encoding_info.back().data(), histo_writer);
     allotment.ReclaimAndCharge(histo_writer, 0, nullptr);
   }
 
   // Encode histograms.
-  total_bits += builder.BuildAndStoreEntropyCodes(
+  total_bits +=
-      params, tokens, codes, context_map, writer, layer, aux_out);
+      builder.BuildAndStoreEntropyCodes(memory_manager, params, tokens, codes,
+                                        context_map, writer, layer, aux_out);
   allotment.FinishedHistogram(writer);
   allotment.ReclaimAndCharge(writer, layer, aux_out);
 

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

@@ -9,6 +9,8 @@
 // Library to encode the ANS population counts to the bit-stream and encode
 // symbols based on the respective distributions.
 
+#include <jxl/memory_manager.h>
+
 #include <cstddef>
 #include <cstdint>
 #include <vector>
@@ -106,13 +108,11 @@ void EncodeHistograms(const std::vector<uint8_t>& context_map,
 // estimate of the total bits used for encoding the stream. If `writer` ==
 // nullptr, the bit estimate will not take into account the context map (which
 // does not get written if `num_contexts` == 1).
-size_t BuildAndEncodeHistograms(const HistogramParams& params,
+size_t BuildAndEncodeHistograms(
-                                size_t num_contexts,
+    JxlMemoryManager* memory_manager, const HistogramParams& params,
-                                std::vector<std::vector<Token>>& tokens,
+    size_t num_contexts, std::vector<std::vector<Token>>& tokens,
-                                EntropyEncodingData* codes,
+    EntropyEncodingData* codes, std::vector<uint8_t>* context_map,
-                                std::vector<uint8_t>* context_map,
+    BitWriter* writer, size_t layer, AuxOut* aux_out);
-                                BitWriter* writer, size_t layer,
-                                AuxOut* aux_out);
 
 // Write the tokens to a string.
 void WriteTokens(const std::vector<Token>& tokens,

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

@@ -1,318 +0,0 @@
-// 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_ar_control_field.h"
-
-#include <algorithm>
-#include <cstdint>
-#include <cstdlib>
-
-#undef HWY_TARGET_INCLUDE
-#define HWY_TARGET_INCLUDE "lib/jxl/enc_ar_control_field.cc"
-#include <hwy/foreach_target.h>
-#include <hwy/highway.h>
-
-#include "lib/jxl/ac_strategy.h"
-#include "lib/jxl/base/compiler_specific.h"
-#include "lib/jxl/base/rect.h"
-#include "lib/jxl/base/status.h"
-#include "lib/jxl/enc_params.h"
-#include "lib/jxl/image.h"
-#include "lib/jxl/image_ops.h"
-
-HWY_BEFORE_NAMESPACE();
-namespace jxl {
-namespace HWY_NAMESPACE {
-namespace {
-
-// These templates are not found via ADL.
-using hwy::HWY_NAMESPACE::Add;
-using hwy::HWY_NAMESPACE::GetLane;
-using hwy::HWY_NAMESPACE::Mul;
-using hwy::HWY_NAMESPACE::MulAdd;
-using hwy::HWY_NAMESPACE::Sqrt;
-
-Status ProcessTile(const CompressParams& cparams,
-                   const FrameHeader& frame_header, const Image3F& opsin,
-                   const Rect& opsin_rect, const ImageF& quant_field,
-                   const AcStrategyImage& ac_strategy, ImageB* epf_sharpness,
-                   const Rect& rect,
-                   ArControlFieldHeuristics::TempImages* temp_image) {
-  JXL_ASSERT(opsin_rect.x0() % 8 == 0);
-  JXL_ASSERT(opsin_rect.y0() % 8 == 0);
-  JXL_ASSERT(opsin_rect.xsize() % 8 == 0);
-  JXL_ASSERT(opsin_rect.ysize() % 8 == 0);
-  constexpr size_t N = kBlockDim;
-  if (cparams.butteraugli_distance < kMinButteraugliForDynamicAR ||
-      cparams.speed_tier > SpeedTier::kWombat ||
-      frame_header.loop_filter.epf_iters == 0) {
-    FillPlane(static_cast<uint8_t>(4), epf_sharpness, rect);
-    return true;
-  }
-
-  // Likely better to have a higher X weight, like:
-  // const float kChannelWeights[3] = {47.0f, 4.35f, 0.287f};
-  const float kChannelWeights[3] = {4.35f, 4.35f, 0.287f};
-  const float kChannelWeightsLapNeg[3] = {-0.125f * kChannelWeights[0],
-                                          -0.125f * kChannelWeights[1],
-                                          -0.125f * kChannelWeights[2]};
-  const size_t sharpness_stride =
-      static_cast<size_t>(epf_sharpness->PixelsPerRow());
-
-  size_t by0 = opsin_rect.y0() / 8 + rect.y0();
-  size_t by1 = by0 + rect.ysize();
-  size_t bx0 = opsin_rect.x0() / 8 + rect.x0();
-  size_t bx1 = bx0 + rect.xsize();
-  JXL_RETURN_IF_ERROR(temp_image->InitOnce());
-  ImageF& laplacian_sqrsum = temp_image->laplacian_sqrsum;
-  // Calculate the L2 of the 3x3 Laplacian in an integral transform
-  // (for example 32x32 dct). This relates to transforms ability
-  // to propagate artefacts.
-  size_t y0 = by0 == 0 ? 0 : by0 * N - 2;
-  size_t y1 = by1 * N == opsin.ysize() ? by1 * N : by1 * N + 2;
-  size_t x0 = bx0 == 0 ? 0 : bx0 * N - 2;
-  size_t x1 = bx1 * N == opsin.xsize() ? bx1 * N : bx1 * N + 2;
-  HWY_FULL(float) df;
-  for (size_t y = y0; y < y1; y++) {
-    float* JXL_RESTRICT laplacian_sqrsum_row =
-        laplacian_sqrsum.Row(y + 2 - by0 * N);
-    const float* JXL_RESTRICT in_row_t[3];
-    const float* JXL_RESTRICT in_row[3];
-    const float* JXL_RESTRICT in_row_b[3];
-    for (size_t c = 0; c < 3; c++) {
-      in_row_t[c] = opsin.ConstPlaneRow(c, y > 0 ? y - 1 : y);
-      in_row[c] = opsin.ConstPlaneRow(c, y);
-      in_row_b[c] = opsin.ConstPlaneRow(c, y + 1 < opsin.ysize() ? y + 1 : y);
-    }
-    auto compute_laplacian_scalar = [&](size_t x) {
-      const size_t prevX = x >= 1 ? x - 1 : x;
-      const size_t nextX = x + 1 < opsin.xsize() ? x + 1 : x;
-      float sumsqr = 0;
-      for (size_t c = 0; c < 3; c++) {
-        float laplacian =
-            kChannelWeights[c] * in_row[c][x] +
-            kChannelWeightsLapNeg[c] *
-                (in_row[c][prevX] + in_row[c][nextX] + in_row_b[c][prevX] +
-                 in_row_b[c][x] + in_row_b[c][nextX] + in_row_t[c][prevX] +
-                 in_row_t[c][x] + in_row_t[c][nextX]);
-        sumsqr += laplacian * laplacian;
-      }
-      laplacian_sqrsum_row[x + 2 - bx0 * N] = sumsqr;
-    };
-    size_t x = x0;
-    for (; x < 1; x++) {
-      compute_laplacian_scalar(x);
-    }
-    // Interior. One extra pixel of border as the last pixel is special.
-    for (; x + Lanes(df) <= x1 && x + Lanes(df) + 1 <= opsin.xsize();
-         x += Lanes(df)) {
-      auto sumsqr = Zero(df);
-      for (size_t c = 0; c < 3; c++) {
-        auto laplacian =
-            Mul(LoadU(df, in_row[c] + x), Set(df, kChannelWeights[c]));
-        auto sum_oth0 = LoadU(df, in_row[c] + x - 1);
-        auto sum_oth1 = LoadU(df, in_row[c] + x + 1);
-        auto sum_oth2 = LoadU(df, in_row_t[c] + x - 1);
-        auto sum_oth3 = LoadU(df, in_row_t[c] + x);
-        sum_oth0 = Add(sum_oth0, LoadU(df, in_row_t[c] + x + 1));
-        sum_oth1 = Add(sum_oth1, LoadU(df, in_row_b[c] + x - 1));
-        sum_oth2 = Add(sum_oth2, LoadU(df, in_row_b[c] + x));
-        sum_oth3 = Add(sum_oth3, LoadU(df, in_row_b[c] + x + 1));
-        sum_oth0 = Add(sum_oth0, sum_oth1);
-        sum_oth2 = Add(sum_oth2, sum_oth3);
-        sum_oth0 = Add(sum_oth0, sum_oth2);
-        laplacian =
-            MulAdd(Set(df, kChannelWeightsLapNeg[c]), sum_oth0, laplacian);
-        sumsqr = MulAdd(laplacian, laplacian, sumsqr);
-      }
-      StoreU(sumsqr, df, laplacian_sqrsum_row + x + 2 - bx0 * N);
-    }
-    for (; x < x1; x++) {
-      compute_laplacian_scalar(x);
-    }
-  }
-  HWY_CAPPED(float, 4) df4;
-  // Calculate the L2 of the 3x3 Laplacian in 4x4 blocks within the area
-  // of the integral transform. Sample them within the integral transform
-  // with two offsets (0,0) and (-2, -2) pixels (sqrsum_00 and sqrsum_22,
-  //  respectively).
-  ImageF& sqrsum_00 = temp_image->sqrsum_00;
-  size_t sqrsum_00_stride = sqrsum_00.PixelsPerRow();
-  float* JXL_RESTRICT sqrsum_00_row = sqrsum_00.Row(0);
-  for (size_t y = 0; y < rect.ysize() * 2; y++) {
-    const float* JXL_RESTRICT rows_in[4];
-    for (size_t iy = 0; iy < 4; iy++) {
-      rows_in[iy] = laplacian_sqrsum.ConstRow(y * 4 + iy + 2);
-    }
-    float* JXL_RESTRICT row_out = sqrsum_00_row + y * sqrsum_00_stride;
-    for (size_t x = 0; x < rect.xsize() * 2; x++) {
-      auto sum = Zero(df4);
-      for (auto& row : rows_in) {
-        for (size_t ix = 0; ix < 4; ix += Lanes(df4)) {
-          sum = Add(sum, LoadU(df4, row + x * 4 + ix + 2));
-        }
-      }
-      row_out[x] = GetLane(Sqrt(SumOfLanes(df4, sum))) * (1.0f / 4.0f);
-    }
-  }
-  // Indexing iy and ix is a bit tricky as we include a 2 pixel border
-  // around the block for evenness calculations. This is similar to what
-  // we did in guetzli for the observability of artefacts, except there
-  // the element is a sliding 5x5, not sparsely sampled 4x4 box like here.
-  ImageF& sqrsum_22 = temp_image->sqrsum_22;
-  size_t sqrsum_22_stride = sqrsum_22.PixelsPerRow();
-  float* JXL_RESTRICT sqrsum_22_row = sqrsum_22.Row(0);
-  for (size_t y = 0; y < rect.ysize() * 2 + 1; y++) {
-    const float* JXL_RESTRICT rows_in[4];
-    for (size_t iy = 0; iy < 4; iy++) {
-      rows_in[iy] = laplacian_sqrsum.ConstRow(y * 4 + iy);
-    }
-    float* JXL_RESTRICT row_out = sqrsum_22_row + y * sqrsum_22_stride;
-    // ignore pixels outside the image.
-    // Y coordinates are relative to by0*8+y*4.
-    size_t sy = y * 4 + by0 * 8 > 0 ? 0 : 2;
-    size_t ey = y * 4 + by0 * 8 + 2 <= opsin.ysize()
-                    ? 4
-                    : opsin.ysize() - y * 4 - by0 * 8 + 2;
-    for (size_t x = 0; x < rect.xsize() * 2 + 1; x++) {
-      // ignore pixels outside the image.
-      // X coordinates are relative to bx0*8.
-      size_t sx = x * 4 + bx0 * 8 > 0 ? x * 4 : x * 4 + 2;
-      size_t ex = x * 4 + bx0 * 8 + 2 <= opsin.xsize()
-                      ? x * 4 + 4
-                      : opsin.xsize() - bx0 * 8 + 2;
-      if (ex - sx == 4 && ey - sy == 4) {
-        auto sum = Zero(df4);
-        for (size_t iy = sy; iy < ey; iy++) {
-          for (size_t ix = sx; ix < ex; ix += Lanes(df4)) {
-            sum = Add(sum, Load(df4, rows_in[iy] + ix));
-          }
-        }
-        row_out[x] = GetLane(Sqrt(SumOfLanes(df4, sum))) * (1.0f / 4.0f);
-      } else {
-        float sum = 0;
-        for (size_t iy = sy; iy < ey; iy++) {
-          for (size_t ix = sx; ix < ex; ix++) {
-            sum += rows_in[iy][ix];
-          }
-        }
-        row_out[x] = std::sqrt(sum / ((ex - sx) * (ey - sy)));
-      }
-    }
-  }
-  for (size_t by = rect.y0(); by < rect.y1(); by++) {
-    AcStrategyRow acs_row = ac_strategy.ConstRow(by);
-    uint8_t* JXL_RESTRICT out_row = epf_sharpness->Row(by);
-    const float* JXL_RESTRICT quant_row = quant_field.Row(by);
-    for (size_t bx = rect.x0(); bx < rect.x1(); bx++) {
-      AcStrategy acs = acs_row[bx];
-      if (!acs.IsFirstBlock()) continue;
-      // The errors are going to be linear to the quantization value in this
-      // locality. We only have access to the initial quant field here.
-      float quant_val = 1.0f / quant_row[bx];
-
-      const auto sq00 = [&](size_t y, size_t x) {
-        return sqrsum_00_row[((by - rect.y0()) * 2 + y) * sqrsum_00_stride +
-                             (bx - rect.x0()) * 2 + x];
-      };
-      const auto sq22 = [&](size_t y, size_t x) {
-        return sqrsum_22_row[((by - rect.y0()) * 2 + y) * sqrsum_22_stride +
-                             (bx - rect.x0()) * 2 + x];
-      };
-      float sqrsum_integral_transform = 0;
-      for (size_t iy = 0; iy < acs.covered_blocks_y() * 2; iy++) {
-        for (size_t ix = 0; ix < acs.covered_blocks_x() * 2; ix++) {
-          sqrsum_integral_transform += sq00(iy, ix) * sq00(iy, ix);
-        }
-      }
-      sqrsum_integral_transform /=
-          4 * acs.covered_blocks_x() * acs.covered_blocks_y();
-      sqrsum_integral_transform = std::sqrt(sqrsum_integral_transform);
-      // If masking is high or amplitude of the artefacts is low, then no
-      // smoothing is needed.
-      for (size_t iy = 0; iy < acs.covered_blocks_y(); iy++) {
-        for (size_t ix = 0; ix < acs.covered_blocks_x(); ix++) {
-          // Five 4x4 blocks for masking estimation, all within the
-          // 8x8 area.
-          float minval_1 = std::min(sq00(2 * iy + 0, 2 * ix + 0),
-                                    sq00(2 * iy + 0, 2 * ix + 1));
-          float minval_2 = std::min(sq00(2 * iy + 1, 2 * ix + 0),
-                                    sq00(2 * iy + 1, 2 * ix + 1));
-          float minval = std::min(minval_1, minval_2);
-          minval = std::min(minval, sq22(2 * iy + 1, 2 * ix + 1));
-          // Nine more 4x4 blocks for masking estimation, includes
-          // the 2 pixel area around the 8x8 block being controlled.
-          float minval2_1 = std::min(sq22(2 * iy + 0, 2 * ix + 0),
-                                     sq22(2 * iy + 0, 2 * ix + 1));
-          float minval2_2 = std::min(sq22(2 * iy + 0, 2 * ix + 2),
-                                     sq22(2 * iy + 1, 2 * ix + 0));
-          float minval2_3 = std::min(sq22(2 * iy + 1, 2 * ix + 1),
-                                     sq22(2 * iy + 1, 2 * ix + 2));
-          float minval2_4 = std::min(sq22(2 * iy + 2, 2 * ix + 0),
-                                     sq22(2 * iy + 2, 2 * ix + 1));
-          float minval2_5 = std::min(minval2_1, minval2_2);
-          float minval2_6 = std::min(minval2_3, minval2_4);
-          float minval2 = std::min(minval2_5, minval2_6);
-          minval2 = std::min(minval2, sq22(2 * iy + 2, 2 * ix + 2));
-          float minval3 = std::min(minval, minval2);
-          minval *= 0.125f;
-          minval += 0.625f * minval3;
-          minval +=
-              0.125f * std::min(1.5f * minval3, sq22(2 * iy + 1, 2 * ix + 1));
-          minval += 0.125f * minval2;
-          // Larger kBias, less smoothing for low intensity changes.
-          float kDeltaLimit = 3.2;
-          float bias = 0.0625f * quant_val;
-          float delta =
-              (sqrsum_integral_transform + (kDeltaLimit + 0.05) * bias) /
-              (minval + bias);
-          int out = 4;
-          if (delta > kDeltaLimit) {
-            out = 4;  // smooth
-          } else {
-            out = 0;
-          }
-          // 'threshold' is separate from 'bias' for easier tuning of these
-          // heuristics.
-          float threshold = 0.0625f * quant_val;
-          const float kSmoothLimit = 0.085f;
-          float smooth = 0.20f * (sq00(2 * iy + 0, 2 * ix + 0) +
-                                  sq00(2 * iy + 0, 2 * ix + 1) +
-                                  sq00(2 * iy + 1, 2 * ix + 0) +
-                                  sq00(2 * iy + 1, 2 * ix + 1) + minval);
-          if (smooth < kSmoothLimit * threshold) {
-            out = 4;
-          }
-          out_row[bx + sharpness_stride * iy + ix] = out;
-        }
-      }
-    }
-  }
-  return true;
-}
-
-}  // namespace
-// NOLINTNEXTLINE(google-readability-namespace-comments)
-}  // namespace HWY_NAMESPACE
-}  // namespace jxl
-HWY_AFTER_NAMESPACE();
-
-#if HWY_ONCE
-namespace jxl {
-HWY_EXPORT(ProcessTile);
-
-Status ArControlFieldHeuristics::RunRect(
-    const CompressParams& cparams, const FrameHeader& frame_header,
-    const Rect& block_rect, const Image3F& opsin, const Rect& opsin_rect,
-    const ImageF& quant_field, const AcStrategyImage& ac_strategy,
-    ImageB* epf_sharpness, size_t thread) {
-  return HWY_DYNAMIC_DISPATCH(ProcessTile)(
-      cparams, frame_header, opsin, opsin_rect, quant_field, ac_strategy,
-      epf_sharpness, block_rect, &temp_images[thread]);
-}
-
-}  // namespace jxl
-
-#endif

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

@@ -1,56 +0,0 @@
-// 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_ENC_AR_CONTROL_FIELD_H_
-#define LIB_JXL_ENC_AR_CONTROL_FIELD_H_
-
-#include <cstddef>
-#include <vector>
-
-#include "lib/jxl/ac_strategy.h"
-#include "lib/jxl/base/rect.h"
-#include "lib/jxl/base/status.h"
-#include "lib/jxl/enc_params.h"
-#include "lib/jxl/frame_header.h"
-#include "lib/jxl/image.h"
-
-namespace jxl {
-
-struct PassesEncoderState;
-
-struct ArControlFieldHeuristics {
-  struct TempImages {
-    Status InitOnce() {
-      if (laplacian_sqrsum.xsize() != 0) return true;
-      JXL_ASSIGN_OR_RETURN(laplacian_sqrsum,
-                           ImageF::Create(kEncTileDim + 4, kEncTileDim + 4));
-      JXL_ASSIGN_OR_RETURN(sqrsum_00,
-                           ImageF::Create(kEncTileDim / 4, kEncTileDim / 4));
-      JXL_ASSIGN_OR_RETURN(
-          sqrsum_22, ImageF::Create(kEncTileDim / 4 + 1, kEncTileDim / 4 + 1));
-      return true;
-    }
-
-    ImageF laplacian_sqrsum;
-    ImageF sqrsum_00;
-    ImageF sqrsum_22;
-  };
-
-  void PrepareForThreads(size_t num_threads) {
-    temp_images.resize(num_threads);
-  }
-
-  Status RunRect(const CompressParams& cparams, const FrameHeader& frame_header,
-                 const Rect& block_rect, const Image3F& opsin,
-                 const Rect& opsin_rect, const ImageF& quant_field,
-                 const AcStrategyImage& ac_strategy, ImageB* epf_sharpness,
-                 size_t thread);
-
-  std::vector<TempImages> temp_images;
-};
-
-}  // namespace jxl
-
-#endif  // LIB_JXL_AR_ENC_CONTROL_FIELD_H_

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

@@ -6,11 +6,11 @@
 #include "lib/jxl/enc_bit_writer.h"
 
 #include <jxl/types.h>
-#include <string.h>  // memcpy
+
+#include <cstring>  // memcpy
 
 #include "lib/jxl/base/byte_order.h"
 #include "lib/jxl/base/printf_macros.h"
-#include "lib/jxl/dec_bit_reader.h"
 #include "lib/jxl/enc_aux_out.h"
 
 namespace jxl {

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

@@ -8,9 +8,10 @@
 
 // BitWriter class: unbuffered writes using unaligned 64-bit stores.
 
-#include <stddef.h>
+#include <jxl/memory_manager.h>
-#include <stdint.h>
 
+#include <cstddef>
+#include <cstdint>
 #include <utility>
 #include <vector>
 
@@ -32,7 +33,8 @@ struct BitWriter {
   // yet zero-initialized).
   static constexpr size_t kMaxBitsPerCall = 56;
 
-  BitWriter() : bits_written_(0) {}
+  explicit BitWriter(JxlMemoryManager* memory_manager)
+      : bits_written_(0), storage_(memory_manager) {}
 
   // Disallow copying - may lead to bugs.
   BitWriter(const BitWriter&) = delete;
@@ -42,6 +44,8 @@ struct BitWriter {
 
   size_t BitsWritten() const { return bits_written_; }
 
+  JxlMemoryManager* memory_manager() const { return storage_.memory_manager(); }
+
   Span<const uint8_t> GetSpan() const {
     // Callers must ensure byte alignment to avoid uninitialized bits.
     JXL_ASSERT(bits_written_ % kBitsPerByte == 0);

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

@@ -16,8 +16,9 @@ JxlButteraugliComparator::JxlButteraugliComparator(
 
 Status JxlButteraugliComparator::SetReferenceImage(const ImageBundle& ref) {
   const ImageBundle* ref_linear_srgb;
+  JxlMemoryManager* memory_manager = ref.memory_manager();
   ImageMetadata metadata = *ref.metadata();
-  ImageBundle store(&metadata);
+  ImageBundle store(memory_manager, &metadata);
   if (!TransformIfNeeded(ref, ColorEncoding::LinearSRGB(ref.IsGray()), cms_,
                          /*pool=*/nullptr, &store, &ref_linear_srgb)) {
     return false;
@@ -46,17 +47,19 @@ Status JxlButteraugliComparator::CompareWith(const ImageBundle& actual,
   if (xsize_ != actual.xsize() || ysize_ != actual.ysize()) {
     return JXL_FAILURE("Images must have same size");
   }
+  JxlMemoryManager* memory_manager = actual.memory_manager();
 
   const ImageBundle* actual_linear_srgb;
   ImageMetadata metadata = *actual.metadata();
-  ImageBundle store(&metadata);
+  ImageBundle store(memory_manager, &metadata);
   if (!TransformIfNeeded(actual, ColorEncoding::LinearSRGB(actual.IsGray()),
                          cms_,
                          /*pool=*/nullptr, &store, &actual_linear_srgb)) {
     return false;
   }
 
-  JXL_ASSIGN_OR_RETURN(ImageF temp_diffmap, ImageF::Create(xsize_, ysize_));
+  JXL_ASSIGN_OR_RETURN(ImageF temp_diffmap,
+                       ImageF::Create(memory_manager, xsize_, ysize_));
   JXL_RETURN_IF_ERROR(
       comparator_->Diffmap(actual_linear_srgb->color(), temp_diffmap));
 

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

@@ -33,6 +33,31 @@
 
 namespace jxl {
 
+Status ComputeACMetadata(ThreadPool* pool, PassesEncoderState* enc_state,
+                         ModularFrameEncoder* modular_frame_encoder) {
+  PassesSharedState& shared = enc_state->shared;
+  std::atomic<bool> has_error{false};
+  auto compute_ac_meta = [&](int group_index, int /* thread */) {
+    const Rect r = shared.frame_dim.DCGroupRect(group_index);
+    int modular_group_index = group_index;
+    if (enc_state->streaming_mode) {
+      JXL_ASSERT(group_index == 0);
+      modular_group_index = enc_state->dc_group_index;
+    }
+    if (!modular_frame_encoder->AddACMetadata(r, modular_group_index,
+                                              /*jpeg_transcode=*/false,
+                                              enc_state)) {
+      has_error = true;
+      return;
+    }
+  };
+  JXL_RETURN_IF_ERROR(RunOnPool(pool, 0, shared.frame_dim.num_dc_groups,
+                                ThreadPool::NoInit, compute_ac_meta,
+                                "Compute AC Metadata"));
+  if (has_error) return JXL_FAILURE("Compute AC Metadata failed");
+  return true;
+}
+
 Status InitializePassesEncoder(const FrameHeader& frame_header,
                                const Image3F& opsin, const Rect& rect,
                                const JxlCmsInterface& cms, ThreadPool* pool,
@@ -40,6 +65,7 @@ Status InitializePassesEncoder(const FrameHeader& frame_header,
                                ModularFrameEncoder* modular_frame_encoder,
                                AuxOut* aux_out) {
   PassesSharedState& JXL_RESTRICT shared = enc_state->shared;
+  JxlMemoryManager* memory_manager = enc_state->memory_manager();
 
   enc_state->x_qm_multiplier = std::pow(1.25f, frame_header.x_qm_scale - 2.0f);
   enc_state->b_qm_multiplier = std::pow(1.25f, frame_header.b_qm_scale - 2.0f);
@@ -51,7 +77,7 @@ Status InitializePassesEncoder(const FrameHeader& frame_header,
       // Allocate enough coefficients for each group on every row.
       JXL_ASSIGN_OR_RETURN(
           std::unique_ptr<ACImageT<int32_t>> coeffs,
-          ACImageT<int32_t>::Make(kGroupDim * kGroupDim,
+          ACImageT<int32_t>::Make(memory_manager, kGroupDim * kGroupDim,
                                   shared.frame_dim.num_groups));
       enc_state->coeffs.emplace_back(std::move(coeffs));
     }
@@ -63,13 +89,13 @@ Status InitializePassesEncoder(const FrameHeader& frame_header,
   if (enc_state->initialize_global_state) {
     float scale =
         shared.quantizer.ScaleGlobalScale(enc_state->cparams.quant_ac_rescale);
-    DequantMatricesScaleDC(&shared.matrices, scale);
+    DequantMatricesScaleDC(memory_manager, &shared.matrices, scale);
     shared.quantizer.RecomputeFromGlobalScale();
   }
 
-  JXL_ASSIGN_OR_RETURN(Image3F dc,
+  JXL_ASSIGN_OR_RETURN(
-                       Image3F::Create(shared.frame_dim.xsize_blocks,
+      Image3F dc, Image3F::Create(memory_manager, shared.frame_dim.xsize_blocks,
-                                       shared.frame_dim.ysize_blocks));
+                                  shared.frame_dim.ysize_blocks));
   JXL_RETURN_IF_ERROR(RunOnPool(
       pool, 0, shared.frame_dim.num_groups, ThreadPool::NoInit,
       [&](size_t group_idx, size_t _) {
@@ -110,7 +136,7 @@ Status InitializePassesEncoder(const FrameHeader& frame_header,
           std::max(kMinButteraugliDistance,
                    enc_state->cparams.butteraugli_distance * 0.1f);
     }
-    ImageBundle ib(&shared.metadata->m);
+    ImageBundle ib(memory_manager, &shared.metadata->m);
     // This is a lie - dc is in XYB
     // (but EncodeFrame will skip RGB->XYB conversion anyway)
     ib.SetFromImage(
@@ -124,7 +150,8 @@ Status InitializePassesEncoder(const FrameHeader& frame_header,
       std::vector<ImageF> extra_channels;
       extra_channels.reserve(ib.metadata()->extra_channel_info.size());
       for (size_t i = 0; i < ib.metadata()->extra_channel_info.size(); i++) {
-        JXL_ASSIGN_OR_RETURN(ImageF ch, ImageF::Create(ib.xsize(), ib.ysize()));
+        JXL_ASSIGN_OR_RETURN(
+            ImageF ch, ImageF::Create(memory_manager, ib.xsize(), ib.ysize()));
         extra_channels.emplace_back(std::move(ch));
         // Must initialize the image with data to not affect blending with
         // uninitialized memory.
@@ -134,15 +161,16 @@ Status InitializePassesEncoder(const FrameHeader& frame_header,
       }
       ib.SetExtraChannels(std::move(extra_channels));
     }
-    auto special_frame = std::unique_ptr<BitWriter>(new BitWriter());
+    auto special_frame =
+        std::unique_ptr<BitWriter>(new BitWriter(memory_manager));
     FrameInfo dc_frame_info;
     dc_frame_info.frame_type = FrameType::kDCFrame;
     dc_frame_info.dc_level = frame_header.dc_level + 1;
     dc_frame_info.ib_needs_color_transform = false;
     dc_frame_info.save_before_color_transform = true;  // Implicitly true
     AuxOut dc_aux_out;
-    JXL_CHECK(EncodeFrame(cparams, dc_frame_info, shared.metadata, ib, cms,
+    JXL_CHECK(EncodeFrame(memory_manager, cparams, dc_frame_info,
-                          pool, special_frame.get(),
+                          shared.metadata, ib, cms, pool, special_frame.get(),
                           aux_out ? &dc_aux_out : nullptr));
     if (aux_out) {
       for (const auto& l : dc_aux_out.layers) {
@@ -152,9 +180,9 @@ Status InitializePassesEncoder(const FrameHeader& frame_header,
     const Span<const uint8_t> encoded = special_frame->GetSpan();
     enc_state->special_frames.emplace_back(std::move(special_frame));
 
-    ImageBundle decoded(&shared.metadata->m);
+    ImageBundle decoded(memory_manager, &shared.metadata->m);
     std::unique_ptr<PassesDecoderState> dec_state =
-        jxl::make_unique<PassesDecoderState>();
+        jxl::make_unique<PassesDecoderState>(memory_manager);
     JXL_CHECK(
         dec_state->output_encoding_info.SetFromMetadata(*shared.metadata));
     const uint8_t* frame_start = encoded.data();
@@ -173,8 +201,9 @@ Status InitializePassesEncoder(const FrameHeader& frame_header,
     // outputs multiple frames, this assumption could be wrong.
     const Image3F& dc_frame =
         dec_state->shared->dc_frames[frame_header.dc_level];
-    JXL_ASSIGN_OR_RETURN(shared.dc_storage,
+    JXL_ASSIGN_OR_RETURN(
-                         Image3F::Create(dc_frame.xsize(), dc_frame.ysize()));
+        shared.dc_storage,
+        Image3F::Create(memory_manager, dc_frame.xsize(), dc_frame.ysize()));
     CopyImageTo(dc_frame, &shared.dc_storage);
     ZeroFillImage(&shared.quant_dc);
     shared.dc = &shared.dc_storage;
@@ -203,29 +232,10 @@ Status InitializePassesEncoder(const FrameHeader& frame_header,
     if (has_error) return JXL_FAILURE("Compute DC coeffs failed");
     // TODO(veluca): this is only useful in tests and if inspection is enabled.
     if (!(frame_header.flags & FrameHeader::kSkipAdaptiveDCSmoothing)) {
-      JXL_RETURN_IF_ERROR(AdaptiveDCSmoothing(shared.quantizer.MulDC(),
+      JXL_RETURN_IF_ERROR(AdaptiveDCSmoothing(
-                                              &shared.dc_storage, pool));
+          memory_manager, shared.quantizer.MulDC(), &shared.dc_storage, pool));
     }
   }
-  std::atomic<bool> has_error{false};
-  auto compute_ac_meta = [&](int group_index, int /* thread */) {
-    const Rect r = enc_state->shared.frame_dim.DCGroupRect(group_index);
-    int modular_group_index = group_index;
-    if (enc_state->streaming_mode) {
-      JXL_ASSERT(group_index == 0);
-      modular_group_index = enc_state->dc_group_index;
-    }
-    if (!modular_frame_encoder->AddACMetadata(r, modular_group_index,
-                                              /*jpeg_transcode=*/false,
-                                              enc_state)) {
-      has_error = true;
-      return;
-    }
-  };
-  JXL_RETURN_IF_ERROR(RunOnPool(pool, 0, shared.frame_dim.num_dc_groups,
-                                ThreadPool::NoInit, compute_ac_meta,
-                                "Compute AC Metadata"));
-  if (has_error) return JXL_FAILURE("Compute AC Metadata failed");
   return true;
 }
 

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

@@ -7,6 +7,7 @@
 #define LIB_JXL_ENC_CACHE_H_
 
 #include <jxl/cms_interface.h>
+#include <jxl/memory_manager.h>
 
 #include <cstddef>
 #include <cstdint>
@@ -32,6 +33,9 @@ struct AuxOut;
 
 // Contains encoder state.
 struct PassesEncoderState {
+  explicit PassesEncoderState(JxlMemoryManager* memory_manager)
+      : shared(memory_manager) {}
+
   PassesSharedState shared;
 
   bool streaming_mode = false;
@@ -66,6 +70,10 @@ struct PassesEncoderState {
   // Multiplier to be applied to the quant matrices of the x channel.
   float x_qm_multiplier = 1.0f;
   float b_qm_multiplier = 1.0f;
+
+  ImageF initial_quant_masking1x1;
+
+  JxlMemoryManager* memory_manager() const { return shared.memory_manager; }
 };
 
 // Initialize per-frame information.
@@ -77,6 +85,9 @@ Status InitializePassesEncoder(const FrameHeader& frame_header,
                                ModularFrameEncoder* modular_frame_encoder,
                                AuxOut* aux_out);
 
+Status ComputeACMetadata(ThreadPool* pool, PassesEncoderState* enc_state,
+                         ModularFrameEncoder* modular_frame_encoder);
+
 }  // namespace jxl
 
 #endif  // LIB_JXL_ENC_CACHE_H_

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

@@ -5,6 +5,8 @@
 
 #include "lib/jxl/enc_chroma_from_luma.h"
 
+#include <jxl/memory_manager.h>
+
 #include <algorithm>
 #include <cfloat>
 #include <cmath>
@@ -170,13 +172,15 @@ int32_t FindBestMultiplier(const float* values_m, const float* values_s,
   return std::max(-128.0f, std::min(127.0f, roundf(x)));
 }
 
-Status InitDCStorage(size_t num_blocks, ImageF* dc_values) {
+Status InitDCStorage(JxlMemoryManager* memory_manager, size_t num_blocks,
+                     ImageF* dc_values) {
   // First row: Y channel
   // Second row: X channel
   // Third row: Y channel
   // Fourth row: B channel
-  JXL_ASSIGN_OR_RETURN(*dc_values,
+  JXL_ASSIGN_OR_RETURN(
-                       ImageF::Create(RoundUpTo(num_blocks, Lanes(df)), 4));
+      *dc_values,
+      ImageF::Create(memory_manager, RoundUpTo(num_blocks, Lanes(df)), 4));
 
   JXL_ASSERT(dc_values->xsize() != 0);
   // Zero-fill the last lanes
@@ -349,11 +353,11 @@ namespace jxl {
 HWY_EXPORT(InitDCStorage);
 HWY_EXPORT(ComputeTile);
 
-Status CfLHeuristics::Init(const Rect& rect) {
+Status CfLHeuristics::Init(JxlMemoryManager* memory_manager, const Rect& rect) {
   size_t xsize_blocks = rect.xsize() / kBlockDim;
   size_t ysize_blocks = rect.ysize() / kBlockDim;
-  return HWY_DYNAMIC_DISPATCH(InitDCStorage)(xsize_blocks * ysize_blocks,
+  return HWY_DYNAMIC_DISPATCH(InitDCStorage)(
-                                             &dc_values);
+      memory_manager, xsize_blocks * ysize_blocks, &dc_values);
 }
 
 void CfLHeuristics::ComputeTile(const Rect& r, const Image3F& opsin,

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

@@ -9,6 +9,8 @@
 // Chroma-from-luma, computed using heuristics to determine the best linear
 // model for the X and B channels from the Y channel.
 
+#include <jxl/memory_manager.h>
+
 #include <cstddef>
 #include <hwy/aligned_allocator.h>
 
@@ -31,7 +33,7 @@ void ColorCorrelationMapEncodeDC(const ColorCorrelationMap& map,
                                  AuxOut* aux_out);
 
 struct CfLHeuristics {
-  Status Init(const Rect& rect);
+  Status Init(JxlMemoryManager* memory_manager, const Rect& rect);
 
   void PrepareForThreads(size_t num_threads) {
     mem = hwy::AllocateAligned<float>(num_threads * ItemsPerThread());

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

@@ -3,6 +3,8 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+#include <jxl/memory_manager.h>
+
 #include <algorithm>
 #include <cmath>
 #include <cstdint>
@@ -234,12 +236,14 @@ void TokenizePermutation(const coeff_order_t* JXL_RESTRICT order, size_t skip,
 void EncodePermutation(const coeff_order_t* JXL_RESTRICT order, size_t skip,
                        size_t size, BitWriter* writer, int layer,
                        AuxOut* aux_out) {
+  JxlMemoryManager* memory_manager = writer->memory_manager();
   std::vector<std::vector<Token>> tokens(1);
   TokenizePermutation(order, skip, size, tokens.data());
   std::vector<uint8_t> context_map;
   EntropyEncodingData codes;
-  BuildAndEncodeHistograms(HistogramParams(), kPermutationContexts, tokens,
+  BuildAndEncodeHistograms(memory_manager, HistogramParams(),
-                           &codes, &context_map, writer, layer, aux_out);
+                           kPermutationContexts, tokens, &codes, &context_map,
+                           writer, layer, aux_out);
   WriteTokens(tokens[0], codes, context_map, 0, writer, layer, aux_out);
 }
 
@@ -260,6 +264,7 @@ void EncodeCoeffOrders(uint16_t used_orders,
                        const coeff_order_t* JXL_RESTRICT order,
                        BitWriter* writer, size_t layer,
                        AuxOut* JXL_RESTRICT aux_out) {
+  JxlMemoryManager* memory_manager = writer->memory_manager();
   auto mem = hwy::AllocateAligned<coeff_order_t>(AcStrategy::kMaxCoeffArea);
   uint16_t computed = 0;
   std::vector<std::vector<Token>> tokens(1);
@@ -283,8 +288,9 @@ void EncodeCoeffOrders(uint16_t used_orders,
   if (used_orders != 0) {
     std::vector<uint8_t> context_map;
     EntropyEncodingData codes;
-    BuildAndEncodeHistograms(HistogramParams(), kPermutationContexts, tokens,
+    BuildAndEncodeHistograms(memory_manager, HistogramParams(),
-                             &codes, &context_map, writer, layer, aux_out);
+                             kPermutationContexts, tokens, &codes, &context_map,
+                             writer, layer, aux_out);
     WriteTokens(tokens[0], codes, context_map, 0, writer, layer, aux_out);
   }
 }

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

@@ -5,10 +5,10 @@
 
 #include "lib/jxl/enc_comparator.h"
 
-#include <stddef.h>
+#include <jxl/memory_manager.h>
-#include <stdint.h>
 
 #include <algorithm>
+#include <cstddef>
 
 #include "lib/jxl/base/compiler_specific.h"
 #include "lib/jxl/enc_gamma_correct.h"
@@ -70,14 +70,15 @@ Status ComputeScore(const ImageBundle& rgb0, const ImageBundle& rgb1,
                     Comparator* comparator, const JxlCmsInterface& cms,
                     float* score, ImageF* diffmap, ThreadPool* pool,
                     bool ignore_alpha) {
+  JxlMemoryManager* memory_manager = rgb0.memory_manager();
   // Convert to linear sRGB (unless already in that space)
   ImageMetadata metadata0 = *rgb0.metadata();
-  ImageBundle store0(&metadata0);
+  ImageBundle store0(memory_manager, &metadata0);
   const ImageBundle* linear_srgb0;
   JXL_CHECK(TransformIfNeeded(rgb0, ColorEncoding::LinearSRGB(rgb0.IsGray()),
                               cms, pool, &store0, &linear_srgb0));
   ImageMetadata metadata1 = *rgb1.metadata();
-  ImageBundle store1(&metadata1);
+  ImageBundle store1(memory_manager, &metadata1);
   const ImageBundle* linear_srgb1;
   JXL_CHECK(TransformIfNeeded(rgb1, ColorEncoding::LinearSRGB(rgb1.IsGray()),
                               cms, pool, &store1, &linear_srgb1));
@@ -115,7 +116,8 @@ Status ComputeScore(const ImageBundle& rgb0, const ImageBundle& rgb1,
   if (diffmap != nullptr) {
     const size_t xsize = rgb0.xsize();
     const size_t ysize = rgb0.ysize();
-    JXL_ASSIGN_OR_RETURN(*diffmap, ImageF::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(*diffmap,
+                         ImageF::Create(memory_manager, xsize, ysize));
     for (size_t y = 0; y < ysize; ++y) {
       const float* JXL_RESTRICT row_black = diffmap_black.ConstRow(y);
       const float* JXL_RESTRICT row_white = diffmap_white.ConstRow(y);

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

@@ -7,6 +7,7 @@
 
 #include "lib/jxl/enc_context_map.h"
 
+#include <jxl/memory_manager.h>
 #include <jxl/types.h>
 #include <stdint.h>
 
@@ -70,6 +71,7 @@ void EncodeContextMap(const std::vector<uint8_t>& context_map,
     return;
   }
 
+  JxlMemoryManager* memory_manager = writer->memory_manager();
   std::vector<uint8_t> transformed_symbols = MoveToFrontTransform(context_map);
   std::vector<std::vector<Token>> tokens(1);
   std::vector<std::vector<Token>> mtf_tokens(1);
@@ -86,14 +88,16 @@ void EncodeContextMap(const std::vector<uint8_t>& context_map,
   {
     EntropyEncodingData codes;
     std::vector<uint8_t> sink_context_map;
-    ans_cost = BuildAndEncodeHistograms(params, 1, tokens, &codes,
+    ans_cost =
-                                        &sink_context_map, nullptr, 0, nullptr);
+        BuildAndEncodeHistograms(memory_manager, params, 1, tokens, &codes,
+                                 &sink_context_map, nullptr, 0, nullptr);
   }
   {
     EntropyEncodingData codes;
     std::vector<uint8_t> sink_context_map;
-    mtf_cost = BuildAndEncodeHistograms(params, 1, mtf_tokens, &codes,
+    mtf_cost =
-                                        &sink_context_map, nullptr, 0, nullptr);
+        BuildAndEncodeHistograms(memory_manager, params, 1, mtf_tokens, &codes,
+                                 &sink_context_map, nullptr, 0, nullptr);
   }
   bool use_mtf = mtf_cost < ans_cost;
   // Rebuild token list.
@@ -118,8 +122,8 @@ void EncodeContextMap(const std::vector<uint8_t>& context_map,
     writer->Write(1, TO_JXL_BOOL(use_mtf));  // Use/don't use MTF.
     EntropyEncodingData codes;
     std::vector<uint8_t> sink_context_map;
-    BuildAndEncodeHistograms(params, 1, tokens, &codes, &sink_context_map,
+    BuildAndEncodeHistograms(memory_manager, params, 1, tokens, &codes,
-                             writer, layer, aux_out);
+                             &sink_context_map, writer, layer, aux_out);
     WriteTokens(tokens[0], codes, sink_context_map, 0, writer);
     allotment.ReclaimAndCharge(writer, layer, aux_out);
   }

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

@@ -6,9 +6,8 @@
 #ifndef LIB_JXL_ENC_CONTEXT_MAP_H_
 #define LIB_JXL_ENC_CONTEXT_MAP_H_
 
-#include <stddef.h>
+#include <cstddef>
-#include <stdint.h>
+#include <cstdint>
-
 #include <vector>
 
 #include "lib/jxl/ac_context.h"

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

@@ -5,6 +5,8 @@
 
 #include "lib/jxl/enc_debug_image.h"
 
+#include <jxl/memory_manager.h>
+
 #include <cstddef>
 #include <cstdint>
 
@@ -24,7 +26,9 @@ StatusOr<Image3F> ConvertToFloat(const Image3<From>& from) {
   if (std::is_same<From, double>::value || std::is_same<From, float>::value) {
     factor = 1.0f;
   }
-  JXL_ASSIGN_OR_RETURN(Image3F to, Image3F::Create(from.xsize(), from.ysize()));
+  JxlMemoryManager* memory_manager = from.memory_manager();
+  JXL_ASSIGN_OR_RETURN(
+      Image3F to, Image3F::Create(memory_manager, from.xsize(), from.ysize()));
   for (size_t c = 0; c < 3; ++c) {
     for (size_t y = 0; y < from.ysize(); ++y) {
       const From* const JXL_RESTRICT row_from = from.ConstPlaneRow(c, y);
@@ -63,8 +67,11 @@ Status DumpPlaneNormalizedT(const CompressParams& cparams, const char* label,
   T min;
   T max;
   ImageMinMax(image, &min, &max);
-  JXL_ASSIGN_OR_RETURN(Image3B normalized,
+  JxlMemoryManager* memory_manager = image.memory_manager();
-                       Image3B::Create(image.xsize(), image.ysize()));
+
+  JXL_ASSIGN_OR_RETURN(
+      Image3B normalized,
+      Image3B::Create(memory_manager, image.xsize(), image.ysize()));
   for (size_t c = 0; c < 3; ++c) {
     float mul = min == max ? 0 : (255.0f / (max - min));
     for (size_t y = 0; y < image.ysize(); ++y) {
@@ -93,9 +100,11 @@ Status DumpImage(const CompressParams& cparams, const char* label,
 Status DumpXybImage(const CompressParams& cparams, const char* label,
                     const Image3F& image) {
   if (!cparams.debug_image) return true;
+  JxlMemoryManager* memory_manager = image.memory_manager();
 
-  JXL_ASSIGN_OR_RETURN(Image3F linear,
+  JXL_ASSIGN_OR_RETURN(
-                       Image3F::Create(image.xsize(), image.ysize()));
+      Image3F linear,
+      Image3F::Create(memory_manager, image.xsize(), image.ysize()));
   OpsinParams opsin_params;
   opsin_params.Init(kDefaultIntensityTarget);
   OpsinToLinear(image, Rect(linear), nullptr, &linear, opsin_params);

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

@@ -5,6 +5,8 @@
 
 #include "lib/jxl/enc_detect_dots.h"
 
+#include <jxl/memory_manager.h>
+
 #include <algorithm>
 #include <array>
 #include <cmath>
@@ -50,9 +52,11 @@ StatusOr<ImageF> SumOfSquareDifferences(const Image3F& forig,
   const auto color_coef0 = Set(d, 0.0f);
   const auto color_coef1 = Set(d, 10.0f);
   const auto color_coef2 = Set(d, 0.0f);
+  JxlMemoryManager* memory_manager = forig.memory_manager();
 
-  JXL_ASSIGN_OR_RETURN(ImageF sum_of_squares,
+  JXL_ASSIGN_OR_RETURN(
-                       ImageF::Create(forig.xsize(), forig.ysize()));
+      ImageF sum_of_squares,
+      ImageF::Create(memory_manager, forig.xsize(), forig.ysize()));
   JXL_CHECK(RunOnPool(
       pool, 0, forig.ysize(), ThreadPool::NoInit,
       [&](const uint32_t task, size_t thread) {
@@ -145,10 +149,13 @@ const WeightsSeparable5& WeightsSeparable5Gaussian3() {
 
 StatusOr<ImageF> ComputeEnergyImage(const Image3F& orig, Image3F* smooth,
                                     ThreadPool* pool) {
+  JxlMemoryManager* memory_manager = orig.memory_manager();
   // Prepare guidance images for dot selection.
-  JXL_ASSIGN_OR_RETURN(Image3F forig,
+  JXL_ASSIGN_OR_RETURN(
-                       Image3F::Create(orig.xsize(), orig.ysize()));
+      Image3F forig,
-  JXL_ASSIGN_OR_RETURN(*smooth, Image3F::Create(orig.xsize(), orig.ysize()));
+      Image3F::Create(memory_manager, orig.xsize(), orig.ysize()));
+  JXL_ASSIGN_OR_RETURN(
+      *smooth, Image3F::Create(memory_manager, orig.xsize(), orig.ysize()));
   Rect rect(orig);
 
   const auto& weights1 = WeightsSeparable5Gaussian0_65();
@@ -289,8 +296,10 @@ StatusOr<std::vector<ConnectedComponent>> FindCC(const ImageF& energy,
                                                  uint32_t maxWindow,
                                                  double minScore) {
   const int kExtraRect = 4;
-  JXL_ASSIGN_OR_RETURN(ImageF img,
+  JxlMemoryManager* memory_manager = energy.memory_manager();
-                       ImageF::Create(energy.xsize(), energy.ysize()));
+  JXL_ASSIGN_OR_RETURN(
+      ImageF img,
+      ImageF::Create(memory_manager, energy.xsize(), energy.ysize()));
   CopyImageTo(energy, &img);
   std::vector<ConnectedComponent> ans;
   for (size_t y = 0; y < rect.ysize(); y++) {
@@ -537,9 +546,11 @@ GaussianEllipse FitGaussian(const ConnectedComponent& cc, const Rect& rect,
 StatusOr<std::vector<PatchInfo>> DetectGaussianEllipses(
     const Image3F& opsin, const Rect& rect, const GaussianDetectParams& params,
     const EllipseQuantParams& qParams, ThreadPool* pool) {
+  JxlMemoryManager* memory_manager = opsin.memory_manager();
   std::vector<PatchInfo> dots;
-  JXL_ASSIGN_OR_RETURN(Image3F smooth,
+  JXL_ASSIGN_OR_RETURN(
-                       Image3F::Create(opsin.xsize(), opsin.ysize()));
+      Image3F smooth,
+      Image3F::Create(memory_manager, opsin.xsize(), opsin.ysize()));
   JXL_ASSIGN_OR_RETURN(ImageF energy, ComputeEnergyImage(opsin, &smooth, pool));
   JXL_ASSIGN_OR_RETURN(std::vector<ConnectedComponent> components,
                        FindCC(energy, rect, params.t_low, params.t_high,

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

@@ -5,10 +5,11 @@
 
 #include "lib/jxl/enc_external_image.h"
 
+#include <jxl/memory_manager.h>
 #include <jxl/types.h>
-#include <string.h>
 
 #include <atomic>
+#include <cstring>
 #include <utility>
 
 #include "lib/jxl/base/byte_order.h"
@@ -102,6 +103,7 @@ Status ConvertFromExternalNoSizeCheck(const uint8_t* data, size_t xsize,
                                       size_t bits_per_sample,
                                       JxlPixelFormat format, ThreadPool* pool,
                                       ImageBundle* ib) {
+  JxlMemoryManager* memory_manager = ib->memory_manager();
   bool has_alpha = format.num_channels == 2 || format.num_channels == 4;
   if (format.num_channels < color_channels) {
     return JXL_FAILURE("Expected %" PRIuS
@@ -109,7 +111,8 @@ Status ConvertFromExternalNoSizeCheck(const uint8_t* data, size_t xsize,
                        color_channels, format.num_channels);
   }
 
-  JXL_ASSIGN_OR_RETURN(Image3F color, Image3F::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(Image3F color,
+                       Image3F::Create(memory_manager, xsize, ysize));
   for (size_t c = 0; c < color_channels; ++c) {
     JXL_RETURN_IF_ERROR(ConvertFromExternalNoSizeCheck(
         data, xsize, ysize, stride, bits_per_sample, format, c, pool,
@@ -124,7 +127,8 @@ Status ConvertFromExternalNoSizeCheck(const uint8_t* data, size_t xsize,
   // Passing an interleaved image with an alpha channel to an image that doesn't
   // have alpha channel just discards the passed alpha channel.
   if (has_alpha && ib->HasAlpha()) {
-    JXL_ASSIGN_OR_RETURN(ImageF alpha, ImageF::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(ImageF alpha,
+                         ImageF::Create(memory_manager, xsize, ysize));
     JXL_RETURN_IF_ERROR(ConvertFromExternalNoSizeCheck(
         data, xsize, ysize, stride, bits_per_sample, format,
         format.num_channels - 1, pool, &alpha));
@@ -132,7 +136,8 @@ Status ConvertFromExternalNoSizeCheck(const uint8_t* data, size_t xsize,
   } else if (!has_alpha && ib->HasAlpha()) {
     // if alpha is not passed, but it is expected, then assume
     // it is all-opaque
-    JXL_ASSIGN_OR_RETURN(ImageF alpha, ImageF::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(ImageF alpha,
+                         ImageF::Create(memory_manager, xsize, ysize));
     FillImage(1.0f, &alpha);
     ib->SetAlpha(std::move(alpha));
   }
@@ -172,6 +177,7 @@ Status ConvertFromExternal(Span<const uint8_t> bytes, size_t xsize,
                            size_t color_channels, size_t bits_per_sample,
                            JxlPixelFormat format, ThreadPool* pool,
                            ImageBundle* ib) {
+  JxlMemoryManager* memory_manager = ib->memory_manager();
   bool has_alpha = format.num_channels == 2 || format.num_channels == 4;
   if (format.num_channels < color_channels) {
     return JXL_FAILURE("Expected %" PRIuS
@@ -179,7 +185,8 @@ Status ConvertFromExternal(Span<const uint8_t> bytes, size_t xsize,
                        color_channels, format.num_channels);
   }
 
-  JXL_ASSIGN_OR_RETURN(Image3F color, Image3F::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(Image3F color,
+                       Image3F::Create(memory_manager, xsize, ysize));
   for (size_t c = 0; c < color_channels; ++c) {
     JXL_RETURN_IF_ERROR(ConvertFromExternal(bytes.data(), bytes.size(), xsize,
                                             ysize, bits_per_sample, format, c,
@@ -194,7 +201,8 @@ Status ConvertFromExternal(Span<const uint8_t> bytes, size_t xsize,
   // Passing an interleaved image with an alpha channel to an image that doesn't
   // have alpha channel just discards the passed alpha channel.
   if (has_alpha && ib->HasAlpha()) {
-    JXL_ASSIGN_OR_RETURN(ImageF alpha, ImageF::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(ImageF alpha,
+                         ImageF::Create(memory_manager, xsize, ysize));
     JXL_RETURN_IF_ERROR(ConvertFromExternal(
         bytes.data(), bytes.size(), xsize, ysize, bits_per_sample, format,
         format.num_channels - 1, pool, &alpha));
@@ -202,7 +210,8 @@ Status ConvertFromExternal(Span<const uint8_t> bytes, size_t xsize,
   } else if (!has_alpha && ib->HasAlpha()) {
     // if alpha is not passed, but it is expected, then assume
     // it is all-opaque
-    JXL_ASSIGN_OR_RETURN(ImageF alpha, ImageF::Create(xsize, ysize));
+    JXL_ASSIGN_OR_RETURN(ImageF alpha,
+                         ImageF::Create(memory_manager, xsize, ysize));
     FillImage(1.0f, &alpha);
     ib->SetAlpha(std::move(alpha));
   }

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

@@ -3,9 +3,17 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
+#include <jxl/types.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
 #include "benchmark/benchmark.h"
+#include "lib/jxl/base/status.h"
 #include "lib/jxl/enc_external_image.h"
-#include "lib/jxl/image_ops.h"
+#include "lib/jxl/image_metadata.h"
+#include "tools/no_memory_manager.h"
 
 namespace jxl {
 namespace {
@@ -18,7 +26,7 @@ void BM_EncExternalImage_ConvertImageRGBA(benchmark::State& state) {
 
   ImageMetadata im;
   im.SetAlphaBits(8);
-  ImageBundle ib(&im);
+  ImageBundle ib(jpegxl::tools::NoMemoryManager(), &im);
 
   std::vector<uint8_t> interleaved(xsize * ysize * 4);
   JxlPixelFormat format = {4, JXL_TYPE_UINT8, JXL_NATIVE_ENDIAN, 0};

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

@@ -5,14 +5,16 @@
 
 #include "lib/jxl/enc_external_image.h"
 
-#include <array>
+#include <jxl/types.h>
-#include <new>
 
-#include "lib/jxl/base/compiler_specific.h"
+#include <cstddef>
-#include "lib/jxl/base/data_parallel.h"
+#include <cstdint>
+
+#include "lib/jxl/base/span.h"
 #include "lib/jxl/color_encoding_internal.h"
-#include "lib/jxl/image_ops.h"
+#include "lib/jxl/image_bundle.h"
-#include "lib/jxl/image_test_utils.h"
+#include "lib/jxl/image_metadata.h"
+#include "lib/jxl/test_utils.h"
 #include "lib/jxl/testing.h"
 
 namespace jxl {
@@ -22,7 +24,7 @@ namespace {
 TEST(ExternalImageTest, InvalidSize) {
   ImageMetadata im;
   im.SetAlphaBits(8);
-  ImageBundle ib(&im);
+  ImageBundle ib(jxl::test::MemoryManager(), &im);
 
   JxlPixelFormat format = {4, JXL_TYPE_UINT16, JXL_BIG_ENDIAN, 0};
   const uint8_t buf[10 * 100 * 8] = {};
@@ -44,7 +46,7 @@ TEST(ExternalImageTest, InvalidSize) {
 TEST(ExternalImageTest, AlphaMissing) {
   ImageMetadata im;
   im.SetAlphaBits(0);  // No alpha
-  ImageBundle ib(&im);
+  ImageBundle ib(jxl::test::MemoryManager(), &im);
 
   const size_t xsize = 10;
   const size_t ysize = 20;
@@ -63,7 +65,7 @@ TEST(ExternalImageTest, AlphaPremultiplied) {
   ImageMetadata im;
   im.SetAlphaBits(8, true);
 
-  ImageBundle ib(&im);
+  ImageBundle ib(jxl::test::MemoryManager(), &im);
   const size_t xsize = 10;
   const size_t ysize = 20;
   const size_t size = xsize * ysize * 8;

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

@@ -5,6 +5,8 @@
 
 #include "lib/jxl/enc_frame.h"
 
+#include <jxl/memory_manager.h>
+
 #include <algorithm>
 #include <array>
 #include <atomic>
@@ -36,7 +38,6 @@
 #include "lib/jxl/enc_ac_strategy.h"
 #include "lib/jxl/enc_adaptive_quantization.h"
 #include "lib/jxl/enc_ans.h"
-#include "lib/jxl/enc_ar_control_field.h"
 #include "lib/jxl/enc_aux_out.h"
 #include "lib/jxl/enc_bit_writer.h"
 #include "lib/jxl/enc_cache.h"
@@ -260,10 +261,14 @@ Status LoopFilterFromParams(const CompressParams& cparams, bool streaming_mode,
   loop_filter->gab = ApplyOverride(
       cparams.gaborish, cparams.speed_tier <= SpeedTier::kHare &&
                             frame_header->encoding == FrameEncoding::kVarDCT &&
-                            cparams.decoding_speed_tier < 4);
+                            cparams.decoding_speed_tier < 4 &&
+                            !cparams.disable_percepeptual_optimizations);
 
   if (cparams.epf != -1) {
     loop_filter->epf_iters = cparams.epf;
+  } else if (cparams.disable_percepeptual_optimizations) {
+    loop_filter->epf_iters = 0;
+    return true;
   } else {
     if (frame_header->encoding == FrameEncoding::kModular) {
       loop_filter->epf_iters = 0;
@@ -732,6 +737,7 @@ Status ComputeJPEGTranscodingData(const jpeg::JPEGData& jpeg_data,
                                   ModularFrameEncoder* enc_modular,
                                   PassesEncoderState* enc_state) {
   PassesSharedState& shared = enc_state->shared;
+  JxlMemoryManager* memory_manager = enc_state->memory_manager();
   const FrameDimensions& frame_dim = shared.frame_dim;
 
   const size_t xsize = frame_dim.xsize_padded;
@@ -740,8 +746,8 @@ Status ComputeJPEGTranscodingData(const jpeg::JPEGData& jpeg_data,
   const size_t ysize_blocks = frame_dim.ysize_blocks;
 
   // no-op chroma from luma
-  JXL_ASSIGN_OR_RETURN(shared.cmap,
+  JXL_ASSIGN_OR_RETURN(shared.cmap, ColorCorrelationMap::Create(
-                       ColorCorrelationMap::Create(xsize, ysize, false));
+                                        memory_manager, xsize, ysize, false));
   shared.ac_strategy.FillDCT8();
   FillImage(static_cast<uint8_t>(0), &shared.epf_sharpness);
 
@@ -749,7 +755,8 @@ Status ComputeJPEGTranscodingData(const jpeg::JPEGData& jpeg_data,
   while (enc_state->coeffs.size() < enc_state->passes.size()) {
     JXL_ASSIGN_OR_RETURN(
         std::unique_ptr<ACImageT<int32_t>> coeffs,
-        ACImageT<int32_t>::Make(kGroupDim * kGroupDim, frame_dim.num_groups));
+        ACImageT<int32_t>::Make(memory_manager, kGroupDim * kGroupDim,
+                                frame_dim.num_groups));
     enc_state->coeffs.emplace_back(std::move(coeffs));
   }
 
@@ -775,7 +782,7 @@ Status ComputeJPEGTranscodingData(const jpeg::JPEGData& jpeg_data,
       }
     }
   }
-  DequantMatricesSetCustomDC(&shared.matrices, dcquantization);
+  DequantMatricesSetCustomDC(memory_manager, &shared.matrices, dcquantization);
   float dcquantization_r[3] = {1.0f / dcquantization[0],
                                1.0f / dcquantization[1],
                                1.0f / dcquantization[2]};
@@ -893,7 +900,8 @@ Status ComputeJPEGTranscodingData(const jpeg::JPEGData& jpeg_data,
     }
   }
 
-  JXL_ASSIGN_OR_RETURN(Image3F dc, Image3F::Create(xsize_blocks, ysize_blocks));
+  JXL_ASSIGN_OR_RETURN(
+      Image3F dc, Image3F::Create(memory_manager, xsize_blocks, ysize_blocks));
   if (!frame_header.chroma_subsampling.Is444()) {
     ZeroFillImage(&dc);
     for (auto& coeff : enc_state->coeffs) {
@@ -1062,12 +1070,27 @@ Status ComputeVarDCTEncodingData(const FrameHeader& frame_header,
                                  AuxOut* aux_out) {
   JXL_ASSERT((rect.xsize() % kBlockDim) == 0 &&
              (rect.ysize() % kBlockDim) == 0);
+  JxlMemoryManager* memory_manager = enc_state->memory_manager();
+  // Save pre-Gaborish opsin for AR control field heuristics computation.
+  Image3F orig_opsin;
+  JXL_ASSIGN_OR_RETURN(
+      orig_opsin, Image3F::Create(memory_manager, rect.xsize(), rect.ysize()));
+  CopyImageTo(rect, *opsin, Rect(orig_opsin), &orig_opsin);
+  orig_opsin.ShrinkTo(enc_state->shared.frame_dim.xsize,
+                      enc_state->shared.frame_dim.ysize);
+
   JXL_RETURN_IF_ERROR(LossyFrameHeuristics(frame_header, enc_state, enc_modular,
                                            linear, opsin, rect, cms, pool,
                                            aux_out));
 
   JXL_RETURN_IF_ERROR(InitializePassesEncoder(
       frame_header, *opsin, rect, cms, pool, enc_state, enc_modular, aux_out));
+
+  JXL_RETURN_IF_ERROR(
+      ComputeARHeuristics(frame_header, enc_state, orig_opsin, rect, pool));
+
+  JXL_RETURN_IF_ERROR(ComputeACMetadata(pool, enc_state, enc_modular));
+
   return true;
 }
 
@@ -1090,10 +1113,11 @@ void ComputeAllCoeffOrders(PassesEncoderState& enc_state,
 // Working area for TokenizeCoefficients (per-group!)
 struct EncCache {
   // Allocates memory when first called.
-  Status InitOnce() {
+  Status InitOnce(JxlMemoryManager* memory_manager) {
     if (num_nzeroes.xsize() == 0) {
-      JXL_ASSIGN_OR_RETURN(
+      JXL_ASSIGN_OR_RETURN(num_nzeroes,
-          num_nzeroes, Image3I::Create(kGroupDimInBlocks, kGroupDimInBlocks));
+                           Image3I::Create(memory_manager, kGroupDimInBlocks,
+                                           kGroupDimInBlocks));
     }
     return true;
   }
@@ -1106,6 +1130,7 @@ Status TokenizeAllCoefficients(const FrameHeader& frame_header,
                                PassesEncoderState* enc_state) {
   PassesSharedState& shared = enc_state->shared;
   std::vector<EncCache> group_caches;
+  JxlMemoryManager* memory_manager = enc_state->memory_manager();
   const auto tokenize_group_init = [&](const size_t num_threads) {
     group_caches.resize(num_threads);
     return true;
@@ -1124,7 +1149,7 @@ Status TokenizeAllCoefficients(const FrameHeader& frame_header,
           enc_state->coeffs[idx_pass]->PlaneRow(2, group_index, 0).ptr32,
       };
       // Ensure group cache is initialized.
-      if (!group_caches[thread].InitOnce()) {
+      if (!group_caches[thread].InitOnce(memory_manager)) {
         has_error = true;
         return;
       }
@@ -1160,8 +1185,10 @@ Status EncodeGlobalDCInfo(const PassesSharedState& shared, BitWriter* writer,
 Status EncodeGlobalACInfo(PassesEncoderState* enc_state, BitWriter* writer,
                           ModularFrameEncoder* enc_modular, AuxOut* aux_out) {
   PassesSharedState& shared = enc_state->shared;
-  JXL_RETURN_IF_ERROR(DequantMatricesEncode(shared.matrices, writer,
+  JxlMemoryManager* memory_manager = enc_state->memory_manager();
-                                            kLayerQuant, aux_out, enc_modular));
+  JXL_RETURN_IF_ERROR(DequantMatricesEncode(memory_manager, shared.matrices,
+                                            writer, kLayerQuant, aux_out,
+                                            enc_modular));
   size_t num_histo_bits = CeilLog2Nonzero(shared.frame_dim.num_groups);
   if (!enc_state->streaming_mode && num_histo_bits != 0) {
     BitWriter::Allotment allotment(writer, num_histo_bits);
@@ -1213,7 +1240,7 @@ Status EncodeGlobalACInfo(PassesEncoderState* enc_state, BitWriter* writer,
     hist_params.streaming_mode = enc_state->streaming_mode;
     hist_params.initialize_global_state = enc_state->initialize_global_state;
     BuildAndEncodeHistograms(
-        hist_params,
+        memory_manager, hist_params,
         num_histogram_groups * shared.block_ctx_map.NumACContexts(),
         enc_state->passes[i].ac_tokens, &enc_state->passes[i].codes,
         &enc_state->passes[i].context_map, writer, kLayerAC, aux_out);
@@ -1225,8 +1252,10 @@ Status EncodeGlobalACInfo(PassesEncoderState* enc_state, BitWriter* writer,
 Status EncodeGroups(const FrameHeader& frame_header,
                     PassesEncoderState* enc_state,
                     ModularFrameEncoder* enc_modular, ThreadPool* pool,
-                    std::vector<BitWriter>* group_codes, AuxOut* aux_out) {
+                    std::vector<std::unique_ptr<BitWriter>>* group_codes,
+                    AuxOut* aux_out) {
   const PassesSharedState& shared = enc_state->shared;
+  JxlMemoryManager* memory_manager = shared.memory_manager;
   const FrameDimensions& frame_dim = shared.frame_dim;
   const size_t num_groups = frame_dim.num_groups;
   const size_t num_passes = enc_state->progressive_splitter.GetNumPasses();
@@ -1237,10 +1266,14 @@ Status EncodeGroups(const FrameHeader& frame_header,
       is_small_image ? 1
                      : AcGroupIndex(0, 0, num_groups, frame_dim.num_dc_groups) +
                            num_groups * num_passes;
-  group_codes->resize(num_toc_entries);
+  JXL_ASSERT(group_codes->empty());
+  group_codes->reserve(num_toc_entries);
+  for (size_t i = 0; i < num_toc_entries; ++i) {
+    group_codes->emplace_back(jxl::make_unique<BitWriter>(memory_manager));
+  }
 
-  const auto get_output = [&](const size_t index) {
+  const auto get_output = [&](const size_t index) -> BitWriter* {
-    return &(*group_codes)[is_small_image ? 0 : index];
+    return (*group_codes)[is_small_image ? 0 : index].get();
   };
   auto ac_group_code = [&](size_t pass, size_t group) {
     return get_output(AcGroupIndex(pass, group, frame_dim.num_groups,
@@ -1377,10 +1410,10 @@ Status EncodeGroups(const FrameHeader& frame_header,
   // Resizing aux_outs to 0 also Assimilates the array.
   static_cast<void>(resize_aux_outs(0));
 
-  for (BitWriter& bw : *group_codes) {
+  for (std::unique_ptr<BitWriter>& bw : *group_codes) {
-    BitWriter::Allotment allotment(&bw, 8);
+    BitWriter::Allotment allotment(bw.get(), 8);
-    bw.ZeroPadToByte();  // end of group.
+    bw->ZeroPadToByte();  // end of group.
-    allotment.ReclaimAndCharge(&bw, kLayerAC, aux_out);
+    allotment.ReclaimAndCharge(bw.get(), kLayerAC, aux_out);
   }
   return true;
 }
@@ -1391,10 +1424,11 @@ Status ComputeEncodingData(
     const jpeg::JPEGData* jpeg_data, size_t x0, size_t y0, size_t xsize,
     size_t ysize, const JxlCmsInterface& cms, ThreadPool* pool,
     FrameHeader& mutable_frame_header, ModularFrameEncoder& enc_modular,
-    PassesEncoderState& enc_state, std::vector<BitWriter>* group_codes,
+    PassesEncoderState& enc_state,
-    AuxOut* aux_out) {
+    std::vector<std::unique_ptr<BitWriter>>* group_codes, AuxOut* aux_out) {
   JXL_ASSERT(x0 + xsize <= frame_data.xsize);
   JXL_ASSERT(y0 + ysize <= frame_data.ysize);
+  JxlMemoryManager* memory_manager = enc_state.memory_manager();
   const FrameHeader& frame_header = mutable_frame_header;
   PassesSharedState& shared = enc_state.shared;
   shared.metadata = metadata;
@@ -1412,26 +1446,29 @@ Status ComputeEncodingData(
   const FrameDimensions& frame_dim = shared.frame_dim;
   JXL_ASSIGN_OR_RETURN(
       shared.ac_strategy,
-      AcStrategyImage::Create(frame_dim.xsize_blocks, frame_dim.ysize_blocks));
+      AcStrategyImage::Create(memory_manager, frame_dim.xsize_blocks,
+                              frame_dim.ysize_blocks));
+  JXL_ASSIGN_OR_RETURN(shared.raw_quant_field,
+                       ImageI::Create(memory_manager, frame_dim.xsize_blocks,
+                                      frame_dim.ysize_blocks));
+  JXL_ASSIGN_OR_RETURN(shared.epf_sharpness,
+                       ImageB::Create(memory_manager, frame_dim.xsize_blocks,
+                                      frame_dim.ysize_blocks));
   JXL_ASSIGN_OR_RETURN(
-      shared.raw_quant_field,
+      shared.cmap, ColorCorrelationMap::Create(memory_manager, frame_dim.xsize,
-      ImageI::Create(frame_dim.xsize_blocks, frame_dim.ysize_blocks));
+                                               frame_dim.ysize));
-  JXL_ASSIGN_OR_RETURN(
-      shared.epf_sharpness,
-      ImageB::Create(frame_dim.xsize_blocks, frame_dim.ysize_blocks));
-  JXL_ASSIGN_OR_RETURN(shared.cmap, ColorCorrelationMap::Create(
-                                        frame_dim.xsize, frame_dim.ysize));
   shared.coeff_order_size = kCoeffOrderMaxSize;
   if (frame_header.encoding == FrameEncoding::kVarDCT) {
     shared.coeff_orders.resize(frame_header.passes.num_passes *
                                kCoeffOrderMaxSize);
   }
 
-  JXL_ASSIGN_OR_RETURN(shared.quant_dc, ImageB::Create(frame_dim.xsize_blocks,
+  JXL_ASSIGN_OR_RETURN(shared.quant_dc,
-                                                       frame_dim.ysize_blocks));
+                       ImageB::Create(memory_manager, frame_dim.xsize_blocks,
-  JXL_ASSIGN_OR_RETURN(
+                                      frame_dim.ysize_blocks));
-      shared.dc_storage,
+  JXL_ASSIGN_OR_RETURN(shared.dc_storage,
-      Image3F::Create(frame_dim.xsize_blocks, frame_dim.ysize_blocks));
+                       Image3F::Create(memory_manager, frame_dim.xsize_blocks,
+                                       frame_dim.ysize_blocks));
   shared.dc = &shared.dc_storage;
 
   const size_t num_extra_channels = metadata->m.num_extra_channels;
@@ -1451,14 +1488,16 @@ Status ComputeEncodingData(
   JXL_ASSERT(patch_rect.IsInside(frame_rect));
 
   // Allocating a large enough image avoids a copy when padding.
-  JXL_ASSIGN_OR_RETURN(Image3F color,
+  JXL_ASSIGN_OR_RETURN(
-                       Image3F::Create(RoundUpToBlockDim(patch_rect.xsize()),
+      Image3F color,
-                                       RoundUpToBlockDim(patch_rect.ysize())));
+      Image3F::Create(memory_manager, RoundUpToBlockDim(patch_rect.xsize()),
+                      RoundUpToBlockDim(patch_rect.ysize())));
   color.ShrinkTo(patch_rect.xsize(), patch_rect.ysize());
   std::vector<ImageF> extra_channels(num_extra_channels);
   for (auto& extra_channel : extra_channels) {
     JXL_ASSIGN_OR_RETURN(
-        extra_channel, ImageF::Create(patch_rect.xsize(), patch_rect.ysize()));
+        extra_channel,
+        ImageF::Create(memory_manager, patch_rect.xsize(), patch_rect.ysize()));
   }
   ImageF* alpha = alpha_eci ? &extra_channels[alpha_idx] : nullptr;
   ImageF* black = black_eci ? &extra_channels[black_idx] : nullptr;
@@ -1484,9 +1523,9 @@ Status ComputeEncodingData(
         frame_info.ib_needs_color_transform) {
       if (frame_header.encoding == FrameEncoding::kVarDCT &&
           cparams.speed_tier <= SpeedTier::kKitten) {
-        JXL_ASSIGN_OR_RETURN(
+        JXL_ASSIGN_OR_RETURN(linear_storage,
-            linear_storage,
+                             Image3F::Create(memory_manager, patch_rect.xsize(),
-            Image3F::Create(patch_rect.xsize(), patch_rect.ysize()));
+                                             patch_rect.ysize()));
         linear = &linear_storage;
       }
       ToXYB(c_enc, metadata->m.IntensityTarget(), black, pool, &color, cms,
@@ -1500,8 +1539,9 @@ Status ComputeEncodingData(
     bool lossless = cparams.IsLossless();
     if (alpha && !alpha_eci->alpha_associated &&
         frame_header.frame_type == FrameType::kRegularFrame &&
-        !ApplyOverride(cparams.keep_invisible, true) &&
+        !ApplyOverride(cparams.keep_invisible, cparams.IsLossless()) &&
-        cparams.ec_resampling == cparams.resampling) {
+        cparams.ec_resampling == cparams.resampling &&
+        !cparams.disable_percepeptual_optimizations) {
       // simplify invisible pixels
       SimplifyInvisible(&color, *alpha, lossless);
       if (linear) {
@@ -1596,7 +1636,7 @@ Status ComputeEncodingData(
 Status PermuteGroups(const CompressParams& cparams,
                      const FrameDimensions& frame_dim, size_t num_passes,
                      std::vector<coeff_order_t>* permutation,
-                     std::vector<BitWriter>* group_codes) {
+                     std::vector<std::unique_ptr<BitWriter>>* group_codes) {
   const size_t num_groups = frame_dim.num_groups;
   if (!cparams.centerfirst || (num_passes == 1 && num_groups == 1)) {
     return true;
@@ -1665,11 +1705,11 @@ Status PermuteGroups(const CompressParams& cparams,
       permutation->push_back(pass_start + v);
     }
   }
-  std::vector<BitWriter> new_group_codes(group_codes->size());
+  std::vector<std::unique_ptr<BitWriter>> new_group_codes(group_codes->size());
   for (size_t i = 0; i < permutation->size(); i++) {
     new_group_codes[(*permutation)[i]] = std::move((*group_codes)[i]);
   }
-  *group_codes = std::move(new_group_codes);
+  group_codes->swap(new_group_codes);
   return true;
 }
 
@@ -1790,8 +1830,9 @@ size_t TOCSize(const std::vector<size_t>& group_sizes) {
   return (toc_bits + 7) / 8;
 }
 
-PaddedBytes EncodeTOC(const std::vector<size_t>& group_sizes, AuxOut* aux_out) {
+PaddedBytes EncodeTOC(JxlMemoryManager* memory_manager,
-  BitWriter writer;
+                      const std::vector<size_t>& group_sizes, AuxOut* aux_out) {
+  BitWriter writer{memory_manager};
   BitWriter::Allotment allotment(&writer, 32 * group_sizes.size());
   for (size_t group_size : group_sizes) {
     JXL_CHECK(U32Coder::Write(kTocDist, group_size, &writer));
@@ -1831,17 +1872,17 @@ size_t ComputeDcGlobalPadding(const std::vector<size_t>& group_sizes,
   return group_data_offset - actual_offset;
 }
 
-Status OutputGroups(std::vector<BitWriter>&& group_codes,
+Status OutputGroups(std::vector<std::unique_ptr<BitWriter>>&& group_codes,
                     std::vector<size_t>* group_sizes,
                     JxlEncoderOutputProcessorWrapper* output_processor) {
   JXL_ASSERT(group_codes.size() >= 4);
   {
-    PaddedBytes dc_group = std::move(group_codes[1]).TakeBytes();
+    PaddedBytes dc_group = std::move(*group_codes[1]).TakeBytes();
     group_sizes->push_back(dc_group.size());
     JXL_RETURN_IF_ERROR(AppendData(*output_processor, dc_group));
   }
   for (size_t i = 3; i < group_codes.size(); ++i) {
-    PaddedBytes ac_group = std::move(group_codes[i]).TakeBytes();
+    PaddedBytes ac_group = std::move(*group_codes[i]).TakeBytes();
     group_sizes->push_back(ac_group.size());
     JXL_RETURN_IF_ERROR(AppendData(*output_processor, ac_group));
   }
@@ -1879,7 +1920,8 @@ Status OutputAcGlobal(PassesEncoderState& enc_state,
                       JxlEncoderOutputProcessorWrapper* output_processor,
                       AuxOut* aux_out) {
   JXL_ASSERT(frame_dim.num_groups > 1);
-  BitWriter writer;
+  JxlMemoryManager* memory_manager = enc_state.memory_manager();
+  BitWriter writer{memory_manager};
   {
     size_t num_histo_bits = CeilLog2Nonzero(frame_dim.num_groups);
     BitWriter::Allotment allotment(&writer, num_histo_bits + 1);
@@ -1917,14 +1959,15 @@ Status OutputAcGlobal(PassesEncoderState& enc_state,
   return true;
 }
 
-Status EncodeFrameStreaming(const CompressParams& cparams,
+Status EncodeFrameStreaming(JxlMemoryManager* memory_manager,
+                            const CompressParams& cparams,
                             const FrameInfo& frame_info,
                             const CodecMetadata* metadata,
                             JxlEncoderChunkedFrameAdapter& frame_data,
                             const JxlCmsInterface& cms, ThreadPool* pool,
                             JxlEncoderOutputProcessorWrapper* output_processor,
                             AuxOut* aux_out) {
-  PassesEncoderState enc_state;
+  PassesEncoderState enc_state{memory_manager};
   SetProgressiveMode(cparams, &enc_state.progressive_splitter);
   FrameHeader frame_header(metadata);
   std::unique_ptr<jpeg::JPEGData> jpeg_data;
@@ -1936,7 +1979,7 @@ Status EncodeFrameStreaming(const CompressParams& cparams,
                                       frame_info, jpeg_data.get(), true,
                                       &frame_header));
   const size_t num_passes = enc_state.progressive_splitter.GetNumPasses();
-  ModularFrameEncoder enc_modular(frame_header, cparams, true);
+  ModularFrameEncoder enc_modular(memory_manager, frame_header, cparams, true);
   std::vector<coeff_order_t> permutation;
   std::vector<size_t> dc_group_order;
   size_t group_size = frame_header.ToFrameDimensions().group_dim;
@@ -1947,8 +1990,8 @@ Status EncodeFrameStreaming(const CompressParams& cparams,
   size_t dc_group_xsize = DivCeil(frame_data.xsize, dc_group_size);
   size_t min_dc_global_size = 0;
   size_t group_data_offset = 0;
-  PaddedBytes frame_header_bytes;
+  PaddedBytes frame_header_bytes{memory_manager};
-  PaddedBytes dc_global_bytes;
+  PaddedBytes dc_global_bytes{memory_manager};
   std::vector<size_t> group_sizes;
   size_t start_pos = output_processor->CurrentPosition();
   for (size_t i = 0; i < dc_group_order.size(); ++i) {
@@ -1970,14 +2013,14 @@ Status EncodeFrameStreaming(const CompressParams& cparams,
     enc_state.initialize_global_state = (i == 0);
     enc_state.dc_group_index = dc_ix;
     enc_state.histogram_idx = std::vector<size_t>(group_xsize * group_ysize, i);
-    std::vector<BitWriter> group_codes;
+    std::vector<std::unique_ptr<BitWriter>> group_codes;
     JXL_RETURN_IF_ERROR(ComputeEncodingData(
         cparams, frame_info, metadata, frame_data, jpeg_data.get(), x0, y0,
         xsize, ysize, cms, pool, frame_header, enc_modular, enc_state,
         &group_codes, aux_out));
     JXL_ASSERT(enc_state.special_frames.empty());
     if (i == 0) {
-      BitWriter writer;
+      BitWriter writer{memory_manager};
       JXL_RETURN_IF_ERROR(WriteFrameHeader(frame_header, &writer, aux_out));
       BitWriter::Allotment allotment(&writer, 8);
       writer.Write(1, 1);  // write permutation
@@ -1986,7 +2029,7 @@ Status EncodeFrameStreaming(const CompressParams& cparams,
       writer.ZeroPadToByte();
       allotment.ReclaimAndCharge(&writer, kLayerHeader, aux_out);
       frame_header_bytes = std::move(writer).TakeBytes();
-      dc_global_bytes = std::move(group_codes[0]).TakeBytes();
+      dc_global_bytes = std::move(*group_codes[0]).TakeBytes();
       ComputeGroupDataOffset(frame_header_bytes.size(), dc_global_bytes.size(),
                              permutation.size(), min_dc_global_size,
                              group_data_offset);
@@ -2015,7 +2058,7 @@ Status EncodeFrameStreaming(const CompressParams& cparams,
       ComputeDcGlobalPadding(group_sizes, frame_header_bytes.size(),
                              group_data_offset, min_dc_global_size);
   group_sizes[0] += padding_size;
-  PaddedBytes toc_bytes = EncodeTOC(group_sizes, aux_out);
+  PaddedBytes toc_bytes = EncodeTOC(memory_manager, group_sizes, aux_out);
   std::vector<uint8_t> padding_bytes(padding_size);
   JXL_RETURN_IF_ERROR(AppendData(*output_processor, frame_header_bytes));
   JXL_RETURN_IF_ERROR(AppendData(*output_processor, toc_bytes));
@@ -2029,16 +2072,16 @@ Status EncodeFrameStreaming(const CompressParams& cparams,
   return true;
 }
 
-Status EncodeFrameOneShot(const CompressParams& cparams,
+Status EncodeFrameOneShot(JxlMemoryManager* memory_manager,
+                          const CompressParams& cparams,
                           const FrameInfo& frame_info,
                           const CodecMetadata* metadata,
                           JxlEncoderChunkedFrameAdapter& frame_data,
                           const JxlCmsInterface& cms, ThreadPool* pool,
                           JxlEncoderOutputProcessorWrapper* output_processor,
                           AuxOut* aux_out) {
-  PassesEncoderState enc_state;
+  PassesEncoderState enc_state{memory_manager};
   SetProgressiveMode(cparams, &enc_state.progressive_splitter);
-  std::vector<BitWriter> group_codes;
   FrameHeader frame_header(metadata);
   std::unique_ptr<jpeg::JPEGData> jpeg_data;
   if (frame_data.IsJPEG()) {
@@ -2049,13 +2092,14 @@ Status EncodeFrameOneShot(const CompressParams& cparams,
                                       frame_info, jpeg_data.get(), false,
                                       &frame_header));
   const size_t num_passes = enc_state.progressive_splitter.GetNumPasses();
-  ModularFrameEncoder enc_modular(frame_header, cparams, false);
+  ModularFrameEncoder enc_modular(memory_manager, frame_header, cparams, false);
+  std::vector<std::unique_ptr<BitWriter>> group_codes;
   JXL_RETURN_IF_ERROR(ComputeEncodingData(
       cparams, frame_info, metadata, frame_data, jpeg_data.get(), 0, 0,
       frame_data.xsize, frame_data.ysize, cms, pool, frame_header, enc_modular,
       enc_state, &group_codes, aux_out));
 
-  BitWriter writer;
+  BitWriter writer{memory_manager};
   writer.AppendByteAligned(enc_state.special_frames);
   JXL_RETURN_IF_ERROR(WriteFrameHeader(frame_header, &writer, aux_out));
 
@@ -2075,7 +2119,8 @@ Status EncodeFrameOneShot(const CompressParams& cparams,
 
 }  // namespace
 
-Status EncodeFrame(const CompressParams& cparams_orig,
+Status EncodeFrame(JxlMemoryManager* memory_manager,
+                   const CompressParams& cparams_orig,
                    const FrameInfo& frame_info, const CodecMetadata* metadata,
                    JxlEncoderChunkedFrameAdapter& frame_data,
                    const JxlCmsInterface& cms, ThreadPool* pool,
@@ -2146,8 +2191,9 @@ Status EncodeFrame(const CompressParams& cparams_orig,
           size_t avail_out = output.size();
           JxlEncoderOutputProcessorWrapper local_output;
           local_output.SetAvailOut(&next_out, &avail_out);
-          if (!EncodeFrame(all_params[task], frame_info, metadata, frame_data,
+          if (!EncodeFrame(memory_manager, all_params[task], frame_info,
-                           cms, nullptr, &local_output, aux_out)) {
+                           metadata, frame_data, cms, nullptr, &local_output,
+                           aux_out)) {
             has_error = true;
             return;
           }
@@ -2215,15 +2261,17 @@ Status EncodeFrame(const CompressParams& cparams_orig,
   }
 
   if (CanDoStreamingEncoding(cparams, frame_info, *metadata, frame_data)) {
-    return EncodeFrameStreaming(cparams, frame_info, metadata, frame_data, cms,
+    return EncodeFrameStreaming(memory_manager, cparams, frame_info, metadata,
-                                pool, output_processor, aux_out);
+                                frame_data, cms, pool, output_processor,
+                                aux_out);
   } else {
-    return EncodeFrameOneShot(cparams, frame_info, metadata, frame_data, cms,
+    return EncodeFrameOneShot(memory_manager, cparams, frame_info, metadata,
-                              pool, output_processor, aux_out);
+                              frame_data, cms, pool, output_processor, aux_out);
   }
 }
 
-Status EncodeFrame(const CompressParams& cparams_orig,
+Status EncodeFrame(JxlMemoryManager* memory_manager,
+                   const CompressParams& cparams_orig,
                    const FrameInfo& frame_info, const CodecMetadata* metadata,
                    const ImageBundle& ib, const JxlCmsInterface& cms,
                    ThreadPool* pool, BitWriter* writer, AuxOut* aux_out) {
@@ -2268,8 +2316,9 @@ Status EncodeFrame(const CompressParams& cparams_orig,
   size_t avail_out = output.size();
   JxlEncoderOutputProcessorWrapper output_processor;
   output_processor.SetAvailOut(&next_out, &avail_out);
-  JXL_RETURN_IF_ERROR(EncodeFrame(cparams_orig, fi, metadata, frame_data, cms,
+  JXL_RETURN_IF_ERROR(EncodeFrame(memory_manager, cparams_orig, fi, metadata,
-                                  pool, &output_processor, aux_out));
+                                  frame_data, cms, pool, &output_processor,
+                                  aux_out));
   output_processor.SetFinalizedPosition();
   output_processor.CopyOutput(output, next_out, avail_out);
   writer->AppendByteAligned(Bytes(output));

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

@@ -7,6 +7,7 @@
 #define LIB_JXL_ENC_FRAME_H_
 
 #include <jxl/cms_interface.h>
+#include <jxl/memory_manager.h>
 #include <jxl/types.h>
 
 #include <cstddef>
@@ -91,14 +92,16 @@ Status ParamsPostInit(CompressParams* p);
 // be processed in parallel by `pool`. metadata is the ImageMetadata encoded in
 // the codestream, and must be used for the FrameHeaders, do not use
 // ib.metadata.
-Status EncodeFrame(const CompressParams& cparams_orig,
+Status EncodeFrame(JxlMemoryManager* memory_manager,
+                   const CompressParams& cparams_orig,
                    const FrameInfo& frame_info, const CodecMetadata* metadata,
                    JxlEncoderChunkedFrameAdapter& frame_data,
                    const JxlCmsInterface& cms, ThreadPool* pool,
                    JxlEncoderOutputProcessorWrapper* output_processor,
                    AuxOut* aux_out);
 
-Status EncodeFrame(const CompressParams& cparams_orig,
+Status EncodeFrame(JxlMemoryManager* memory_manager,
+                   const CompressParams& cparams_orig,
                    const FrameInfo& frame_info, const CodecMetadata* metadata,
                    const ImageBundle& ib, const JxlCmsInterface& cms,
                    ThreadPool* pool, BitWriter* writer, AuxOut* aux_out);

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

@@ -5,6 +5,8 @@
 
 #include "lib/jxl/enc_gaborish.h"
 
+#include <jxl/memory_manager.h>
+
 #include <hwy/base.h>
 
 #include "lib/jxl/base/data_parallel.h"
@@ -18,6 +20,7 @@ namespace jxl {
 
 Status GaborishInverse(Image3F* in_out, const Rect& rect, const float mul[3],
                        ThreadPool* pool) {
+  JxlMemoryManager* memory_manager = in_out->memory_manager();
   WeightsSymmetric5 weights[3];
   // Only an approximation. One or even two 3x3, and rank-1 (separable) 5x5
   // are insufficient. The numbers here have been obtained by butteraugli
@@ -49,8 +52,9 @@ Status GaborishInverse(Image3F* in_out, const Rect& rect, const float mul[3],
   // have planes of different stride. Instead, we copy one plane in a temporary
   // image and reuse the existing planes of the in/out image.
   ImageF temp;
-  JXL_ASSIGN_OR_RETURN(
+  JXL_ASSIGN_OR_RETURN(temp,
-      temp, ImageF::Create(in_out->Plane(2).xsize(), in_out->Plane(2).ysize()));
+                       ImageF::Create(memory_manager, in_out->Plane(2).xsize(),
+                                      in_out->Plane(2).ysize()));
   CopyImageTo(in_out->Plane(2), &temp);
   Rect xrect = rect.Extend(3, Rect(*in_out));
   Symmetric5(in_out->Plane(0), xrect, weights[0], pool, &in_out->Plane(2),

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

@@ -17,6 +17,7 @@
 #include "lib/jxl/image.h"
 #include "lib/jxl/image_ops.h"
 #include "lib/jxl/image_test_utils.h"
+#include "lib/jxl/test_utils.h"
 #include "lib/jxl/testing.h"
 
 namespace jxl {
@@ -43,7 +44,8 @@ void ConvolveGaborish(const ImageF& in, float weight1, float weight2,
 }
 
 void TestRoundTrip(const Image3F& in, float max_l1) {
-  JXL_ASSIGN_OR_DIE(Image3F fwd, Image3F::Create(in.xsize(), in.ysize()));
+  JXL_ASSIGN_OR_DIE(Image3F fwd, Image3F::Create(jxl::test::MemoryManager(),
+                                                 in.xsize(), in.ysize()));
   ThreadPool* null_pool = nullptr;
   ConvolveGaborish(in.Plane(0), 0, 0, null_pool, &fwd.Plane(0));
   ConvolveGaborish(in.Plane(1), 0, 0, null_pool, &fwd.Plane(1));
@@ -59,7 +61,8 @@ void TestRoundTrip(const Image3F& in, float max_l1) {
 }
 
 TEST(GaborishTest, TestZero) {
-  JXL_ASSIGN_OR_DIE(Image3F in, Image3F::Create(20, 20));
+  JXL_ASSIGN_OR_DIE(Image3F in,
+                    Image3F::Create(jxl::test::MemoryManager(), 20, 20));
   ZeroFillImage(&in);
   TestRoundTrip(in, 0.0f);
 }
@@ -67,7 +70,8 @@ TEST(GaborishTest, TestZero) {
 // Disabled: large difference.
 #if JXL_FALSE
 TEST(GaborishTest, TestDirac) {
-  JXL_ASSIGN_OR_DIE(Image3F in, Image3F::Create(20, 20));
+  JXL_ASSIGN_OR_DIE(Image3F in,
+                    Image3F::Create(jxl::test::MemoryManager(), 20, 20));
   ZeroFillImage(&in);
   in.PlaneRow(1, 10)[10] = 10.0f;
   TestRoundTrip(in, 0.26f);
@@ -75,7 +79,8 @@ TEST(GaborishTest, TestDirac) {
 #endif
 
 TEST(GaborishTest, TestFlat) {
-  JXL_ASSIGN_OR_DIE(Image3F in, Image3F::Create(20, 20));
+  JXL_ASSIGN_OR_DIE(Image3F in,
+                    Image3F::Create(jxl::test::MemoryManager(), 20, 20));
   FillImage(1.0f, &in);
   TestRoundTrip(in, 1E-5f);
 }

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

@@ -6,6 +6,7 @@
 #include "lib/jxl/enc_heuristics.h"
 
 #include <jxl/cms_interface.h>
+#include <jxl/memory_manager.h>
 
 #include <algorithm>
 #include <cstddef>
@@ -30,10 +31,11 @@
 #include "lib/jxl/chroma_from_luma.h"
 #include "lib/jxl/coeff_order.h"
 #include "lib/jxl/coeff_order_fwd.h"
+#include "lib/jxl/common.h"
+#include "lib/jxl/dec_group.h"
 #include "lib/jxl/dec_xyb.h"
 #include "lib/jxl/enc_ac_strategy.h"
 #include "lib/jxl/enc_adaptive_quantization.h"
-#include "lib/jxl/enc_ar_control_field.h"
 #include "lib/jxl/enc_cache.h"
 #include "lib/jxl/enc_chroma_from_luma.h"
 #include "lib/jxl/enc_gaborish.h"
@@ -43,6 +45,7 @@
 #include "lib/jxl/enc_patch_dictionary.h"
 #include "lib/jxl/enc_quant_weights.h"
 #include "lib/jxl/enc_splines.h"
+#include "lib/jxl/epf.h"
 #include "lib/jxl/frame_dimensions.h"
 #include "lib/jxl/frame_header.h"
 #include "lib/jxl/image.h"
@@ -193,26 +196,27 @@ void FindBestBlockEntropyModel(const CompressParams& cparams, const ImageI& rqf,
 
 namespace {
 
-Status FindBestDequantMatrices(const CompressParams& cparams,
+Status FindBestDequantMatrices(JxlMemoryManager* memory_manager,
+                               const CompressParams& cparams,
                                ModularFrameEncoder* modular_frame_encoder,
                                DequantMatrices* dequant_matrices) {
   // TODO(veluca): quant matrices for no-gaborish.
   // TODO(veluca): heuristics for in-bitstream quant tables.
   *dequant_matrices = DequantMatrices();
-  if (cparams.max_error_mode) {
+  if (cparams.max_error_mode || cparams.disable_percepeptual_optimizations) {
+    constexpr float kMSEWeights[3] = {0.001, 0.001, 0.001};
+    const float* wp = cparams.disable_percepeptual_optimizations
+                          ? kMSEWeights
+                          : cparams.max_error;
     // Set numerators of all quantization matrices to constant values.
-    float weights[3][1] = {{1.0f / cparams.max_error[0]},
+    float weights[3][1] = {{1.0f / wp[0]}, {1.0f / wp[1]}, {1.0f / wp[2]}};
-                           {1.0f / cparams.max_error[1]},
-                           {1.0f / cparams.max_error[2]}};
     DctQuantWeightParams dct_params(weights);
     std::vector<QuantEncoding> encodings(DequantMatrices::kNum,
                                          QuantEncoding::DCT(dct_params));
     JXL_RETURN_IF_ERROR(DequantMatricesSetCustom(dequant_matrices, encodings,
                                                  modular_frame_encoder));
-    float dc_weights[3] = {1.0f / cparams.max_error[0],
+    float dc_weights[3] = {1.0f / wp[0], 1.0f / wp[1], 1.0f / wp[2]};
-                           1.0f / cparams.max_error[1],
+    DequantMatricesSetCustomDC(memory_manager, dequant_matrices, dc_weights);
-                           1.0f / cparams.max_error[2]};
-    DequantMatricesSetCustomDC(dequant_matrices, dc_weights);
   }
   return true;
 }
@@ -265,6 +269,7 @@ void CreateMask(const ImageF& image, ImageF& mask) {
 Status DownsampleImage2_Sharper(const ImageF& input, ImageF* output) {
   const int64_t kernelx = 12;
   const int64_t kernely = 12;
+  JxlMemoryManager* memory_manager = input.memory_manager();
 
   static const float kernel[144] = {
       -0.000314256996835, -0.000314256996835, -0.000897597057705,
@@ -319,12 +324,14 @@ Status DownsampleImage2_Sharper(const ImageF& input, ImageF* output) {
   int64_t xsize = input.xsize();
   int64_t ysize = input.ysize();
 
-  JXL_ASSIGN_OR_RETURN(ImageF box_downsample, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF box_downsample,
+                       ImageF::Create(memory_manager, xsize, ysize));
   CopyImageTo(input, &box_downsample);
   JXL_ASSIGN_OR_RETURN(box_downsample, DownsampleImage(box_downsample, 2));
 
-  JXL_ASSIGN_OR_RETURN(ImageF mask, ImageF::Create(box_downsample.xsize(),
+  JXL_ASSIGN_OR_RETURN(ImageF mask,
-                                                   box_downsample.ysize()));
+                       ImageF::Create(memory_manager, box_downsample.xsize(),
+                                      box_downsample.ysize()));
   CreateMask(box_downsample, mask);
 
   for (size_t y = 0; y < output->ysize(); y++) {
@@ -391,9 +398,11 @@ Status DownsampleImage2_Sharper(const ImageF& input, ImageF* output) {
 
 Status DownsampleImage2_Sharper(Image3F* opsin) {
   // Allocate extra space to avoid a reallocation when padding.
-  JXL_ASSIGN_OR_RETURN(Image3F downsampled,
+  JxlMemoryManager* memory_manager = opsin->memory_manager();
-                       Image3F::Create(DivCeil(opsin->xsize(), 2) + kBlockDim,
+  JXL_ASSIGN_OR_RETURN(
-                                       DivCeil(opsin->ysize(), 2) + kBlockDim));
+      Image3F downsampled,
+      Image3F::Create(memory_manager, DivCeil(opsin->xsize(), 2) + kBlockDim,
+                      DivCeil(opsin->ysize(), 2) + kBlockDim));
   downsampled.ShrinkTo(downsampled.xsize() - kBlockDim,
                        downsampled.ysize() - kBlockDim);
 
@@ -637,30 +646,37 @@ Status DownsampleImage2_Iterative(const ImageF& orig, ImageF* output) {
   int64_t ysize = orig.ysize();
   int64_t xsize2 = DivCeil(orig.xsize(), 2);
   int64_t ysize2 = DivCeil(orig.ysize(), 2);
+  JxlMemoryManager* memory_manager = orig.memory_manager();
 
-  JXL_ASSIGN_OR_RETURN(ImageF box_downsample, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF box_downsample,
+                       ImageF::Create(memory_manager, xsize, ysize));
   CopyImageTo(orig, &box_downsample);
   JXL_ASSIGN_OR_RETURN(box_downsample, DownsampleImage(box_downsample, 2));
-  JXL_ASSIGN_OR_RETURN(ImageF mask, ImageF::Create(box_downsample.xsize(),
+  JXL_ASSIGN_OR_RETURN(ImageF mask,
-                                                   box_downsample.ysize()));
+                       ImageF::Create(memory_manager, box_downsample.xsize(),
+                                      box_downsample.ysize()));
   CreateMask(box_downsample, mask);
 
   output->ShrinkTo(xsize2, ysize2);
 
   // Initial result image using the sharper downsampling.
   // Allocate extra space to avoid a reallocation when padding.
-  JXL_ASSIGN_OR_RETURN(ImageF initial,
+  JXL_ASSIGN_OR_RETURN(
-                       ImageF::Create(DivCeil(orig.xsize(), 2) + kBlockDim,
+      ImageF initial,
-                                      DivCeil(orig.ysize(), 2) + kBlockDim));
+      ImageF::Create(memory_manager, DivCeil(orig.xsize(), 2) + kBlockDim,
+                     DivCeil(orig.ysize(), 2) + kBlockDim));
   initial.ShrinkTo(initial.xsize() - kBlockDim, initial.ysize() - kBlockDim);
   JXL_RETURN_IF_ERROR(DownsampleImage2_Sharper(orig, &initial));
 
-  JXL_ASSIGN_OR_RETURN(ImageF down,
+  JXL_ASSIGN_OR_RETURN(
-                       ImageF::Create(initial.xsize(), initial.ysize()));
+      ImageF down,
+      ImageF::Create(memory_manager, initial.xsize(), initial.ysize()));
   CopyImageTo(initial, &down);
-  JXL_ASSIGN_OR_RETURN(ImageF up, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF up, ImageF::Create(memory_manager, xsize, ysize));
-  JXL_ASSIGN_OR_RETURN(ImageF corr, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF corr,
-  JXL_ASSIGN_OR_RETURN(ImageF corr2, ImageF::Create(xsize2, ysize2));
+                       ImageF::Create(memory_manager, xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF corr2,
+                       ImageF::Create(memory_manager, xsize2, ysize2));
 
   // In the weights map, relatively higher values will allow less ringing but
   // also less sharpness. With all constant values, it optimizes equally
@@ -669,14 +685,16 @@ Status DownsampleImage2_Iterative(const ImageF& orig, ImageF* output) {
   // TODO(lode): Make use of the weights field for anti-ringing and clamping,
   // the values are all set to 1 for now, but it is intended to be used for
   // reducing ringing based on the mask, and taking clamping into account.
-  JXL_ASSIGN_OR_RETURN(ImageF weights, ImageF::Create(xsize, ysize));
+  JXL_ASSIGN_OR_RETURN(ImageF weights,
+                       ImageF::Create(memory_manager, xsize, ysize));
   for (size_t y = 0; y < weights.ysize(); y++) {
     auto* row = weights.Row(y);
     for (size_t x = 0; x < weights.xsize(); x++) {
       row[x] = 1;
     }
   }
-  JXL_ASSIGN_OR_RETURN(ImageF weights2, ImageF::Create(xsize2, ysize2));
+  JXL_ASSIGN_OR_RETURN(ImageF weights2,
+                       ImageF::Create(memory_manager, xsize2, ysize2));
   AntiUpsample(weights, &weights2);
 
   const size_t num_it = 3;
@@ -706,27 +724,32 @@ Status DownsampleImage2_Iterative(const ImageF& orig, ImageF* output) {
 }  // namespace
 
 Status DownsampleImage2_Iterative(Image3F* opsin) {
+  JxlMemoryManager* memory_manager = opsin->memory_manager();
   // Allocate extra space to avoid a reallocation when padding.
-  JXL_ASSIGN_OR_RETURN(Image3F downsampled,
+  JXL_ASSIGN_OR_RETURN(
-                       Image3F::Create(DivCeil(opsin->xsize(), 2) + kBlockDim,
+      Image3F downsampled,
-                                       DivCeil(opsin->ysize(), 2) + kBlockDim));
+      Image3F::Create(memory_manager, DivCeil(opsin->xsize(), 2) + kBlockDim,
+                      DivCeil(opsin->ysize(), 2) + kBlockDim));
   downsampled.ShrinkTo(downsampled.xsize() - kBlockDim,
                        downsampled.ysize() - kBlockDim);
 
-  JXL_ASSIGN_OR_RETURN(Image3F rgb,
+  JXL_ASSIGN_OR_RETURN(
-                       Image3F::Create(opsin->xsize(), opsin->ysize()));
+      Image3F rgb,
+      Image3F::Create(memory_manager, opsin->xsize(), opsin->ysize()));
   OpsinParams opsin_params;  // TODO(user): use the ones that are actually used
   opsin_params.Init(kDefaultIntensityTarget);
   OpsinToLinear(*opsin, Rect(rgb), nullptr, &rgb, opsin_params);
 
-  JXL_ASSIGN_OR_RETURN(ImageF mask,
+  JXL_ASSIGN_OR_RETURN(
-                       ImageF::Create(opsin->xsize(), opsin->ysize()));
+      ImageF mask,
+      ImageF::Create(memory_manager, opsin->xsize(), opsin->ysize()));
   ButteraugliParams butter_params;
   JXL_ASSIGN_OR_RETURN(std::unique_ptr<ButteraugliComparator> butter,
                        ButteraugliComparator::Make(rgb, butter_params));
   JXL_RETURN_IF_ERROR(butter->Mask(&mask));
-  JXL_ASSIGN_OR_RETURN(ImageF mask_fuzzy,
+  JXL_ASSIGN_OR_RETURN(
-                       ImageF::Create(opsin->xsize(), opsin->ysize()));
+      ImageF mask_fuzzy,
+      ImageF::Create(memory_manager, opsin->xsize(), opsin->ysize()));
 
   for (size_t c = 0; c < 3; c++) {
     JXL_RETURN_IF_ERROR(
@@ -736,10 +759,220 @@ Status DownsampleImage2_Iterative(Image3F* opsin) {
   return true;
 }
 
+StatusOr<Image3F> ReconstructImage(
+    const FrameHeader& orig_frame_header, const PassesSharedState& shared,
+    const std::vector<std::unique_ptr<ACImage>>& coeffs, ThreadPool* pool) {
+  const FrameDimensions& frame_dim = shared.frame_dim;
+  JxlMemoryManager* memory_manager = shared.memory_manager;
+
+  FrameHeader frame_header = orig_frame_header;
+  frame_header.UpdateFlag(shared.image_features.patches.HasAny(),
+                          FrameHeader::kPatches);
+  frame_header.UpdateFlag(shared.image_features.splines.HasAny(),
+                          FrameHeader::kSplines);
+  frame_header.color_transform = ColorTransform::kNone;
+
+  const bool is_gray = shared.metadata->m.color_encoding.IsGray();
+  PassesDecoderState dec_state(memory_manager);
+  JXL_RETURN_IF_ERROR(
+      dec_state.output_encoding_info.SetFromMetadata(*shared.metadata));
+  JXL_RETURN_IF_ERROR(dec_state.output_encoding_info.MaybeSetColorEncoding(
+      ColorEncoding::LinearSRGB(is_gray)));
+  dec_state.shared = &shared;
+  JXL_CHECK(dec_state.Init(frame_header));
+
+  ImageBundle decoded(memory_manager, &shared.metadata->m);
+  decoded.origin = frame_header.frame_origin;
+  JXL_ASSIGN_OR_RETURN(
+      Image3F tmp,
+      Image3F::Create(memory_manager, frame_dim.xsize, frame_dim.ysize));
+  decoded.SetFromImage(std::move(tmp),
+                       dec_state.output_encoding_info.color_encoding);
+
+  PassesDecoderState::PipelineOptions options;
+  options.use_slow_render_pipeline = false;
+  options.coalescing = false;
+  options.render_spotcolors = false;
+  options.render_noise = true;
+
+  JXL_CHECK(dec_state.PreparePipeline(frame_header, &shared.metadata->m,
+                                      &decoded, options));
+
+  hwy::AlignedUniquePtr<GroupDecCache[]> group_dec_caches;
+  const auto allocate_storage = [&](const size_t num_threads) -> Status {
+    JXL_RETURN_IF_ERROR(
+        dec_state.render_pipeline->PrepareForThreads(num_threads,
+                                                     /*use_group_ids=*/false));
+    group_dec_caches = hwy::MakeUniqueAlignedArray<GroupDecCache>(num_threads);
+    return true;
+  };
+  std::atomic<bool> has_error{false};
+  const auto process_group = [&](const uint32_t group_index,
+                                 const size_t thread) {
+    if (has_error) return;
+    if (frame_header.loop_filter.epf_iters > 0) {
+      ComputeSigma(frame_header.loop_filter,
+                   frame_dim.BlockGroupRect(group_index), &dec_state);
+    }
+    RenderPipelineInput input =
+        dec_state.render_pipeline->GetInputBuffers(group_index, thread);
+    JXL_CHECK(DecodeGroupForRoundtrip(frame_header, coeffs, group_index,
+                                      &dec_state, &group_dec_caches[thread],
+                                      thread, input, nullptr, nullptr));
+    if (!input.Done()) {
+      has_error = true;
+      return;
+    }
+  };
+  JXL_CHECK(RunOnPool(pool, 0, frame_dim.num_groups, allocate_storage,
+                      process_group, "ReconstructImage"));
+  if (has_error) return JXL_FAILURE("ReconstructImage failure");
+  return std::move(*decoded.color());
+}
+
+float ComputeBlockL2Distance(const Image3F& a, const Image3F& b,
+                             const ImageF& mask1x1, size_t by, size_t bx) {
+  Rect rect(bx * kBlockDim, by * kBlockDim, kBlockDim, kBlockDim, a.xsize(),
+            a.ysize());
+  float err2 = 0.0f;
+  static const float kXYBWeights[] = {36.0f, 1.0f, 0.2f};
+  for (size_t y = 0; y < rect.ysize(); ++y) {
+    const float* row_a_x = rect.ConstPlaneRow(a, 0, y);
+    const float* row_a_y = rect.ConstPlaneRow(a, 1, y);
+    const float* row_a_b = rect.ConstPlaneRow(a, 2, y);
+    const float* row_b_x = rect.ConstPlaneRow(b, 0, y);
+    const float* row_b_y = rect.ConstPlaneRow(b, 1, y);
+    const float* row_b_b = rect.ConstPlaneRow(b, 2, y);
+    const float* row_mask = rect.ConstRow(mask1x1, y);
+
+    for (size_t x = 0; x < rect.xsize(); ++x) {
+      float mask = row_mask[x];
+      for (size_t c = 0; c < 3; ++c) {
+        float diff_x = row_a_x[x] - row_b_x[x];
+        float diff_y = row_a_y[x] - row_b_y[x];
+        float diff_b = row_a_b[x] - row_b_b[x];
+        err2 += (kXYBWeights[0] * diff_x * diff_x +
+                 kXYBWeights[1] * diff_y * diff_y +
+                 kXYBWeights[2] * diff_b * diff_b) *
+                mask * mask;
+      }
+    }
+  }
+  return err2;
+}
+
+Status ComputeARHeuristics(const FrameHeader& frame_header,
+                           PassesEncoderState* enc_state,
+                           const Image3F& orig_opsin, const Rect& rect,
+                           ThreadPool* pool) {
+  const CompressParams& cparams = enc_state->cparams;
+  PassesSharedState& shared = enc_state->shared;
+  const FrameDimensions& frame_dim = shared.frame_dim;
+  const ImageF& initial_quant_masking1x1 = enc_state->initial_quant_masking1x1;
+  ImageB& epf_sharpness = shared.epf_sharpness;
+  JxlMemoryManager* memory_manager = enc_state->memory_manager();
+
+  if (cparams.butteraugli_distance < kMinButteraugliForDynamicAR ||
+      cparams.speed_tier > SpeedTier::kWombat ||
+      frame_header.loop_filter.epf_iters == 0) {
+    FillPlane(static_cast<uint8_t>(4), &epf_sharpness, Rect(epf_sharpness));
+    return true;
+  }
+
+  std::vector<uint8_t> epf_steps;
+  if (cparams.butteraugli_distance > 4.5f) {
+    epf_steps.push_back(0);
+    epf_steps.push_back(4);
+  } else {
+    epf_steps.push_back(0);
+    epf_steps.push_back(3);
+    epf_steps.push_back(7);
+  }
+  static const int kNumEPFVals = 8;
+  std::array<ImageF, kNumEPFVals> error_images;
+  for (uint8_t val : epf_steps) {
+    FillPlane(val, &epf_sharpness, Rect(epf_sharpness));
+    JXL_ASSIGN_OR_RETURN(
+        Image3F decoded,
+        ReconstructImage(frame_header, shared, enc_state->coeffs, pool));
+    JXL_ASSIGN_OR_RETURN(error_images[val],
+                         ImageF::Create(memory_manager, frame_dim.xsize_blocks,
+                                        frame_dim.ysize_blocks));
+    for (size_t by = 0; by < frame_dim.ysize_blocks; by++) {
+      float* error_row = error_images[val].Row(by);
+      for (size_t bx = 0; bx < frame_dim.xsize_blocks; bx++) {
+        error_row[bx] = ComputeBlockL2Distance(
+            orig_opsin, decoded, initial_quant_masking1x1, by, bx);
+      }
+    }
+  }
+  std::vector<std::vector<size_t>> histo(4, std::vector<size_t>(kNumEPFVals));
+  std::vector<size_t> totals(4, 1);
+  for (size_t by = 0; by < frame_dim.ysize_blocks; by++) {
+    uint8_t* JXL_RESTRICT out_row = epf_sharpness.Row(by);
+    uint8_t* JXL_RESTRICT prev_row = epf_sharpness.Row(by > 0 ? by - 1 : 0);
+    for (size_t bx = 0; bx < frame_dim.xsize_blocks; bx++) {
+      uint8_t best_val = 0;
+      float best_error = std::numeric_limits<float>::max();
+      uint8_t top_val = by > 0 ? prev_row[bx] : 0;
+      uint8_t left_val = bx > 0 ? out_row[bx - 1] : 0;
+      float top_error = error_images[top_val].Row(by)[bx];
+      float left_error = error_images[left_val].Row(by)[bx];
+      for (uint8_t val : epf_steps) {
+        float error = error_images[val].Row(by)[bx];
+        if (val == 0) {
+          error *= 0.97f;
+        }
+        if (error < best_error) {
+          best_val = val;
+          best_error = error;
+        }
+      }
+      if (best_error < 0.995 * std::min(top_error, left_error)) {
+        out_row[bx] = best_val;
+      } else if (top_error < left_error) {
+        out_row[bx] = top_val;
+      } else {
+        out_row[bx] = left_val;
+      }
+      int context = ((top_val > 3) ? 2 : 0) + ((left_val > 3) ? 1 : 0);
+      ++histo[context][out_row[bx]];
+      ++totals[context];
+    }
+  }
+  const float context_weight =
+      0.14f + 0.007f * std::min(10.0f, cparams.butteraugli_distance);
+  for (size_t by = 0; by < frame_dim.ysize_blocks; by++) {
+    uint8_t* JXL_RESTRICT out_row = epf_sharpness.Row(by);
+    uint8_t* JXL_RESTRICT prev_row = epf_sharpness.Row(by > 0 ? by - 1 : 0);
+    for (size_t bx = 0; bx < frame_dim.xsize_blocks; bx++) {
+      uint8_t best_val = 0;
+      float best_error = std::numeric_limits<float>::max();
+      uint8_t top_val = by > 0 ? prev_row[bx] : 0;
+      uint8_t left_val = bx > 0 ? out_row[bx - 1] : 0;
+      int context = ((top_val > 3) ? 2 : 0) + ((left_val > 3) ? 1 : 0);
+      const auto& ctx_histo = histo[context];
+      for (uint8_t val : epf_steps) {
+        float error =
+            error_images[val].Row(by)[bx] /
+            (1 + std::log1p(ctx_histo[val] * context_weight / totals[context]));
+        if (val == 0) error *= 0.93f;
+        if (error < best_error) {
+          best_val = val;
+          best_error = error;
+        }
+      }
+      out_row[bx] = best_val;
+    }
+  }
+
+  return true;
+}
+
 Status LossyFrameHeuristics(const FrameHeader& frame_header,
                             PassesEncoderState* enc_state,
                             ModularFrameEncoder* modular_frame_encoder,
-                            const Image3F* original_pixels, Image3F* opsin,
+                            const Image3F* linear, Image3F* opsin,
                             const Rect& rect, const JxlCmsInterface& cms,
                             ThreadPool* pool, AuxOut* aux_out) {
   const CompressParams& cparams = enc_state->cparams;
@@ -750,11 +983,12 @@ Status LossyFrameHeuristics(const FrameHeader& frame_header,
   ImageFeatures& image_features = shared.image_features;
   DequantMatrices& matrices = shared.matrices;
   Quantizer& quantizer = shared.quantizer;
+  ImageF& initial_quant_masking1x1 = enc_state->initial_quant_masking1x1;
   ImageI& raw_quant_field = shared.raw_quant_field;
   ColorCorrelationMap& cmap = shared.cmap;
   AcStrategyImage& ac_strategy = shared.ac_strategy;
-  ImageB& epf_sharpness = shared.epf_sharpness;
   BlockCtxMap& block_ctx_map = shared.block_ctx_map;
+  JxlMemoryManager* memory_manager = enc_state->memory_manager();
 
   // Find and subtract splines.
   if (cparams.custom_splines.HasAny()) {
@@ -799,27 +1033,33 @@ Status LossyFrameHeuristics(const FrameHeader& frame_header,
   //
   // output: Gaborished XYB, CfL, ACS, raw quant field, EPF control field.
 
-  ArControlFieldHeuristics ar_heuristics;
   AcStrategyHeuristics acs_heuristics(cparams);
   CfLHeuristics cfl_heuristics;
   ImageF initial_quant_field;
   ImageF initial_quant_masking;
-  ImageF initial_quant_masking1x1;
 
   // Compute an initial estimate of the quantization field.
   // Call InitialQuantField only in Hare mode or slower. Otherwise, rely
   // on simple heuristics in FindBestAcStrategy, or set a constant for Falcon
   // mode.
-  if (cparams.speed_tier > SpeedTier::kHare) {
+  if (cparams.speed_tier > SpeedTier::kHare ||
-    JXL_ASSIGN_OR_RETURN(
+      cparams.disable_percepeptual_optimizations) {
-        initial_quant_field,
+    JXL_ASSIGN_OR_RETURN(initial_quant_field,
-        ImageF::Create(frame_dim.xsize_blocks, frame_dim.ysize_blocks));
+                         ImageF::Create(memory_manager, frame_dim.xsize_blocks,
-    JXL_ASSIGN_OR_RETURN(
+                                        frame_dim.ysize_blocks));
-        initial_quant_masking,
+    JXL_ASSIGN_OR_RETURN(initial_quant_masking,
-        ImageF::Create(frame_dim.xsize_blocks, frame_dim.ysize_blocks));
+                         ImageF::Create(memory_manager, frame_dim.xsize_blocks,
+                                        frame_dim.ysize_blocks));
     float q = 0.79 / cparams.butteraugli_distance;
     FillImage(q, &initial_quant_field);
-    FillImage(1.0f / (q + 0.001f), &initial_quant_masking);
+    float masking = 1.0f / (q + 0.001f);
+    FillImage(masking, &initial_quant_masking);
+    if (cparams.disable_percepeptual_optimizations) {
+      JXL_ASSIGN_OR_RETURN(
+          initial_quant_masking1x1,
+          ImageF::Create(memory_manager, frame_dim.xsize, frame_dim.ysize));
+      FillImage(masking, &initial_quant_masking1x1);
+    }
     quantizer.ComputeGlobalScaleAndQuant(quant_dc, q, 0);
   } else {
     // Call this here, as it relies on pre-gaborish values.
@@ -850,17 +1090,15 @@ Status LossyFrameHeuristics(const FrameHeader& frame_header,
   }
 
   if (initialize_global_state) {
-    JXL_RETURN_IF_ERROR(
+    JXL_RETURN_IF_ERROR(FindBestDequantMatrices(
-        FindBestDequantMatrices(cparams, modular_frame_encoder, &matrices));
+        memory_manager, cparams, modular_frame_encoder, &matrices));
   }
 
-  JXL_RETURN_IF_ERROR(cfl_heuristics.Init(rect));
+  JXL_RETURN_IF_ERROR(cfl_heuristics.Init(memory_manager, rect));
   acs_heuristics.Init(*opsin, rect, initial_quant_field, initial_quant_masking,
                       initial_quant_masking1x1, &matrices);
 
-  std::atomic<bool> has_error{false};
   auto process_tile = [&](const uint32_t tid, const size_t thread) {
-    if (has_error) return;
     size_t n_enc_tiles = DivCeil(frame_dim.xsize_blocks, kEncTileDimInBlocks);
     size_t tx = tid % n_enc_tiles;
     size_t ty = tid / n_enc_tiles;
@@ -885,15 +1123,6 @@ Status LossyFrameHeuristics(const FrameHeader& frame_header,
     // Choose block sizes.
     acs_heuristics.ProcessRect(r, cmap, &ac_strategy, thread);
 
-    // Choose amount of post-processing smoothing.
-    // TODO(veluca): should this go *after* AdjustQuantField?
-    if (!ar_heuristics.RunRect(cparams, frame_header, r, *opsin, rect,
-                               initial_quant_field, ac_strategy, &epf_sharpness,
-                               thread)) {
-      has_error = true;
-      return;
-    }
-
     // Always set the initial quant field, so we can compute the CfL map with
     // more accuracy. The initial quant field might change in slower modes, but
     // adjusting the quant field with butteraugli when all the other encoding
@@ -915,18 +1144,18 @@ Status LossyFrameHeuristics(const FrameHeader& frame_header,
           DivCeil(frame_dim.ysize_blocks, kEncTileDimInBlocks),
       [&](const size_t num_threads) {
         acs_heuristics.PrepareForThreads(num_threads);
-        ar_heuristics.PrepareForThreads(num_threads);
         cfl_heuristics.PrepareForThreads(num_threads);
         return true;
       },
       process_tile, "Enc Heuristics"));
-  if (has_error) return JXL_FAILURE("Enc Heuristics failed");
 
   JXL_RETURN_IF_ERROR(acs_heuristics.Finalize(frame_dim, ac_strategy, aux_out));
 
   // Refine quantization levels.
-  if (!streaming_mode) {
+  if (!streaming_mode && !cparams.disable_percepeptual_optimizations) {
-    JXL_RETURN_IF_ERROR(FindBestQuantizer(frame_header, original_pixels, *opsin,
+    ImageB& epf_sharpness = shared.epf_sharpness;
+    FillPlane(static_cast<uint8_t>(4), &epf_sharpness, Rect(epf_sharpness));
+    JXL_RETURN_IF_ERROR(FindBestQuantizer(frame_header, linear, *opsin,
                                           initial_quant_field, enc_state, cms,
                                           pool, aux_out));
   }

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

@@ -32,10 +32,15 @@ class ModularFrameEncoder;
 Status LossyFrameHeuristics(const FrameHeader& frame_header,
                             PassesEncoderState* enc_state,
                             ModularFrameEncoder* modular_frame_encoder,
-                            const Image3F* original_pixels, Image3F* opsin,
+                            const Image3F* linear, Image3F* opsin,
                             const Rect& rect, const JxlCmsInterface& cms,
                             ThreadPool* pool, AuxOut* aux_out);
 
+Status ComputeARHeuristics(const FrameHeader& frame_header,
+                           PassesEncoderState* enc_state,
+                           const Image3F& orig_opsin, const Rect& rect,
+                           ThreadPool* pool);
+
 void FindBestBlockEntropyModel(PassesEncoderState& enc_state);
 
 Status DownsampleImage2_Iterative(Image3F* opsin);

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

@@ -5,14 +5,13 @@
 
 #include "lib/jxl/enc_icc_codec.h"
 
-#include <stdint.h>
+#include <jxl/memory_manager.h>
 
+#include <cstdint>
 #include <limits>
 #include <map>
-#include <string>
 #include <vector>
 
-#include "lib/jxl/base/byte_order.h"
 #include "lib/jxl/color_encoding_internal.h"
 #include "lib/jxl/enc_ans.h"
 #include "lib/jxl/enc_aux_out.h"
@@ -33,9 +32,10 @@ namespace {
 // elements at the bottom of the rightmost column. The input is the input matrix
 // in scanline order, the output is the result matrix in scanline order, with
 // missing elements skipped over (this may occur at multiple positions).
-void Unshuffle(uint8_t* data, size_t size, size_t width) {
+void Unshuffle(JxlMemoryManager* memory_manager, uint8_t* data, size_t size,
+               size_t width) {
   size_t height = (size + width - 1) / width;  // amount of rows of input
-  PaddedBytes result(size);
+  PaddedBytes result(memory_manager, size);
   // i = input index, j output index
   size_t s = 0;
   size_t j = 0;
@@ -57,6 +57,7 @@ Status PredictAndShuffle(size_t stride, size_t width, int order, size_t num,
                          const uint8_t* data, size_t size, size_t* pos,
                          PaddedBytes* result) {
   JXL_RETURN_IF_ERROR(CheckOutOfBounds(*pos, num, size));
+  JxlMemoryManager* memory_manager = result->memory_manager();
   // Required by the specification, see decoder. stride * 4 must be < *pos.
   if (!*pos || ((*pos - 1u) >> 2u) < stride) {
     return JXL_FAILURE("Invalid stride");
@@ -69,7 +70,7 @@ Status PredictAndShuffle(size_t stride, size_t width, int order, size_t num,
     result->push_back(data[*pos + i] - predicted);
   }
   *pos += num;
-  if (width > 1) Unshuffle(result->data() + start, num, width);
+  if (width > 1) Unshuffle(memory_manager, result->data() + start, num, width);
   return true;
 }
 
@@ -105,8 +106,9 @@ constexpr size_t kSizeLimit = std::numeric_limits<uint32_t>::max() >> 2;
 // form that is easier to compress (more zeroes, ...) and will compress better
 // with brotli.
 Status PredictICC(const uint8_t* icc, size_t size, PaddedBytes* result) {
-  PaddedBytes commands;
+  JxlMemoryManager* memory_manager = result->memory_manager();
-  PaddedBytes data;
+  PaddedBytes commands{memory_manager};
+  PaddedBytes data{memory_manager};
 
   static_assert(sizeof(size_t) >= 4, "size_t is too short");
   // Fuzzer expects that PredictICC can accept any input,
@@ -118,7 +120,7 @@ Status PredictICC(const uint8_t* icc, size_t size, PaddedBytes* result) {
   EncodeVarInt(size, result);
 
   // Header
-  PaddedBytes header;
+  PaddedBytes header{memory_manager};
   header.append(ICCInitialHeaderPrediction());
   EncodeUint32(0, size, &header);
   for (size_t i = 0; i < kICCHeaderSize && i < size; i++) {
@@ -256,8 +258,8 @@ Status PredictICC(const uint8_t* icc, size_t size, PaddedBytes* result) {
   // but will not predict as well.
   while (pos <= size) {
     size_t last1 = pos;
-    PaddedBytes commands_add;
+    PaddedBytes commands_add{memory_manager};
-    PaddedBytes data_add;
+    PaddedBytes data_add{memory_manager};
 
     // This means the loop brought the position beyond the tag end.
     // If tagsize is nonsensical, any pos looks "ok-ish".
@@ -285,7 +287,7 @@ Status PredictICC(const uint8_t* icc, size_t size, PaddedBytes* result) {
           data_add.push_back(icc[pos]);
           pos++;
         }
-        Unshuffle(data_add.data() + start, num, 2);
+        Unshuffle(memory_manager, data_add.data() + start, num, 2);
       }
 
       if (tag == kCurvTag && tag_sane() && pos + tagsize <= size &&
@@ -430,7 +432,8 @@ Status PredictICC(const uint8_t* icc, size_t size, PaddedBytes* result) {
 Status WriteICC(const IccBytes& icc, BitWriter* JXL_RESTRICT writer,
                 size_t layer, AuxOut* JXL_RESTRICT aux_out) {
   if (icc.empty()) return JXL_FAILURE("ICC must be non-empty");
-  PaddedBytes enc;
+  JxlMemoryManager* memory_manager = writer->memory_manager();
+  PaddedBytes enc{memory_manager};
   JXL_RETURN_IF_ERROR(PredictICC(icc.data(), icc.size(), &enc));
   std::vector<std::vector<Token>> tokens(1);
   BitWriter::Allotment allotment(writer, 128);
@@ -448,8 +451,8 @@ Status WriteICC(const IccBytes& icc, BitWriter* JXL_RESTRICT writer,
   EntropyEncodingData code;
   std::vector<uint8_t> context_map;
   params.force_huffman = true;
-  BuildAndEncodeHistograms(params, kNumICCContexts, tokens, &code, &context_map,
+  BuildAndEncodeHistograms(memory_manager, params, kNumICCContexts, tokens,
-                           writer, layer, aux_out);
+                           &code, &context_map, writer, layer, aux_out);
   WriteTokens(tokens[0], code, context_map, 0, writer, layer, aux_out);
   return true;
 }

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

@@ -8,9 +8,8 @@
 
 // Compressed representation of ICC profiles.
 
-#include <stddef.h>
+#include <cstddef>
-#include <stdint.h>
+#include <cstdint>
-
 #include <vector>
 
 #include "lib/jxl/base/compiler_specific.h"

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

@@ -6,6 +6,7 @@
 #include "lib/jxl/enc_image_bundle.h"
 
 #include <jxl/cms_interface.h>
+#include <jxl/memory_manager.h>
 
 #include <atomic>
 #include <utility>
@@ -27,8 +28,10 @@ Status ApplyColorTransform(const ColorEncoding& c_current,
   // Changing IsGray is probably a bug.
   JXL_CHECK(c_current.IsGray() == c_desired.IsGray());
   bool is_gray = c_current.IsGray();
+  JxlMemoryManager* memory_amanger = color.memory_manager();
   if (out->xsize() < rect.xsize() || out->ysize() < rect.ysize()) {
-    JXL_ASSIGN_OR_RETURN(*out, Image3F::Create(rect.xsize(), rect.ysize()));
+    JXL_ASSIGN_OR_RETURN(
+        *out, Image3F::Create(memory_amanger, rect.xsize(), rect.ysize()));
   } else {
     out->ShrinkTo(rect.xsize(), rect.ysize());
   }
@@ -131,10 +134,12 @@ Status TransformIfNeeded(const ImageBundle& in, const ColorEncoding& c_desired,
     *out = &in;
     return true;
   }
+  JxlMemoryManager* memory_manager = in.memory_manager();
   // TODO(janwas): avoid copying via createExternal+copyBackToIO
   // instead of copy+createExternal+copyBackToIO
-  JXL_ASSIGN_OR_RETURN(Image3F color,
+  JXL_ASSIGN_OR_RETURN(
-                       Image3F::Create(in.color().xsize(), in.color().ysize()));
+      Image3F color,
+      Image3F::Create(memory_manager, in.color().xsize(), in.color().ysize()));
   CopyImageTo(in.color(), &color);
   store->SetFromImage(std::move(color), in.c_current());
 
@@ -142,8 +147,9 @@ Status TransformIfNeeded(const ImageBundle& in, const ColorEncoding& c_desired,
   if (in.HasExtraChannels()) {
     std::vector<ImageF> extra_channels;
     for (const ImageF& extra_channel : in.extra_channels()) {
-      JXL_ASSIGN_OR_RETURN(ImageF ec, ImageF::Create(extra_channel.xsize(),
+      JXL_ASSIGN_OR_RETURN(ImageF ec,
-                                                     extra_channel.ysize()));
+                           ImageF::Create(memory_manager, extra_channel.xsize(),
+                                          extra_channel.ysize()));
       CopyImageTo(extra_channel, &ec);
       extra_channels.emplace_back(std::move(ec));
     }

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

@@ -17,7 +17,11 @@ void ConvertToDiagonal(const Matrix2x2& A, Vector2& diag, Matrix2x2& U) {
   JXL_ASSERT(std::abs(A[0][1] - A[1][0]) < 1e-15);
 #endif
 
-  if (std::abs(A[0][1]) < 1e-15) {
+  double b = -(A[0][0] + A[1][1]);
+  double c = A[0][0] * A[1][1] - A[0][1] * A[0][1];
+  double d = b * b - 4.0 * c;
+
+  if (std::abs(A[0][1]) < 1e-10 || d < 0) {
     // Already diagonal.
     diag[0] = A[0][0];
     diag[1] = A[1][1];
@@ -25,9 +29,7 @@ void ConvertToDiagonal(const Matrix2x2& A, Vector2& diag, Matrix2x2& U) {
     U[0][1] = U[1][0] = 0.0;
     return;
   }
-  double b = -(A[0][0] + A[1][1]);
+
-  double c = A[0][0] * A[1][1] - A[0][1] * A[0][1];
-  double d = b * b - 4.0 * c;
   double sqd = std::sqrt(d);
   double l1 = (-b - sqd) * 0.5;
   double l2 = (-b + sqd) * 0.5;

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

@@ -73,14 +73,24 @@ TEST(LinAlgTest, ConvertToDiagonal) {
     VerifyOrthogonal(U, 1e-12);
     VerifyMatrixEqual(A, MatMul(U, MatMul(Diagonal(d), Transpose(U))), 1e-12);
   }
+  {
+    Matrix2x2 A;
+    A[0] = {0.000208649, 1.13687e-12};
+    A[1] = {1.13687e-12, 0.000208649};
+    Matrix2x2 U;
+    Vector2 d;
+    ConvertToDiagonal(A, d, U);
+    VerifyOrthogonal(U, 1e-12);
+    VerifyMatrixEqual(A, MatMul(U, MatMul(Diagonal(d), Transpose(U))), 1e-11);
+  }
   Rng rng(0);
-  for (size_t i = 0; i < 100; ++i) {
+  for (size_t i = 0; i < 1000000; ++i) {
     Matrix2x2 A = RandomSymmetricMatrix(rng, -1.0, 1.0);
     Matrix2x2 U;
     Vector2 d;
     ConvertToDiagonal(A, d, U);
     VerifyOrthogonal(U, 1e-12);
-    VerifyMatrixEqual(A, MatMul(U, MatMul(Diagonal(d), Transpose(U))), 1e-12);
+    VerifyMatrixEqual(A, MatMul(U, MatMul(Diagonal(d), Transpose(U))), 5e-10);
   }
 }
 

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

@@ -5,6 +5,8 @@
 
 #include "lib/jxl/enc_modular.h"
 
+#include <jxl/memory_manager.h>
+
 #include <array>
 #include <atomic>
 #include <cstddef>
@@ -425,10 +427,13 @@ void try_palettes(Image& gi, int& max_bitdepth, int& maxval,
 
 }  // namespace
 
-ModularFrameEncoder::ModularFrameEncoder(const FrameHeader& frame_header,
+ModularFrameEncoder::ModularFrameEncoder(JxlMemoryManager* memory_manager,
+                                         const FrameHeader& frame_header,
                                          const CompressParams& cparams_orig,
                                          bool streaming_mode)
-    : frame_dim_(frame_header.ToFrameDimensions()), cparams_(cparams_orig) {
+    : memory_manager_(memory_manager),
+      frame_dim_(frame_header.ToFrameDimensions()),
+      cparams_(cparams_orig) {
   size_t num_streams =
       ModularStreamId::Num(frame_dim_, frame_header.passes.num_passes);
   if (cparams_.ModularPartIsLossless()) {
@@ -461,7 +466,9 @@ ModularFrameEncoder::ModularFrameEncoder(const FrameHeader& frame_header,
         ModularOptions::TreeKind::kTrivialTreeNoPredictor;
     cparams_.options.nb_repeats = 0;
   }
-  stream_images_.resize(num_streams);
+  for (size_t i = 0; i < num_streams; ++i) {
+    stream_images_.emplace_back(memory_manager);
+  }
 
   // use a sensible default if nothing explicit is specified:
   // Squeeze for lossy, no squeeze for lossless
@@ -624,6 +631,7 @@ Status ModularFrameEncoder::ComputeEncodingData(
     const Rect& frame_area_rect, PassesEncoderState* JXL_RESTRICT enc_state,
     const JxlCmsInterface& cms, ThreadPool* pool, AuxOut* aux_out,
     bool do_color) {
+  JxlMemoryManager* memory_manager = enc_state->memory_manager();
   JXL_DEBUG_V(6, "Computing modular encoding data for frame %s",
               frame_header.DebugString().c_str());
 
@@ -682,8 +690,8 @@ Status ModularFrameEncoder::ComputeEncodingData(
       do_color ? metadata.bit_depth.bits_per_sample + (fp ? 0 : 1) : 0;
   Image& gi = stream_images_[0];
   JXL_ASSIGN_OR_RETURN(
-      gi, Image::Create(xsize, ysize, metadata.bit_depth.bits_per_sample,
+      gi, Image::Create(memory_manager, xsize, ysize,
-                        nb_chans));
+                        metadata.bit_depth.bits_per_sample, nb_chans));
   int c = 0;
   if (cparams_.color_transform == ColorTransform::kXYB &&
       cparams_.modular_mode == true) {
@@ -696,11 +704,12 @@ Status ModularFrameEncoder::ComputeEncodingData(
       cparams_.butteraugli_distance = 0;
     }
     if (cparams_.manual_xyb_factors.size() == 3) {
-      DequantMatricesSetCustomDC(&enc_state->shared.matrices,
+      DequantMatricesSetCustomDC(memory_manager, &enc_state->shared.matrices,
                                  cparams_.manual_xyb_factors.data());
       // TODO(jon): update max_bitdepth in this case
     } else {
-      DequantMatricesSetCustomDC(&enc_state->shared.matrices, enc_factors);
+      DequantMatricesSetCustomDC(memory_manager, &enc_state->shared.matrices,
+                                 enc_factors);
       max_bitdepth = 12;
     }
   }
@@ -1238,6 +1247,7 @@ Status ModularFrameEncoder::ComputeTokens(ThreadPool* pool) {
 Status ModularFrameEncoder::EncodeGlobalInfo(bool streaming_mode,
                                              BitWriter* writer,
                                              AuxOut* aux_out) {
+  JxlMemoryManager* memory_manager = writer->memory_manager();
   BitWriter::Allotment allotment(writer, 1);
   // If we are using brotli, or not using modular mode.
   if (tree_tokens_.empty() || tree_tokens_[0].empty()) {
@@ -1254,9 +1264,9 @@ Status ModularFrameEncoder::EncodeGlobalInfo(bool streaming_mode,
   {
     EntropyEncodingData tree_code;
     std::vector<uint8_t> tree_context_map;
-    BuildAndEncodeHistograms(params, kNumTreeContexts, tree_tokens_, &tree_code,
+    BuildAndEncodeHistograms(memory_manager, params, kNumTreeContexts,
-                             &tree_context_map, writer, kLayerModularTree,
+                             tree_tokens_, &tree_code, &tree_context_map,
-                             aux_out);
+                             writer, kLayerModularTree, aux_out);
     WriteTokens(tree_tokens_[0], tree_code, tree_context_map, 0, writer,
                 kLayerModularTree, aux_out);
   }
@@ -1264,8 +1274,9 @@ Status ModularFrameEncoder::EncodeGlobalInfo(bool streaming_mode,
   params.add_missing_symbols = streaming_mode;
   params.image_widths = image_widths_;
   // Write histograms.
-  BuildAndEncodeHistograms(params, (tree_.size() + 1) / 2, tokens_, &code_,
+  BuildAndEncodeHistograms(memory_manager, params, (tree_.size() + 1) / 2,
-                           &context_map_, writer, kLayerModularGlobal, aux_out);
+                           tokens_, &code_, &context_map_, writer,
+                           kLayerModularGlobal, aux_out);
   return true;
 }
 
@@ -1292,7 +1303,7 @@ Status ModularFrameEncoder::EncodeStream(BitWriter* writer, AuxOut* aux_out,
 
 void ModularFrameEncoder::ClearStreamData(const ModularStreamId& stream) {
   size_t stream_id = stream.ID(frame_dim_);
-  Image empty_image;
+  Image empty_image(stream_images_[stream_id].memory_manager());
   std::swap(stream_images_[stream_id], empty_image);
 }
 
@@ -1321,12 +1332,13 @@ Status ModularFrameEncoder::PrepareStreamParams(const Rect& rect,
                                                 bool do_color, bool groupwise) {
   size_t stream_id = stream.ID(frame_dim_);
   Image& full_image = stream_images_[0];
+  JxlMemoryManager* memory_manager = full_image.memory_manager();
   const size_t xsize = rect.xsize();
   const size_t ysize = rect.ysize();
   Image& gi = stream_images_[stream_id];
   if (stream_id > 0) {
-    JXL_ASSIGN_OR_RETURN(gi,
+    JXL_ASSIGN_OR_RETURN(gi, Image::Create(memory_manager, xsize, ysize,
-                         Image::Create(xsize, ysize, full_image.bitdepth, 0));
+                                           full_image.bitdepth, 0));
     // start at the first bigger-than-frame_dim.group_dim non-metachannel
     size_t c = full_image.nb_meta_channels;
     if (!groupwise) {
@@ -1344,7 +1356,8 @@ Status ModularFrameEncoder::PrepareStreamParams(const Rect& rect,
              rect.xsize() >> fc.hshift, rect.ysize() >> fc.vshift, fc.w, fc.h);
       if (r.xsize() == 0 || r.ysize() == 0) continue;
       gi_channel_[stream_id].push_back(c);
-      JXL_ASSIGN_OR_RETURN(Channel gc, Channel::Create(r.xsize(), r.ysize()));
+      JXL_ASSIGN_OR_RETURN(
+          Channel gc, Channel::Create(memory_manager, r.xsize(), r.ysize()));
       gc.hshift = fc.hshift;
       gc.vshift = fc.vshift;
       for (size_t y = 0; y < r.ysize(); ++y) {
@@ -1490,6 +1503,7 @@ Status ModularFrameEncoder::AddVarDCTDC(const FrameHeader& frame_header,
                                         size_t group_index, bool nl_dc,
                                         PassesEncoderState* enc_state,
                                         bool jpeg_transcode) {
+  JxlMemoryManager* memory_manager = dc.memory_manager();
   extra_dc_precision[group_index] = nl_dc ? 1 : 0;
   float mul = 1 << extra_dc_precision[group_index];
 
@@ -1515,8 +1529,9 @@ Status ModularFrameEncoder::AddVarDCTDC(const FrameHeader& frame_header,
   stream_options_[stream_id].histogram_params =
       stream_options_[0].histogram_params;
 
-  JXL_ASSIGN_OR_RETURN(stream_images_[stream_id],
+  JXL_ASSIGN_OR_RETURN(
-                       Image::Create(r.xsize(), r.ysize(), 8, 3));
+      stream_images_[stream_id],
+      Image::Create(memory_manager, r.xsize(), r.ysize(), 8, 3));
   if (nl_dc && stream_options_[stream_id].tree_kind ==
                    ModularOptions::TreeKind::kGradientFixedDC) {
     JXL_ASSERT(frame_header.chroma_subsampling.Is444());
@@ -1637,6 +1652,7 @@ Status ModularFrameEncoder::AddVarDCTDC(const FrameHeader& frame_header,
 Status ModularFrameEncoder::AddACMetadata(const Rect& r, size_t group_index,
                                           bool jpeg_transcode,
                                           PassesEncoderState* enc_state) {
+  JxlMemoryManager* memory_manager = enc_state->memory_manager();
   size_t stream_id = ModularStreamId::ACMetadata(group_index).ID(frame_dim_);
   stream_options_[stream_id].max_chan_size = 0xFFFFFF;
   if (stream_options_[stream_id].predictor != Predictor::Weighted) {
@@ -1661,15 +1677,19 @@ Status ModularFrameEncoder::AddACMetadata(const Rect& r, size_t group_index,
       stream_options_[0].histogram_params;
   // YToX, YToB, ACS + QF, EPF
   Image& image = stream_images_[stream_id];
-  JXL_ASSIGN_OR_RETURN(image, Image::Create(r.xsize(), r.ysize(), 8, 4));
+  JXL_ASSIGN_OR_RETURN(
+      image, Image::Create(memory_manager, r.xsize(), r.ysize(), 8, 4));
   static_assert(kColorTileDimInBlocks == 8, "Color tile size changed");
   Rect cr(r.x0() >> 3, r.y0() >> 3, (r.xsize() + 7) >> 3, (r.ysize() + 7) >> 3);
-  JXL_ASSIGN_OR_RETURN(image.channel[0],
+  JXL_ASSIGN_OR_RETURN(
-                       Channel::Create(cr.xsize(), cr.ysize(), 3, 3));
+      image.channel[0],
-  JXL_ASSIGN_OR_RETURN(image.channel[1],
+      Channel::Create(memory_manager, cr.xsize(), cr.ysize(), 3, 3));
-                       Channel::Create(cr.xsize(), cr.ysize(), 3, 3));
+  JXL_ASSIGN_OR_RETURN(
-  JXL_ASSIGN_OR_RETURN(image.channel[2],
+      image.channel[1],
-                       Channel::Create(r.xsize() * r.ysize(), 2, 0, 0));
+      Channel::Create(memory_manager, cr.xsize(), cr.ysize(), 3, 3));
+  JXL_ASSIGN_OR_RETURN(
+      image.channel[2],
+      Channel::Create(memory_manager, r.xsize() * r.ysize(), 2, 0, 0));
   ConvertPlaneAndClamp(cr, enc_state->shared.cmap.ytox_map,
                        Rect(image.channel[0].plane), &image.channel[0].plane);
   ConvertPlaneAndClamp(cr, enc_state->shared.cmap.ytob_map,
@@ -1696,8 +1716,8 @@ Status ModularFrameEncoder::AddACMetadata(const Rect& r, size_t group_index,
 }
 
 Status ModularFrameEncoder::EncodeQuantTable(
-    size_t size_x, size_t size_y, BitWriter* writer,
+    JxlMemoryManager* memory_manager, size_t size_x, size_t size_y,
-    const QuantEncoding& encoding, size_t idx,
+    BitWriter* writer, const QuantEncoding& encoding, size_t idx,
     ModularFrameEncoder* modular_frame_encoder) {
   JXL_ASSERT(encoding.qraw.qtable != nullptr);
   JXL_ASSERT(size_x * size_y * 3 == encoding.qraw.qtable->size());
@@ -1707,7 +1727,8 @@ Status ModularFrameEncoder::EncodeQuantTable(
         writer, nullptr, 0, ModularStreamId::QuantTable(idx)));
     return true;
   }
-  JXL_ASSIGN_OR_RETURN(Image image, Image::Create(size_x, size_y, 8, 3));
+  JXL_ASSIGN_OR_RETURN(Image image,
+                       Image::Create(memory_manager, size_x, size_y, 8, 3));
   for (size_t c = 0; c < 3; c++) {
     for (size_t y = 0; y < size_y; y++) {
       int32_t* JXL_RESTRICT row = image.channel[c].Row(y);
@@ -1728,7 +1749,9 @@ Status ModularFrameEncoder::AddQuantTable(size_t size_x, size_t size_y,
   JXL_ASSERT(encoding.qraw.qtable != nullptr);
   JXL_ASSERT(size_x * size_y * 3 == encoding.qraw.qtable->size());
   Image& image = stream_images_[stream_id];
-  JXL_ASSIGN_OR_RETURN(image, Image::Create(size_x, size_y, 8, 3));
+  JxlMemoryManager* memory_manager = image.memory_manager();
+  JXL_ASSIGN_OR_RETURN(image,
+                       Image::Create(memory_manager, size_x, size_y, 8, 3));
   for (size_t c = 0; c < 3; c++) {
     for (size_t y = 0; y < size_y; y++) {
       int32_t* JXL_RESTRICT row = image.channel[c].Row(y);

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

@@ -7,6 +7,7 @@
 #define LIB_JXL_ENC_MODULAR_H_
 
 #include <jxl/cms_interface.h>
+#include <jxl/memory_manager.h>
 
 #include <cstddef>
 #include <cstdint>
@@ -37,7 +38,8 @@ struct AuxOut;
 
 class ModularFrameEncoder {
  public:
-  ModularFrameEncoder(const FrameHeader& frame_header,
+  ModularFrameEncoder(JxlMemoryManager* memory_manager,
+                      const FrameHeader& frame_header,
                       const CompressParams& cparams_orig, bool streaming_mode);
   Status ComputeEncodingData(
       const FrameHeader& frame_header, const ImageMetadata& metadata,
@@ -77,7 +79,8 @@ class ModularFrameEncoder {
   // null, the quantization table in `encoding` is used, with dimensions `size_x
   // x size_y`. Otherwise, the table with ID `idx` is encoded from the given
   // `modular_frame_encoder`.
-  static Status EncodeQuantTable(size_t size_x, size_t size_y,
+  static Status EncodeQuantTable(JxlMemoryManager* memory_manager,
+                                 size_t size_x, size_t size_y,
                                  BitWriter* writer,
                                  const QuantEncoding& encoding, size_t idx,
                                  ModularFrameEncoder* modular_frame_encoder);
@@ -88,11 +91,14 @@ class ModularFrameEncoder {
   std::vector<size_t> ac_metadata_size;
   std::vector<uint8_t> extra_dc_precision;
 
+  JxlMemoryManager* memory_manager() const { return memory_manager_; }
+
  private:
   Status PrepareStreamParams(const Rect& rect, const CompressParams& cparams,
                              int minShift, int maxShift,
                              const ModularStreamId& stream, bool do_color,
                              bool groupwise);
+  JxlMemoryManager* memory_manager_;
   std::vector<Image> stream_images_;
   std::vector<ModularOptions> stream_options_;
   std::vector<uint32_t> quants_;

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

@@ -5,19 +5,14 @@
 
 #include "lib/jxl/enc_noise.h"
 
-#include <stdint.h>
-#include <stdlib.h>
-
 #include <algorithm>
+#include <cstdint>
+#include <cstdlib>
 #include <numeric>
 #include <utility>
 
-#include "lib/jxl/base/compiler_specific.h"
-#include "lib/jxl/chroma_from_luma.h"
-#include "lib/jxl/convolve.h"
 #include "lib/jxl/enc_aux_out.h"
 #include "lib/jxl/enc_optimize.h"
-#include "lib/jxl/image_ops.h"
 
 namespace jxl {
 namespace {