1 files changed, 46 insertions, 7 deletions

diff --git a/media/libjxl/src/lib/jxl/codec_in_out.h b/media/libjxl/src/lib/jxl/codec_in_out.h
index c1bff66366..23f0a4afeb 100644
--- a/media/libjxl/src/lib/jxl/codec_in_out.h
+++ b/media/libjxl/src/lib/jxl/codec_in_out.h
@@ -13,6 +13,7 @@
 #include <utility>
 #include <vector>
 
+#include "lib/jxl/alpha.h"
 #include "lib/jxl/base/data_parallel.h"
 #include "lib/jxl/common.h"
 #include "lib/jxl/frame_header.h"
@@ -133,18 +134,57 @@ class CodecInOut {
     }
     return true;
   }
-  // Calls PremultiplyAlpha for each ImageBundle (preview/frames).
-  void PremultiplyAlpha() {
+  // Performs "PremultiplyAlpha" for each ImageBundle (preview/frames).
+  bool PremultiplyAlpha() {
+    const auto doPremultiplyAlpha = [](ImageBundle& bundle) {
+      if (!bundle.HasAlpha()) return;
+      if (!bundle.HasColor()) return;
+      auto* color = bundle.color();
+      const auto* alpha = bundle.alpha();
+      JXL_CHECK(color->ysize() == alpha->ysize());
+      JXL_CHECK(color->xsize() == alpha->xsize());
+      for (size_t y = 0; y < color->ysize(); y++) {
+        ::jxl::PremultiplyAlpha(color->PlaneRow(0, y), color->PlaneRow(1, y),
+                                color->PlaneRow(2, y), alpha->Row(y),
+                                color->xsize());
+      }
+    };
     ExtraChannelInfo* eci = metadata.m.Find(ExtraChannel::kAlpha);
-    if (eci == nullptr || eci->alpha_associated) return;  // nothing to do
+    if (eci == nullptr || eci->alpha_associated) return false;
     if (metadata.m.have_preview) {
-      preview_frame.PremultiplyAlpha();
+      doPremultiplyAlpha(preview_frame);
     }
     for (ImageBundle& ib : frames) {
-      ib.PremultiplyAlpha();
+      doPremultiplyAlpha(ib);
     }
     eci->alpha_associated = true;
-    return;
+    return true;
+  }
+
+  bool UnpremultiplyAlpha() {
+    const auto doUnpremultiplyAlpha = [](ImageBundle& bundle) {
+      if (!bundle.HasAlpha()) return;
+      if (!bundle.HasColor()) return;
+      auto* color = bundle.color();
+      const auto* alpha = bundle.alpha();
+      JXL_CHECK(color->ysize() == alpha->ysize());
+      JXL_CHECK(color->xsize() == alpha->xsize());
+      for (size_t y = 0; y < color->ysize(); y++) {
+        ::jxl::UnpremultiplyAlpha(color->PlaneRow(0, y), color->PlaneRow(1, y),
+                                  color->PlaneRow(2, y), alpha->Row(y),
+                                  color->xsize());
+      }
+    };
+    ExtraChannelInfo* eci = metadata.m.Find(ExtraChannel::kAlpha);
+    if (eci == nullptr || !eci->alpha_associated) return false;
+    if (metadata.m.have_preview) {
+      doUnpremultiplyAlpha(preview_frame);
+    }
+    for (ImageBundle& ib : frames) {
+      doUnpremultiplyAlpha(ib);
+    }
+    eci->alpha_associated = false;
+    return true;
   }
 
   // -- DECODER INPUT:
@@ -170,7 +210,6 @@ class CodecInOut {
 
   std::vector<ImageBundle> frames;  // size=1 if !metadata.have_animation
 
-  bool use_sjpeg = false;
   // If the image should be written to a JPEG, use this quality for encoding.
   size_t jpeg_quality;
 };