Add SkJpegGainmapEncoder::EncodeJpegR and EncodeHDRGM

Reland this patch with a default second argument for
skjpeg_destination_mgr, to avoid breaking code that uses that
interface.

Add a SkJpegGainmapEncoder class that is a friend of SkJpegEncoder and
add the functions SkJpegGainmapEncoder::EncodeJpegR and EncodeHDRGM.

This requires adding two new pieces of data to the SkJpegEncoder.
First is a list of new segments. For both JpegR and HDRGM, this
includes the XMP metadata. For HDRM this includes segments for the
encoded gainmap image. Second is a "suffix" to append to the end of
the base image. For JpegR, this is the encoded gainmap image.

Send the new segments to jpeg_write_marker (just like is currently
done for the ICC profile).

Plumb the "suffix" through to the skjpeg_destination_mgr, to be
appended after image is done writing (but before the stream is
flushed).

In both EncodeJpegR and EncodeHDRGM function, hard-code the XMP
metadata string (rather than constructing and serializing an SkDOM).

Include the math to transcode from any gainmap representation to
JpegR.

Add a test that transcodes an MPF gainmap image to a JpegR and
HDRGM image, and ensures that the rendered results are the same (up
to rounding error).

Bug: skia:14031
Change-Id: Ife54b83f703b785ae151106d5f7a01010d21a1d1
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/633459
Commit-Queue: Christopher Cameron <ccameron@google.com>
Reviewed-by: Brian Osman <brianosman@google.com>
diff --git a/BUILD.gn b/BUILD.gn
index 1efe409..b4fc83a 100644
--- a/BUILD.gn
+++ b/BUILD.gn
@@ -1234,6 +1234,10 @@
     "src/images/SkJPEGWriteUtility.cpp",
     "src/images/SkJpegEncoder.cpp",
   ]
+
+  if (skia_use_jpeg_gainmaps) {
+    sources += [ "src/images/SkJpegGainmapEncoder.cpp" ]
+  }
 }
 
 optional("jpegxl_decode") {
diff --git a/include/encode/SkJpegEncoder.h b/include/encode/SkJpegEncoder.h
index 0ce5ed9..8d4ae17 100644
--- a/include/encode/SkJpegEncoder.h
+++ b/include/encode/SkJpegEncoder.h
@@ -99,8 +99,32 @@
     bool onEncodeRows(int numRows) override;
 
 private:
+    friend class SkJpegGainmapEncoder;
     SkJpegEncoder(std::unique_ptr<SkJpegEncoderMgr>, const SkPixmap& src);
 
+    /**
+     *  Create a jpeg encoder that will encode the |src| pixels and |segmentData| to the |dst|
+     *  stream, followed by the data in |suffix|. |options| may be used to control the encoding
+     *  behavior.
+     *
+     *  |segmentCount| lists the number of metadata segments to include. |segmentMarker| lists the
+     *  marker type identifiers for each segment (e.g: 0xE1 for APP1), and |segmentData| lists the
+     *  data for each segment.
+     *
+     *  |dst|, |makerTypes|, |segmentData|, and |suffix| are unowned and must remain valid for the
+     *  lifetime of the object.
+     *
+     *  This returns nullptr on an invalid or unsupported |src|.
+     */
+    static constexpr size_t kSegmentDataMaxSize = 65533;
+    static std::unique_ptr<SkEncoder> Make(SkWStream* dst,
+                                           const SkPixmap& src,
+                                           const Options& options,
+                                           size_t segmentCount,
+                                           uint8_t* segmentMarkers,
+                                           SkData** segmentData,
+                                           SkData* suffix);
+
     std::unique_ptr<SkJpegEncoderMgr> fEncoderMgr;
     using INHERITED = SkEncoder;
 };
diff --git a/include/private/SkJpegGainmapEncoder.h b/include/private/SkJpegGainmapEncoder.h
new file mode 100644
index 0000000..9f9cb14
--- /dev/null
+++ b/include/private/SkJpegGainmapEncoder.h
@@ -0,0 +1,63 @@
+/*
+ * Copyright 2023 Google Inc.
+ *
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#ifndef SkJpegGainmapEncoder_DEFINED
+#define SkJpegGainmapEncoder_DEFINED
+
+#include "include/encode/SkJpegEncoder.h"
+
+class SkPixmap;
+class SkWStream;
+struct SkGainmapInfo;
+
+class SK_API SkJpegGainmapEncoder {
+public:
+    /**
+     *  Encode a JpegR image to |dst|.
+     *
+     *  The base image is specified by |base|, and |baseOptions| controls the encoding behavior for
+     *  the base image.
+     *
+     *  The gainmap image is specified by |gainmap|, and |gainmapOptions| controls the encoding
+     *  behavior for the gainmap image.
+     *
+     *  The rendering behavior of the gainmap image is provided in |gainmapInfo|. Not all gainmap
+     *  based images are compatible with JpegR. If the image is not compatible with JpegR, then
+     *  convert the gainmap to a format that is capable with JpegR. This conversion may result in
+     *  less precise quantization of the gainmap image.
+     *
+     *  Returns true on success.  Returns false on an invalid or unsupported |src|.
+     */
+    static bool EncodeJpegR(SkWStream* dst,
+                            const SkPixmap& base,
+                            const SkJpegEncoder::Options& baseOptions,
+                            const SkPixmap& gainmap,
+                            const SkJpegEncoder::Options& gainmapOptions,
+                            const SkGainmapInfo& gainmapInfo);
+
+    /**
+     *  Encode an HDRGM image to |dst|.
+     *
+     *  The base image is specified by |base|, and |baseOptions| controls the encoding behavior for
+     *  the base image.
+     *
+     *  The gainmap image is specified by |gainmap|, and |gainmapOptions| controls the encoding
+     *  behavior for the gainmap image.
+     *
+     *  The rendering behavior of the gainmap image is provided in |gainmapInfo|.
+     *
+     *  Returns true on success.  Returns false on an invalid or unsupported |src|.
+     */
+    static bool EncodeHDRGM(SkWStream* dst,
+                            const SkPixmap& base,
+                            const SkJpegEncoder::Options& baseOptions,
+                            const SkPixmap& gainmap,
+                            const SkJpegEncoder::Options& gainmapOptions,
+                            const SkGainmapInfo& gainmapInfo);
+};
+
+#endif
diff --git a/src/images/SkJPEGWriteUtility.cpp b/src/images/SkJPEGWriteUtility.cpp
index 29d6bcf..c4c824d 100644
--- a/src/images/SkJPEGWriteUtility.cpp
+++ b/src/images/SkJPEGWriteUtility.cpp
@@ -8,6 +8,7 @@
 
 #include "src/images/SkJPEGWriteUtility.h"
 
+#include "include/core/SkData.h"
 #include "include/core/SkStream.h"
 #include "include/private/base/SkTArray.h"
 #include "src/codec/SkJpegPriv.h"
@@ -54,11 +55,19 @@
             return;
         }
     }
+
+    if (dest->fSuffix) {
+        if (!dest->fStream->write(dest->fSuffix->data(), dest->fSuffix->size())) {
+            ERREXIT(cinfo, JERR_FILE_WRITE);
+            return;
+        }
+    }
+
     dest->fStream->flush();
 }
 
-skjpeg_destination_mgr::skjpeg_destination_mgr(SkWStream* stream)
-        : fStream(stream) {
+skjpeg_destination_mgr::skjpeg_destination_mgr(SkWStream* stream, SkData* suffix)
+        : fStream(stream), fSuffix(suffix) {
     this->init_destination = sk_init_destination;
     this->empty_output_buffer = sk_empty_output_buffer;
     this->term_destination = sk_term_destination;
diff --git a/src/images/SkJPEGWriteUtility.h b/src/images/SkJPEGWriteUtility.h
index fac4281..c607257 100644
--- a/src/images/SkJPEGWriteUtility.h
+++ b/src/images/SkJPEGWriteUtility.h
@@ -20,6 +20,7 @@
     #include "jpeglib.h"
 }
 
+class SkData;
 class SkWStream;
 
 void skjpeg_error_exit(j_common_ptr cinfo);
@@ -29,9 +30,12 @@
  * object.
  */
 struct SK_SPI skjpeg_destination_mgr : jpeg_destination_mgr {
-    skjpeg_destination_mgr(SkWStream* stream);
+    skjpeg_destination_mgr(SkWStream* stream, SkData* suffix = nullptr);
 
-    SkWStream*  fStream;
+    SkWStream* const fStream;
+    // Extra data to write after the Jpeg file's EndOfImage. Used for JpegR
+    // based gainmaps.
+    SkData* const fSuffix;
 
     enum {
         kBufferSize = 1024
diff --git a/src/images/SkJpegEncoder.cpp b/src/images/SkJpegEncoder.cpp
index a769450..f07a4f2 100644
--- a/src/images/SkJpegEncoder.cpp
+++ b/src/images/SkJpegEncoder.cpp
@@ -40,13 +40,12 @@
 
 class SkJpegEncoderMgr final : SkNoncopyable {
 public:
-
     /*
      * Create the decode manager
-     * Does not take ownership of stream
+     * Does not take ownership of stream or suffix.
      */
-    static std::unique_ptr<SkJpegEncoderMgr> Make(SkWStream* stream) {
-        return std::unique_ptr<SkJpegEncoderMgr>(new SkJpegEncoderMgr(stream));
+    static std::unique_ptr<SkJpegEncoderMgr> Make(SkWStream* stream, SkData* suffix) {
+        return std::unique_ptr<SkJpegEncoderMgr>(new SkJpegEncoderMgr(stream, suffix));
     }
 
     bool setParams(const SkImageInfo& srcInfo, const SkJpegEncoder::Options& options);
@@ -62,11 +61,7 @@
     }
 
 private:
-
-    SkJpegEncoderMgr(SkWStream* stream)
-        : fDstMgr(stream)
-        , fProc(nullptr)
-    {
+    SkJpegEncoderMgr(SkWStream* stream, SkData* suffix) : fDstMgr(stream, suffix), fProc(nullptr) {
         fCInfo.err = jpeg_std_error(&fErrMgr);
         fErrMgr.error_exit = skjpeg_error_exit;
         jpeg_create_compress(&fCInfo);
@@ -179,11 +174,21 @@
 
 std::unique_ptr<SkEncoder> SkJpegEncoder::Make(SkWStream* dst, const SkPixmap& src,
                                                const Options& options) {
+    return Make(dst, src, options, 0, nullptr, nullptr, nullptr);
+}
+
+std::unique_ptr<SkEncoder> SkJpegEncoder::Make(SkWStream* dst,
+                                               const SkPixmap& src,
+                                               const Options& options,
+                                               size_t segmentCount,
+                                               uint8_t* segmentMarkers,
+                                               SkData** segmentData,
+                                               SkData* suffix) {
     if (!SkPixmapIsValid(src)) {
         return nullptr;
     }
 
-    std::unique_ptr<SkJpegEncoderMgr> encoderMgr = SkJpegEncoderMgr::Make(dst);
+    std::unique_ptr<SkJpegEncoderMgr> encoderMgr = SkJpegEncoderMgr::Make(dst, suffix);
 
     skjpeg_error_mgr::AutoPushJmpBuf jmp(encoderMgr->errorMgr());
     if (setjmp(jmp)) {
@@ -197,6 +202,14 @@
     jpeg_set_quality(encoderMgr->cinfo(), options.fQuality, TRUE);
     jpeg_start_compress(encoderMgr->cinfo(), TRUE);
 
+    for (size_t i = 0; i < segmentCount; ++i) {
+        SkASSERT(segmentData[i]->size() <= kSegmentDataMaxSize);
+        jpeg_write_marker(encoderMgr->cinfo(),
+                          segmentMarkers[i],
+                          segmentData[i]->bytes(),
+                          segmentData[i]->size());
+    }
+
     sk_sp<SkData> icc =
             icc_from_color_space(src.info(), options.fICCProfile, options.fICCProfileDescription);
     if (icc) {
diff --git a/src/images/SkJpegGainmapEncoder.cpp b/src/images/SkJpegGainmapEncoder.cpp
new file mode 100644
index 0000000..8e7aef4
--- /dev/null
+++ b/src/images/SkJpegGainmapEncoder.cpp
@@ -0,0 +1,391 @@
+/*
+ * Copyright 2023 Google Inc.
+ *
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#include "include/private/SkJpegGainmapEncoder.h"
+
+#ifdef SK_ENCODE_JPEG
+
+#include "include/core/SkBitmap.h"
+#include "include/core/SkPixmap.h"
+#include "include/core/SkStream.h"
+#include "include/encode/SkJpegEncoder.h"
+#include "include/private/SkGainmapInfo.h"
+#include "src/codec/SkJpegPriv.h"
+
+#include <vector>
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// XMP helpers
+
+void xmp_write_per_channel_attr(
+        SkDynamicMemoryWStream& s, const char* attrib, SkScalar r, SkScalar g, SkScalar b) {
+    s.writeText(attrib);
+    s.writeText("=\"");
+    if (r == g && r == b) {
+        s.writeScalarAsText(r);
+    } else {
+        s.writeScalarAsText(r);
+        s.writeText(",");
+        s.writeScalarAsText(g);
+        s.writeText(",");
+        s.writeScalarAsText(b);
+    }
+    s.writeText("\"\n");
+}
+
+void xmp_write_scalar_attr(SkDynamicMemoryWStream& s, const char* attrib, SkScalar value) {
+    s.writeText(attrib);
+    s.writeText("=\"");
+    s.writeScalarAsText(value);
+    s.writeText("\"\n");
+}
+
+void xmp_write_decimal_attr(SkDynamicMemoryWStream& s,
+                            const char* attrib,
+                            int32_t value,
+                            bool newLine = true) {
+    s.writeText(attrib);
+    s.writeText("=\"");
+    s.writeDecAsText(value);
+    s.writeText("\"");
+    if (newLine) {
+        s.writeText("\n");
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// JpegR encoding
+
+static float mix(float a, float b, float amount) { return (b - a) * amount + a; }
+
+static float compute_range_scaling_factor(const SkGainmapInfo& info) {
+    // Find the minimum and maximum log-ratio values that can be encoded. We don't want to encode a
+    // range any larger than this.
+    const float loadLogRatioMaxComponent =
+            std::max({info.fLogRatioMax.fR, info.fLogRatioMax.fG, info.fLogRatioMax.fB});
+    const float loadLogRatioMinComponent =
+            std::min({info.fLogRatioMin.fR, info.fLogRatioMin.fG, info.fLogRatioMin.fB});
+    const float logRatioRSF =
+            sk_float_exp(std::max(loadLogRatioMaxComponent, -loadLogRatioMinComponent));
+
+    // Limit the range to only encode values that could reach the the maximum rendering brightness.
+    float hdrRatioMaxRSF = info.fHdrRatioMax;
+
+    return std::min(logRatioRSF, hdrRatioMaxRSF);
+}
+
+// Ensure that the specified gainmap can be encoded as a JpegR. If it cannot, transform it so that
+// it can.
+void make_jpegr_compatible_if_needed(SkGainmapInfo& info, SkBitmap& bitmap) {
+    // If fLogRatioMin == -fLogRatioMax and bitmap has a single channel then this is already
+    // compatible with JpegR.
+    if (info.fLogRatioMin.fR == -info.fLogRatioMax.fR &&
+        info.fLogRatioMin.fG == -info.fLogRatioMax.fG &&
+        info.fLogRatioMin.fB == -info.fLogRatioMax.fB &&
+        bitmap.colorType() == kGray_8_SkColorType) {
+        return;
+    }
+
+    // If not, transform the gainmap to a JpegR compatible format.
+    SkGainmapInfo oldInfo = info;
+    SkBitmap oldBitmap = bitmap;
+    SkBitmap newBitmap;
+    SkImageInfo newBitmapInfo =
+            SkImageInfo::Make(oldBitmap.dimensions(), kGray_8_SkColorType, kOpaque_SkAlphaType);
+    newBitmap.allocPixels(newBitmapInfo);
+
+    // Compute the new gainmap rangeScalingFactor and its log.
+    const float rangeScalingFactor = compute_range_scaling_factor(oldInfo);
+    const float newLogRatioMax = sk_float_log(rangeScalingFactor);
+    const float newLogRatioMin = -newLogRatioMax;
+
+    // Transform the old gainmap to the new range.
+    // TODO(ccameron): This is not remotely performant. Consider using a blit.
+    {
+        const SkColor4f oldLogRatioMin = oldInfo.fLogRatioMin;
+        const SkColor4f oldLogRatioMax = oldInfo.fLogRatioMax;
+        const SkColor4f gainmapGamma = oldInfo.fGainmapGamma;
+        auto newPixmap = newBitmap.pixmap();
+        for (int y = 0; y < oldBitmap.height(); ++y) {
+            for (int x = 0; x < oldBitmap.width(); ++x) {
+                // Convert the gainmap from its encoded value to oldLogRatio, which is log(HDR/SDR).
+                SkColor4f oldG = oldBitmap.getColor4f(x, y);
+                SkColor4f oldLogRatio = {
+                        mix(oldLogRatioMin.fR,
+                            oldLogRatioMax.fR,
+                            sk_float_pow(oldG.fR, gainmapGamma.fR)),
+                        mix(oldLogRatioMin.fG,
+                            oldLogRatioMax.fG,
+                            sk_float_pow(oldG.fG, gainmapGamma.fG)),
+                        mix(oldLogRatioMin.fB,
+                            oldLogRatioMax.fB,
+                            sk_float_pow(oldG.fB, gainmapGamma.fB)),
+                        1.f,
+                };
+
+                // Undo the log, computing HDR/SDR, and take the average of the components of this.
+                // TODO(ccameron): This assumes that the primaries of the base image are sRGB.
+                float averageLinearRatio = 0.2126f * sk_float_exp(oldLogRatio.fR) +
+                                           0.7152f * sk_float_exp(oldLogRatio.fG) +
+                                           0.0722f * sk_float_exp(oldLogRatio.fB);
+
+                // Compute log(HDR/SDR) for the average HDR/SDR ratio.
+                float newLogRatio = sk_float_log(averageLinearRatio);
+
+                // Convert from log(HDR/SDR) to the JpegR gainmap image encoding.
+                float newG = (newLogRatio - newLogRatioMin) / (newLogRatioMax - newLogRatioMin);
+                *newPixmap.writable_addr8(x, y) =
+                        std::min(std::max(sk_float_round(255.f * newG), 0.f), 255.f);
+            }
+        }
+    }
+
+    // Write the gainmap info for the transformed gainmap.
+    SkGainmapInfo newInfo;
+    newInfo.fLogRatioMin = {newLogRatioMin, newLogRatioMin, newLogRatioMin, 1.f};
+    newInfo.fLogRatioMax = {newLogRatioMax, newLogRatioMax, newLogRatioMax, 1.f};
+    newInfo.fGainmapGamma = {1.f, 1.f, 1.f, 1.f};
+    newInfo.fEpsilonSdr = 0.f;
+    newInfo.fEpsilonHdr = 0.f;
+    newInfo.fHdrRatioMin = 1.f;
+    newInfo.fHdrRatioMax = sk_float_exp(newLogRatioMax);
+    newInfo.fType = SkGainmapInfo::Type::kJpegR_Linear;
+    info = newInfo;
+    bitmap = newBitmap;
+}
+
+// Generate the XMP metadata for a JpegR file.
+sk_sp<SkData> get_jpegr_xmp_data(float rangeScalingFactor,
+                                 int32_t transferFunction,
+                                 int32_t itemLength) {
+    SkDynamicMemoryWStream s;
+    s.write(kXMPSig, sizeof(kXMPSig));
+    s.writeText(
+            "<x:xmpmeta xmlns:x=\"adobe:ns:meta/\" x:xmptk=\"Adobe XMP Core 5.1.2\">\n"
+            "<rdf:RDF xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n"
+            "<rdf:Description xmlns:GContainer=\"http://ns.google.com/photos/1.0/container/\" "
+            "xmlns:RecoveryMap=\"http://ns.google.com/photos/1.0/recoverymap/\">\n"
+            "<GContainer:Version>1</GContainer:Version>\n"
+            "<GContainer:Directory>\n"
+            "<rdf:Seq>\n"
+            "<rdf:li>\n"
+            "<GContainer:Item GContainer:ItemSemantic=\"Primary\"\n"
+            "GContainer:ItemMime=\"image/jpeg\"\n");
+    xmp_write_decimal_attr(s, "RecoveryMap:Version", 1);
+    xmp_write_scalar_attr(s, "RecoveryMap:RangeScalingFactor", rangeScalingFactor);
+    xmp_write_decimal_attr(s, "RecoveryMap:TransferFunction", transferFunction, /*newLine=*/false);
+    s.writeText("/>\n");
+    s.writeText(
+            "</rdf:li>\n"
+            "<rdf:li>\n"
+            "<GContainer:Item GContainer:ItemSemantic=\"RecoveryMap\"\n"
+            "GContainer:ItemMime=\"image/jpeg\"\n");
+    xmp_write_decimal_attr(s, "GContainer:ItemLength", itemLength, /*newLine=*/false);
+    s.writeText("/>\n");
+    s.writeText(
+            "</rdf:li>\n"
+            "</rdf:Seq>\n"
+            "</GContainer:Directory>\n"
+            "</rdf:Description>\n"
+            "</rdf:RDF>\n"
+            "</x:xmpmeta>\n");
+    return s.detachAsData();
+}
+
+bool SkJpegGainmapEncoder::EncodeJpegR(SkWStream* dst,
+                                       const SkPixmap& base,
+                                       const SkJpegEncoder::Options& baseOptions,
+                                       const SkPixmap& gainmap,
+                                       const SkJpegEncoder::Options& gainmapOptions,
+                                       const SkGainmapInfo& gainmapInfo) {
+    // Transform the gainmap to be compatible with JpegR, if needed.
+    SkBitmap gainmapJpegR;
+    gainmapJpegR.installPixels(gainmap);
+    SkGainmapInfo gainmapInfoJpegR = gainmapInfo;
+    make_jpegr_compatible_if_needed(gainmapInfoJpegR, gainmapJpegR);
+
+    // Encode the gainmap as a Jpeg.
+    SkDynamicMemoryWStream gainmapEncodeStream;
+    if (!SkJpegEncoder::Encode(&gainmapEncodeStream, gainmapJpegR.pixmap(), gainmapOptions)) {
+        return false;
+    }
+    sk_sp<SkData> gainmapEncoded = gainmapEncodeStream.detachAsData();
+
+    // Compute the XMP metadata.
+    sk_sp<SkData> xmpMetadata =
+            get_jpegr_xmp_data(gainmapInfoJpegR.fHdrRatioMax, 0, gainmapEncoded->size());
+
+    // Send this to the base image encoder.
+    uint8_t segmentMarker = kXMPMarker;
+    SkData* segmentData = xmpMetadata.get();
+    auto encoder = SkJpegEncoder::Make(
+            dst, base, baseOptions, 1, &segmentMarker, &segmentData, gainmapEncoded.get());
+    return encoder.get() && encoder->encodeRows(base.height());
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// HDRGM encoding
+
+// Generate the XMP metadata for an HDRGM file.
+sk_sp<SkData> get_hdrgm_xmp_data(const SkGainmapInfo& gainmapInfo) {
+    const float kLog2 = sk_float_log(2.f);
+    SkDynamicMemoryWStream s;
+    s.write(kXMPSig, sizeof(kXMPSig));
+    s.writeText(
+            "<x:xmpmeta xmlns:x=\"adobe:ns:meta/\" x:xmptk=\"XMP Core 5.5.0\">\n"
+            "<rdf:RDF xmlns:rdf=\"http://www.w3.org/1999/02/22-rdf-syntax-ns#\">\n"
+            "<rdf:Description rdf:about=\"\"\n"
+            "xmlns:hdrgm=\"http://ns.adobe.com/hdr-gain-map/1.0/\"\n"
+            "hdrgm:Version=\"1.0\"\n");
+    xmp_write_per_channel_attr(s,
+                               "hdrgm:GainMapMin",
+                               gainmapInfo.fLogRatioMin.fR / kLog2,
+                               gainmapInfo.fLogRatioMin.fG / kLog2,
+                               gainmapInfo.fLogRatioMin.fB / kLog2);
+    xmp_write_per_channel_attr(s,
+                               "hdrgm:GainMapMax",
+                               gainmapInfo.fLogRatioMax.fR / kLog2,
+                               gainmapInfo.fLogRatioMax.fG / kLog2,
+                               gainmapInfo.fLogRatioMax.fB / kLog2);
+    xmp_write_per_channel_attr(s,
+                               "hdrgm:Gamma",
+                               gainmapInfo.fGainmapGamma.fR,
+                               gainmapInfo.fGainmapGamma.fG,
+                               gainmapInfo.fGainmapGamma.fB);
+    xmp_write_per_channel_attr(s,
+                               "hdrgm:OffsetSDR",
+                               gainmapInfo.fEpsilonSdr,
+                               gainmapInfo.fEpsilonSdr,
+                               gainmapInfo.fEpsilonSdr);
+    xmp_write_per_channel_attr(s,
+                               "hdrgm:OffsetHDR",
+                               gainmapInfo.fEpsilonHdr,
+                               gainmapInfo.fEpsilonHdr,
+                               gainmapInfo.fEpsilonHdr);
+    xmp_write_scalar_attr(
+            s, "hdrgm:HDRCapacityMin", sk_float_log(gainmapInfo.fHdrRatioMin) / kLog2);
+    xmp_write_scalar_attr(
+            s, "hdrgm:HDRCapacityMax", sk_float_log(gainmapInfo.fHdrRatioMax) / kLog2);
+    s.writeText("hdrgm:BaseRendition=\"");
+    switch (gainmapInfo.fBaseImageType) {
+        case SkGainmapInfo::BaseImageType::kSDR:
+            s.writeText("SDR");
+            break;
+        case SkGainmapInfo::BaseImageType::kHDR:
+            s.writeText("HDR");
+            break;
+    }
+    s.writeText(
+            "\"/>\n"
+            "</rdf:RDF>\n"
+            "</x:xmpmeta>");
+    return s.detachAsData();
+}
+
+// Split an SkData into segments.
+std::vector<sk_sp<SkData>> get_hdrgm_image_segments(sk_sp<SkData> image,
+                                                    size_t segmentMaxDataSize) {
+    // Compute the total size of the header to a gainmap image segment (not including the 2 bytes
+    // for the segment size, which the encoder is responsible for writing).
+    constexpr size_t kGainmapHeaderSize = sizeof(kGainmapSig) + 2 * kGainmapMarkerIndexSize;
+
+    // Compute the payload size for each segment.
+    const size_t kGainmapPayloadSize = segmentMaxDataSize - kGainmapHeaderSize;
+
+    // Compute the number of segments we'll need.
+    const size_t segmentCount = (image->size() + kGainmapPayloadSize - 1) / kGainmapPayloadSize;
+    std::vector<sk_sp<SkData>> result;
+    result.reserve(segmentCount);
+
+    // Move |imageData| through |image| until it hits |imageDataEnd|.
+    const uint8_t* imageData = image->bytes();
+    const uint8_t* imageDataEnd = image->bytes() + image->size();
+    while (imageData < imageDataEnd) {
+        SkDynamicMemoryWStream segmentStream;
+
+        // Write the signature.
+        segmentStream.write(kGainmapSig, sizeof(kGainmapSig));
+
+        // Write the segment index as big-endian.
+        size_t segmentIndex = result.size() + 1;
+        uint8_t segmentIndexBytes[2] = {
+                static_cast<uint8_t>(segmentIndex / 256u),
+                static_cast<uint8_t>(segmentIndex % 256u),
+        };
+        segmentStream.write(segmentIndexBytes, sizeof(segmentIndexBytes));
+
+        // Write the segment count as big-endian.
+        uint8_t segmentCountBytes[2] = {
+                static_cast<uint8_t>(segmentCount / 256u),
+                static_cast<uint8_t>(segmentCount % 256u),
+        };
+        segmentStream.write(segmentCountBytes, sizeof(segmentCountBytes));
+
+        // Verify that our header size math is correct.
+        SkASSERT(segmentStream.bytesWritten() == kGainmapHeaderSize);
+
+        // Write the rest of the segment.
+        size_t bytesToWrite =
+                std::min(imageDataEnd - imageData, static_cast<intptr_t>(kGainmapPayloadSize));
+        segmentStream.write(imageData, bytesToWrite);
+        imageData += bytesToWrite;
+
+        // Verify that our data size math is correct.
+        if (segmentIndex == segmentCount) {
+            SkASSERT(segmentStream.bytesWritten() <= segmentMaxDataSize);
+        } else {
+            SkASSERT(segmentStream.bytesWritten() == segmentMaxDataSize);
+        }
+        result.push_back(segmentStream.detachAsData());
+    }
+
+    // Verify that our segment count math was correct.
+    SkASSERT(imageData == imageDataEnd);
+    SkASSERT(result.size() == segmentCount);
+    return result;
+}
+
+bool SkJpegGainmapEncoder::EncodeHDRGM(SkWStream* dst,
+                                       const SkPixmap& base,
+                                       const SkJpegEncoder::Options& baseOptions,
+                                       const SkPixmap& gainmap,
+                                       const SkJpegEncoder::Options& gainmapOptions,
+                                       const SkGainmapInfo& gainmapInfo) {
+    // Encode the gainmap as a Jpeg, and split it into segments.
+    SkDynamicMemoryWStream gainmapEncodeStream;
+    if (!SkJpegEncoder::Encode(&gainmapEncodeStream, gainmap, gainmapOptions)) {
+        return false;
+    }
+    std::vector<sk_sp<SkData>> gainmapSegments = get_hdrgm_image_segments(
+            gainmapEncodeStream.detachAsData(), SkJpegEncoder::kSegmentDataMaxSize);
+
+    // Compute the XMP metadata.
+    sk_sp<SkData> xmpMetadata = get_hdrgm_xmp_data(gainmapInfo);
+
+    // Merge these into the list of segments to send to the encoder.
+    std::vector<uint8_t> segmentMarker;
+    std::vector<SkData*> segmentData;
+    segmentMarker.push_back(kXMPMarker);
+    segmentData.push_back(xmpMetadata.get());
+    for (auto& gainmapSegment : gainmapSegments) {
+        segmentMarker.push_back(kGainmapMarker);
+        segmentData.push_back(gainmapSegment.get());
+    }
+    SkASSERT(segmentMarker.size() == segmentData.size());
+
+    // Send this to the base image encoder.
+    auto encoder = SkJpegEncoder::Make(dst,
+                                       base,
+                                       baseOptions,
+                                       segmentMarker.size(),
+                                       segmentMarker.data(),
+                                       segmentData.data(),
+                                       nullptr);
+    return encoder.get() && encoder->encodeRows(base.height());
+}
+
+#endif  // SK_ENCODE_JPEG
diff --git a/tests/JpegGainmapTest.cpp b/tests/JpegGainmapTest.cpp
index e06d07e..1959fe1 100644
--- a/tests/JpegGainmapTest.cpp
+++ b/tests/JpegGainmapTest.cpp
@@ -8,11 +8,16 @@
 #include "include/codec/SkAndroidCodec.h"
 #include "include/codec/SkCodec.h"
 #include "include/core/SkBitmap.h"
+#include "include/core/SkCanvas.h"
 #include "include/core/SkColor.h"
+#include "include/core/SkImageEncoder.h"
 #include "include/core/SkSize.h"
 #include "include/core/SkStream.h"
 #include "include/core/SkTypes.h"
+#include "include/encode/SkJpegEncoder.h"
 #include "include/private/SkGainmapInfo.h"
+#include "include/private/SkGainmapShader.h"
+#include "include/private/SkJpegGainmapEncoder.h"
 #include "src/codec/SkJpegMultiPicture.h"
 #include "src/codec/SkJpegSegmentScan.h"
 #include "tests/Test.h"
@@ -122,6 +127,107 @@
     REPORTER_ASSERT(r, bitmaps[2].getColor(575, 767) == 0xFFB5B5B5);
 }
 
+// Decode an image and its gainmap.
+template <typename Reporter>
+void decode_all(Reporter& r,
+                std::unique_ptr<SkStream> stream,
+                SkBitmap& baseBitmap,
+                SkBitmap& gainmapBitmap,
+                SkGainmapInfo& gainmapInfo) {
+    // Decode the base bitmap.
+    std::unique_ptr<SkCodec> baseCodec = SkCodec::MakeFromStream(std::move(stream));
+    REPORTER_ASSERT(r, baseCodec);
+    baseBitmap.allocPixels(baseCodec->getInfo());
+    REPORTER_ASSERT(r,
+                    SkCodec::kSuccess == baseCodec->getPixels(baseBitmap.info(),
+                                                              baseBitmap.getPixels(),
+                                                              baseBitmap.rowBytes()));
+    std::unique_ptr<SkAndroidCodec> androidCodec =
+            SkAndroidCodec::MakeFromCodec(std::move(baseCodec));
+    REPORTER_ASSERT(r, androidCodec);
+
+    // Extract the gainmap info and stream.
+    std::unique_ptr<SkStream> gainmapStream;
+    REPORTER_ASSERT(r, androidCodec->getAndroidGainmap(&gainmapInfo, &gainmapStream));
+    REPORTER_ASSERT(r, gainmapStream);
+
+    // Decode the gainmap bitmap.
+    std::unique_ptr<SkCodec> gainmapCodec = SkCodec::MakeFromStream(std::move(gainmapStream));
+    REPORTER_ASSERT(r, gainmapCodec);
+    SkBitmap bm;
+    bm.allocPixels(gainmapCodec->getInfo());
+    gainmapBitmap.allocPixels(gainmapCodec->getInfo());
+    REPORTER_ASSERT(r,
+                    SkCodec::kSuccess == gainmapCodec->getPixels(gainmapBitmap.info(),
+                                                                 gainmapBitmap.getPixels(),
+                                                                 gainmapBitmap.rowBytes()));
+}
+
+// Render an applied gainmap.
+SkBitmap render_gainmap(const SkImageInfo& renderInfo,
+                        float renderHdrRatio,
+                        const SkBitmap& baseBitmap,
+                        const SkBitmap& gainmapBitmap,
+                        const SkGainmapInfo& gainmapInfo,
+                        int x,
+                        int y) {
+    SkRect baseRect = SkRect::MakeXYWH(x, y, renderInfo.width(), renderInfo.height());
+
+    float scaleX = gainmapBitmap.width() / static_cast<float>(baseBitmap.width());
+    float scaleY = gainmapBitmap.height() / static_cast<float>(baseBitmap.height());
+    SkRect gainmapRect = SkRect::MakeXYWH(baseRect.x() * scaleX,
+                                          baseRect.y() * scaleY,
+                                          baseRect.width() * scaleX,
+                                          baseRect.height() * scaleY);
+
+    SkRect dstRect = SkRect::Make(renderInfo.dimensions());
+
+    sk_sp<SkImage> baseImage = SkImage::MakeFromBitmap(baseBitmap);
+    sk_sp<SkImage> gainmapImage = SkImage::MakeFromBitmap(gainmapBitmap);
+    sk_sp<SkShader> shader = SkGainmapShader::Make(baseImage,
+                                                   baseRect,
+                                                   SkSamplingOptions(),
+                                                   gainmapImage,
+                                                   gainmapRect,
+                                                   SkSamplingOptions(),
+                                                   gainmapInfo,
+                                                   dstRect,
+                                                   renderHdrRatio,
+                                                   renderInfo.refColorSpace());
+
+    SkBitmap result;
+    result.allocPixels(renderInfo);
+    result.eraseColor(SK_ColorTRANSPARENT);
+    SkCanvas canvas(result);
+
+    SkPaint paint;
+    paint.setShader(shader);
+    canvas.drawRect(dstRect, paint);
+
+    return result;
+}
+
+// Render a single pixel of an applied gainmap and return it.
+SkColor4f render_gainmap_pixel(float renderHdrRatio,
+                               const SkBitmap& baseBitmap,
+                               const SkBitmap& gainmapBitmap,
+                               const SkGainmapInfo& gainmapInfo,
+                               int x,
+                               int y) {
+    SkImageInfo testPixelInfo = SkImageInfo::Make(
+            1, 1, kRGBA_F16_SkColorType, kPremul_SkAlphaType, SkColorSpace::MakeSRGB());
+    SkBitmap testPixelBitmap = render_gainmap(
+            testPixelInfo, renderHdrRatio, baseBitmap, gainmapBitmap, gainmapInfo, x, y);
+    return testPixelBitmap.getColor4f(0, 0);
+}
+
+static bool approx_eq(float x, float y, float epsilon) { return std::abs(x - y) < epsilon; }
+
+static bool approx_eq_rgb(const SkColor4f& x, const SkColor4f& y, float epsilon) {
+    return approx_eq(x.fR, y.fR, epsilon) && approx_eq(x.fG, y.fG, epsilon) &&
+           approx_eq(x.fB, y.fB, epsilon);
+}
+
 DEF_TEST(AndroidCodec_jpegGainmap, r) {
     const struct Rec {
         const char* path;
@@ -168,50 +274,131 @@
             auto stream = GetResourceAsStream(rec.path, useFileStream);
             REPORTER_ASSERT(r, stream);
 
-            std::unique_ptr<SkCodec> codec = SkCodec::MakeFromStream(std::move(stream));
-            REPORTER_ASSERT(r, codec);
-
-            std::unique_ptr<SkAndroidCodec> androidCodec =
-                    SkAndroidCodec::MakeFromCodec(std::move(codec));
-            REPORTER_ASSERT(r, androidCodec);
-
+            SkBitmap baseBitmap;
+            SkBitmap gainmapBitmap;
             SkGainmapInfo gainmapInfo;
-            std::unique_ptr<SkStream> gainmapStream;
-            REPORTER_ASSERT(r, androidCodec->getAndroidGainmap(&gainmapInfo, &gainmapStream));
-            REPORTER_ASSERT(r, gainmapStream);
-
-            std::unique_ptr<SkCodec> gainmapCodec =
-                    SkCodec::MakeFromStream(std::move(gainmapStream));
-            REPORTER_ASSERT(r, gainmapCodec);
-
-            SkBitmap bm;
-            bm.allocPixels(gainmapCodec->getInfo());
-            REPORTER_ASSERT(r,
-                            SkCodec::kSuccess == gainmapCodec->getPixels(
-                                                         bm.info(), bm.getPixels(), bm.rowBytes()));
+            decode_all(r,
+                       GetResourceAsStream(rec.path, useFileStream),
+                       baseBitmap,
+                       gainmapBitmap,
+                       gainmapInfo);
 
             // Spot-check the image size and pixels.
-            REPORTER_ASSERT(r, bm.dimensions() == rec.dimensions);
-            REPORTER_ASSERT(r, bm.getColor(0, 0) == rec.originColor);
-            REPORTER_ASSERT(r,
-                            bm.getColor(rec.dimensions.fWidth - 1, rec.dimensions.fHeight - 1) ==
-                                    rec.farCornerColor);
+            REPORTER_ASSERT(r, gainmapBitmap.dimensions() == rec.dimensions);
+            REPORTER_ASSERT(r, gainmapBitmap.getColor(0, 0) == rec.originColor);
+            REPORTER_ASSERT(
+                    r,
+                    gainmapBitmap.getColor(rec.dimensions.fWidth - 1, rec.dimensions.fHeight - 1) ==
+                            rec.farCornerColor);
 
             // Verify the gainmap rendering parameters.
-            auto approxEq = [=](float x, float y) { return std::abs(x - y) < 1e-3f; };
+            constexpr float kEpsilon = 1e-3f;
+            REPORTER_ASSERT(r, approx_eq(gainmapInfo.fLogRatioMin.fR, rec.logRatioMin, kEpsilon));
+            REPORTER_ASSERT(r, approx_eq(gainmapInfo.fLogRatioMin.fG, rec.logRatioMin, kEpsilon));
+            REPORTER_ASSERT(r, approx_eq(gainmapInfo.fLogRatioMin.fB, rec.logRatioMin, kEpsilon));
 
-            REPORTER_ASSERT(r, approxEq(gainmapInfo.fLogRatioMin.fR, rec.logRatioMin));
-            REPORTER_ASSERT(r, approxEq(gainmapInfo.fLogRatioMin.fG, rec.logRatioMin));
-            REPORTER_ASSERT(r, approxEq(gainmapInfo.fLogRatioMin.fB, rec.logRatioMin));
+            REPORTER_ASSERT(r, approx_eq(gainmapInfo.fLogRatioMax.fR, rec.logRatioMax, kEpsilon));
+            REPORTER_ASSERT(r, approx_eq(gainmapInfo.fLogRatioMax.fG, rec.logRatioMax, kEpsilon));
+            REPORTER_ASSERT(r, approx_eq(gainmapInfo.fLogRatioMax.fB, rec.logRatioMax, kEpsilon));
 
-            REPORTER_ASSERT(r, approxEq(gainmapInfo.fLogRatioMax.fR, rec.logRatioMax));
-            REPORTER_ASSERT(r, approxEq(gainmapInfo.fLogRatioMax.fG, rec.logRatioMax));
-            REPORTER_ASSERT(r, approxEq(gainmapInfo.fLogRatioMax.fB, rec.logRatioMax));
-
-            REPORTER_ASSERT(r, approxEq(gainmapInfo.fHdrRatioMin, rec.hdrRatioMin));
-            REPORTER_ASSERT(r, approxEq(gainmapInfo.fHdrRatioMax, rec.hdrRatioMax));
+            REPORTER_ASSERT(r, approx_eq(gainmapInfo.fHdrRatioMin, rec.hdrRatioMin, kEpsilon));
+            REPORTER_ASSERT(r, approx_eq(gainmapInfo.fHdrRatioMax, rec.hdrRatioMax, kEpsilon));
 
             REPORTER_ASSERT(r, gainmapInfo.fType == rec.type);
         }
     }
 }
+
+#ifdef SK_ENCODE_JPEG
+DEF_TEST(AndroidCodec_jpegGainmapTranscode, r) {
+    const char* path = "images/iphone_13_pro.jpeg";
+    SkBitmap baseBitmap[2];
+    SkBitmap gainmapBitmap[2];
+    SkGainmapInfo gainmapInfo[2];
+
+    // Decode an MPF-based gainmap image.
+    decode_all(r, GetResourceAsStream(path), baseBitmap[0], gainmapBitmap[0], gainmapInfo[0]);
+
+    constexpr float kEpsilon = 1e-2f;
+    for (size_t i = 0; i < 2; ++i) {
+        SkDynamicMemoryWStream encodeStream;
+        bool encodeResult = false;
+
+        if (i == 0) {
+            // Transcode to JpegR.
+            encodeResult = SkJpegGainmapEncoder::EncodeJpegR(&encodeStream,
+                                                             baseBitmap[0].pixmap(),
+                                                             SkJpegEncoder::Options(),
+                                                             gainmapBitmap[0].pixmap(),
+                                                             SkJpegEncoder::Options(),
+                                                             gainmapInfo[0]);
+        } else {
+            // Transcode to HDRGM.
+            encodeResult = SkJpegGainmapEncoder::EncodeHDRGM(&encodeStream,
+                                                             baseBitmap[0].pixmap(),
+                                                             SkJpegEncoder::Options(),
+                                                             gainmapBitmap[0].pixmap(),
+                                                             SkJpegEncoder::Options(),
+                                                             gainmapInfo[0]);
+        }
+        REPORTER_ASSERT(r, encodeResult);
+
+        // Decode the just-encoded JpegR or HDRGM.
+        auto decodeStream = std::make_unique<SkMemoryStream>(encodeStream.detachAsData());
+        decode_all(r, std::move(decodeStream), baseBitmap[1], gainmapBitmap[1], gainmapInfo[1]);
+
+        // Verify that the representations are different.
+        REPORTER_ASSERT(r, gainmapInfo[0].fType != gainmapInfo[1].fType);
+        if (i == 0) {
+            // JpegR will have different rendering parameters.
+            REPORTER_ASSERT(r, gainmapInfo[0].fLogRatioMin != gainmapInfo[1].fLogRatioMin);
+        } else {
+            // HDRGM will have the same rendering parameters.
+            REPORTER_ASSERT(
+                    r,
+                    approx_eq_rgb(
+                            gainmapInfo[0].fLogRatioMin, gainmapInfo[1].fLogRatioMin, kEpsilon));
+            REPORTER_ASSERT(
+                    r,
+                    approx_eq_rgb(
+                            gainmapInfo[0].fLogRatioMax, gainmapInfo[1].fLogRatioMax, kEpsilon));
+            REPORTER_ASSERT(
+                    r,
+                    approx_eq_rgb(
+                            gainmapInfo[0].fGainmapGamma, gainmapInfo[1].fGainmapGamma, kEpsilon));
+            REPORTER_ASSERT(
+                    r, approx_eq(gainmapInfo[0].fEpsilonSdr, gainmapInfo[1].fEpsilonSdr, kEpsilon));
+            REPORTER_ASSERT(
+                    r, approx_eq(gainmapInfo[0].fEpsilonHdr, gainmapInfo[1].fEpsilonHdr, kEpsilon));
+            REPORTER_ASSERT(
+                    r,
+                    approx_eq(gainmapInfo[0].fHdrRatioMin, gainmapInfo[1].fHdrRatioMin, kEpsilon));
+            REPORTER_ASSERT(
+                    r,
+                    approx_eq(gainmapInfo[0].fHdrRatioMax, gainmapInfo[1].fHdrRatioMax, kEpsilon));
+        }
+
+#ifdef SK_ENABLE_SKSL
+        // Render a few pixels and verify that they come out the same. Rendering requires SkSL.
+        const struct Rec {
+            int x;
+            int y;
+            float hdrRatio;
+            SkColor4f expectedColor;
+        } recs[] = {
+                {1446, 1603, 1.05f, {0.984375f, 1.004883f, 1.008789f, 1.f}},
+                {1446, 1603, 100.f, {1.147461f, 1.170898f, 1.174805f, 1.f}},
+        };
+
+        for (const auto& rec : recs) {
+            SkColor4f p0 = render_gainmap_pixel(
+                    rec.hdrRatio, baseBitmap[0], gainmapBitmap[0], gainmapInfo[0], rec.x, rec.y);
+            SkColor4f p1 = render_gainmap_pixel(
+                    rec.hdrRatio, baseBitmap[1], gainmapBitmap[1], gainmapInfo[1], rec.x, rec.y);
+
+            REPORTER_ASSERT(r, approx_eq_rgb(p0, p1, kEpsilon));
+        }
+#endif  // SK_ENABLE_SKSL
+    }
+}
+#endif  // SK_ENCODE_JPEG