pixels: Ensure expected formats when accesing bytes of snapshot (#37767)

I introduced snapshot in #36119 to pack raw bytes and metadata together,
now we take the next step and require for user to always specify what
kind of byte data they want when calling `as_bytes` or `to_vec`
(basically joining transform and data). There are also valid usages when
one might require just one property of bytes (textures can generally
handle both RGBA and BGRA). There are also valid usages of using just
raw bytes (when cropping). This PR tries to make such usages more
obvious.

This will make it easier to fix stuff around 2d canvas (we do not want
to assume any bytes properties in abstraction).

Testing: Code is covered by WPT tests.

---------

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>
Co-authored-by: Martin Robinson <mrobinson@igalia.com>

components/canvas/canvas_paint_thread.rs

@@ -17,7 +17,7 @@ use ipc_channel::ipc::{self, IpcSender};
 use ipc_channel::router::ROUTER;
 use log::warn;
 use net_traits::ResourceThreads;
-use pixels::Snapshot;
+use pixels::{Snapshot, SnapshotPixelFormat};
 use style::color::AbsoluteColor;
 use style::properties::style_structs::Font as FontStyleStruct;
 use webrender_api::ImageKey;
@@ -175,14 +175,17 @@ impl<'a> CanvasPaintThread<'a> {
                 .canvas(canvas_id)
                 .is_point_in_path_(&path[..], x, y, fill_rule, chan),
             Canvas2dMsg::DrawImage(snapshot, dest_rect, source_rect, smoothing_enabled) => {
-                let snapshot = snapshot.to_owned();
+                let mut snapshot = snapshot.to_owned();
+                let size = snapshot.size();
+                let (data, alpha_mode, _) =
+                    snapshot.as_bytes(None, Some(SnapshotPixelFormat::BGRA));
                 self.canvas(canvas_id).draw_image(
-                    snapshot.data(),
+                    data,
-                    snapshot.size(),
+                    size,
                     dest_rect,
                     source_rect,
                     smoothing_enabled,
-                    !snapshot.alpha_mode().is_premultiplied(),
+                    alpha_mode.alpha().needs_alpha_multiplication(),
                 )
             },
             Canvas2dMsg::DrawEmptyImage(image_size, dest_rect, source_rect) => {
@@ -202,16 +205,18 @@ impl<'a> CanvasPaintThread<'a> {
                 source_rect,
                 smoothing,
             ) => {
-                let image_data = self
+                let mut snapshot = self
                     .canvas(canvas_id)
                     .read_pixels(Some(source_rect.to_u32()), Some(image_size));
+                let (data, alpha_mode, _) =
+                    snapshot.as_bytes(None, Some(SnapshotPixelFormat::BGRA));
                 self.canvas(other_canvas_id).draw_image(
-                    image_data.data(),
+                    data,
                     source_rect.size.to_u32(),
                     dest_rect,
                     source_rect,
                     smoothing,
-                    false,
+                    alpha_mode.alpha().needs_alpha_multiplication(),
                 );
             },
             Canvas2dMsg::MoveTo(ref point) => self.canvas(canvas_id).move_to(point),
@@ -403,7 +408,7 @@ impl Canvas<'_> {
         dest_rect: Rect<f64>,
         source_rect: Rect<f64>,
         smoothing_enabled: bool,
-        is_premultiplied: bool,
+        premultiply: bool,
     ) {
         match self {
             Canvas::Raqote(canvas_data) => canvas_data.draw_image(
@@ -412,7 +417,7 @@ impl Canvas<'_> {
                 dest_rect,
                 source_rect,
                 smoothing_enabled,
-                is_premultiplied,
+                premultiply,
             ),
         }
     }

components/pixels/snapshot.rs

@@ -18,6 +18,33 @@ pub enum SnapshotPixelFormat {
     BGRA,
 }
 
+#[derive(Clone, Copy, Debug, Deserialize, Eq, MallocSizeOf, PartialEq, Serialize)]
+pub enum Alpha {
+    Premultiplied,
+    NotPremultiplied,
+    /// This is used for opaque textures for which the presence of alpha in the
+    /// output data format does not matter.
+    DontCare,
+}
+
+impl Alpha {
+    pub const fn from_premultiplied(is_premultiplied: bool) -> Self {
+        if is_premultiplied {
+            Self::Premultiplied
+        } else {
+            Self::NotPremultiplied
+        }
+    }
+
+    pub const fn needs_alpha_multiplication(&self) -> bool {
+        match self {
+            Alpha::Premultiplied => false,
+            Alpha::NotPremultiplied => true,
+            Alpha::DontCare => false,
+        }
+    }
+}
+
 #[derive(Clone, Copy, Debug, Deserialize, Eq, MallocSizeOf, PartialEq, Serialize)]
 pub enum SnapshotAlphaMode {
     /// Internal data is opaque (alpha is cleared to 1)
@@ -37,20 +64,17 @@ impl Default for SnapshotAlphaMode {
 }
 
 impl SnapshotAlphaMode {
-    pub const fn is_premultiplied(&self) -> bool {
+    pub const fn alpha(&self) -> Alpha {
         match self {
-            SnapshotAlphaMode::Opaque => true,
+            SnapshotAlphaMode::Opaque => Alpha::DontCare,
-            SnapshotAlphaMode::AsOpaque { premultiplied } => *premultiplied,
+            SnapshotAlphaMode::AsOpaque { premultiplied } => {
-            SnapshotAlphaMode::Transparent { premultiplied } => *premultiplied,
+                Alpha::from_premultiplied(*premultiplied)
+            },
+            SnapshotAlphaMode::Transparent { premultiplied } => {
+                Alpha::from_premultiplied(*premultiplied)
+            },
         }
     }
-
-    pub const fn is_opaque(&self) -> bool {
-        matches!(
-            self,
-            SnapshotAlphaMode::Opaque | SnapshotAlphaMode::AsOpaque { .. }
-        )
-    }
 }
 
 #[derive(Clone, Debug, Deserialize, MallocSizeOf, Serialize)]
@@ -112,14 +136,6 @@ impl<T> Snapshot<T> {
     pub const fn alpha_mode(&self) -> SnapshotAlphaMode {
         self.alpha_mode
     }
-
-    pub const fn is_premultiplied(&self) -> bool {
-        self.alpha_mode().is_premultiplied()
-    }
-
-    pub const fn is_opaque(&self) -> bool {
-        self.alpha_mode().is_opaque()
-    }
 }
 
 impl Snapshot<SnapshotData> {
@@ -181,14 +197,6 @@ impl Snapshot<SnapshotData> {
     }
     */
 
-    pub fn data(&self) -> &[u8] {
-        &self.data
-    }
-
-    pub fn data_mut(&mut self) -> &mut [u8] {
-        &mut self.data
-    }
-
     /// Convert inner data of snapshot to target format and alpha mode.
     /// If data is already in target format and alpha mode no work will be done.
     pub fn transform(
@@ -200,7 +208,7 @@ impl Snapshot<SnapshotData> {
         let multiply = match (self.alpha_mode, target_alpha_mode) {
             (SnapshotAlphaMode::Opaque, _) => Multiply::None,
             (alpha_mode, SnapshotAlphaMode::Opaque) => {
-                if alpha_mode.is_premultiplied() {
+                if alpha_mode.alpha() == Alpha::Premultiplied {
                     Multiply::UnMultiply
                 } else {
                     Multiply::None
@@ -232,6 +240,37 @@ impl Snapshot<SnapshotData> {
         self.format = target_format;
     }
 
+    pub fn as_raw_bytes(&self) -> &[u8] {
+        &self.data
+    }
+
+    pub fn as_raw_bytes_mut(&mut self) -> &mut [u8] {
+        &mut self.data
+    }
+
+    pub fn as_bytes(
+        &mut self,
+        target_alpha_mode: Option<SnapshotAlphaMode>,
+        target_format: Option<SnapshotPixelFormat>,
+    ) -> (&mut [u8], SnapshotAlphaMode, SnapshotPixelFormat) {
+        let target_alpha_mode = target_alpha_mode.unwrap_or(self.alpha_mode);
+        let target_format = target_format.unwrap_or(self.format);
+        self.transform(target_alpha_mode, target_format);
+        (&mut self.data, target_alpha_mode, target_format)
+    }
+
+    pub fn to_vec(
+        mut self,
+        target_alpha_mode: Option<SnapshotAlphaMode>,
+        target_format: Option<SnapshotPixelFormat>,
+    ) -> (Vec<u8>, SnapshotAlphaMode, SnapshotPixelFormat) {
+        let target_alpha_mode = target_alpha_mode.unwrap_or(self.alpha_mode);
+        let target_format = target_format.unwrap_or(self.format);
+        self.transform(target_alpha_mode, target_format);
+        let SnapshotData::Owned(data) = self.data;
+        (data, target_alpha_mode, target_format)
+    }
+
     pub fn as_ipc(self) -> Snapshot<IpcSharedMemory> {
         let Snapshot {
             size,
@@ -250,12 +289,6 @@ impl Snapshot<SnapshotData> {
             alpha_mode,
         }
     }
-
-    pub fn to_vec(self) -> Vec<u8> {
-        match self.data {
-            SnapshotData::Owned(data) => data,
-        }
-    }
 }
 
 impl Snapshot<IpcSharedMemory> {

components/script/canvas_state.rs

@@ -394,14 +394,15 @@ impl CanvasState {
 
         let (sender, receiver) = ipc::channel().unwrap();
         self.send_canvas_2d_msg(Canvas2dMsg::GetImageData(rect, canvas_size, sender));
-        let mut snapshot = receiver.recv().unwrap().to_owned();
+        let snapshot = receiver.recv().unwrap().to_owned();
-        snapshot.transform(
+        snapshot
-            SnapshotAlphaMode::Transparent {
+            .to_vec(
+                Some(SnapshotAlphaMode::Transparent {
                     premultiplied: false,
-            },
+                }),
-            SnapshotPixelFormat::RGBA,
+                Some(SnapshotPixelFormat::RGBA),
-        );
+            )
-        snapshot.to_vec()
+            .0
     }
 
     ///
@@ -1177,7 +1178,14 @@ impl CanvasState {
             let size = snapshot.size();
             Ok(Some(CanvasPattern::new(
                 global,
-                snapshot.to_vec(),
+                snapshot
+                    .to_vec(
+                        Some(SnapshotAlphaMode::Transparent {
+                            premultiplied: true,
+                        }),
+                        Some(SnapshotPixelFormat::BGRA),
+                    )
+                    .0, // TODO: send snapshot
                 size.cast(),
                 rep,
                 self.is_origin_clean(image),

components/script/dom/htmlcanvaselement.rs

@@ -390,16 +390,27 @@ impl HTMLCanvasElement {
         &self,
         image_type: &EncodedImageType,
         quality: Option<f64>,
-        snapshot: &Snapshot,
+        mut snapshot: Snapshot,
         encoder: &mut W,
     ) -> Result<(), ImageError> {
         // We can't use self.Width() or self.Height() here, since the size of the canvas
         // may have changed since the snapshot was created. Truncating the dimensions to a
         // u32 can't panic, since the data comes from a canvas which is always smaller than
         // u32::MAX.
-        let canvas_data = snapshot.data();
         let width = snapshot.size().width;
         let height = snapshot.size().height;
+        let (canvas_data, _, _) = snapshot.as_bytes(
+            if *image_type == EncodedImageType::Jpeg {
+                Some(SnapshotAlphaMode::AsOpaque {
+                    premultiplied: true,
+                })
+            } else {
+                Some(SnapshotAlphaMode::Transparent {
+                    premultiplied: false,
+                })
+            },
+            Some(SnapshotPixelFormat::RGBA),
+        );
 
         match image_type {
             EncodedImageType::Png => {
@@ -537,23 +548,12 @@ impl HTMLCanvasElementMethods<crate::DomTypeHolder> for HTMLCanvasElement {
 
         // Step 3: Let file be a serialization of this canvas element's bitmap as a file,
         // passing type and quality if given.
-        let Some(mut snapshot) = self.get_image_data() else {
+        let Some(snapshot) = self.get_image_data() else {
             return Ok(USVString("data:,".into()));
         };
 
         let image_type = EncodedImageType::from(mime_type);
-        snapshot.transform(
+
-            if image_type == EncodedImageType::Jpeg {
-                SnapshotAlphaMode::AsOpaque {
-                    premultiplied: true,
-                }
-            } else {
-                SnapshotAlphaMode::Transparent {
-                    premultiplied: false,
-                }
-            },
-            SnapshotPixelFormat::RGBA,
-        );
         let mut url = format!("data:{};base64,", image_type.as_mime_type());
 
         let mut encoder = base64::write::EncoderStringWriter::from_consumer(
@@ -565,7 +565,7 @@ impl HTMLCanvasElementMethods<crate::DomTypeHolder> for HTMLCanvasElement {
             .encode_for_mime_type(
                 &image_type,
                 Self::maybe_quality(quality),
-                &snapshot,
+                snapshot,
                 &mut encoder,
             )
             .is_err()
@@ -622,16 +622,11 @@ impl HTMLCanvasElementMethods<crate::DomTypeHolder> for HTMLCanvasElement {
                     return error!("Expected blob callback, but found none!");
                 };
 
-                let Some(mut snapshot) = result else {
+                let Some(snapshot) = result else {
                     let _ = callback.Call__(None, ExceptionHandling::Report, CanGc::note());
                     return;
                 };
 
-                snapshot.transform(
-                    SnapshotAlphaMode::Transparent { premultiplied: false },
-                    SnapshotPixelFormat::RGBA
-                );
-
                 // Step 4.1: If result is non-null, then set result to a serialization of
                 // result as a file with type and quality if given.
                 // Step 4.2: Queue an element task on the canvas blob serialization task
@@ -639,7 +634,7 @@ impl HTMLCanvasElementMethods<crate::DomTypeHolder> for HTMLCanvasElement {
                 let mut encoded: Vec<u8> = vec![];
                 let blob_impl;
                 let blob;
-                let result = match this.encode_for_mime_type(&image_type, quality, &snapshot, &mut encoded) {
+                let result = match this.encode_for_mime_type(&image_type, quality, snapshot, &mut encoded) {
                    Ok(..) => {
                        // Step 4.2.1: If result is non-null, then set result to a new Blob
                        // object, created in the relevant realm of this canvas element,

components/script/dom/imagebitmap.rs

@@ -198,10 +198,10 @@ impl ImageBitmap {
         pixels::copy_rgba8_image(
             input.size(),
             input_rect_cropped.cast(),
-            input.data(),
+            input.as_raw_bytes(),
             source.size(),
             source_rect_cropped.cast(),
-            source.data_mut(),
+            source.as_raw_bytes_mut(),
         );
 
         // Step 7. Scale output to the size specified by outputWidth and outputHeight.
@@ -213,9 +213,12 @@ impl ImageBitmap {
                 ResizeQuality::High => pixels::FilterQuality::High,
             };
 
-            let Some(output_data) =
+            let Some(output_data) = pixels::scale_rgba8_image(
-                pixels::scale_rgba8_image(source.size(), source.data(), output_size, quality)
+                source.size(),
-            else {
+                source.as_raw_bytes(),
+                output_size,
+                quality,
+            ) else {
                 log::warn!(
                     "Failed to scale the bitmap of size {:?} to required size {:?}",
                     source.size(),
@@ -240,7 +243,7 @@ impl ImageBitmap {
         // output must be flipped vertically, disregarding any image orientation metadata
         // of the source (such as EXIF metadata), if any.
         if options.imageOrientation == ImageOrientation::FlipY {
-            pixels::flip_y_rgba8_image_inplace(output.size(), output.data_mut());
+            pixels::flip_y_rgba8_image_inplace(output.size(), output.as_raw_bytes_mut());
         }
 
         // TODO: Step 9. If image is an img element or a Blob object, let val be the value

components/script/dom/webgl2renderingcontext.rs

@@ -21,7 +21,7 @@ use js::jsapi::{JSObject, Type};
 use js::jsval::{BooleanValue, DoubleValue, Int32Value, NullValue, ObjectValue, UInt32Value};
 use js::rust::{CustomAutoRooterGuard, HandleObject, MutableHandleValue};
 use js::typedarray::{ArrayBufferView, CreateWith, Float32, Int32Array, Uint32, Uint32Array};
-use pixels::Snapshot;
+use pixels::{Alpha, Snapshot};
 use script_bindings::interfaces::WebGL2RenderingContextHelpers;
 use servo_config::pref;
 use url::Host;
@@ -3257,7 +3257,8 @@ impl WebGL2RenderingContextMethods<crate::DomTypeHolder> for WebGL2RenderingCont
 
         let size = Size2D::new(width, height);
 
-        let (alpha_treatment, y_axis_treatment) = self.base.get_current_unpack_state(false);
+        let (alpha_treatment, y_axis_treatment) =
+            self.base.get_current_unpack_state(Alpha::NotPremultiplied);
 
         self.base.tex_image_2d(
             &texture,

components/script/dom/webglrenderingcontext.rs

@@ -29,7 +29,7 @@ use js::typedarray::{
     ArrayBufferView, CreateWith, Float32, Float32Array, Int32, Int32Array, TypedArray,
     TypedArrayElementCreator, Uint32Array,
 };
-use pixels::{self, PixelFormat, Snapshot, SnapshotPixelFormat};
+use pixels::{self, Alpha, PixelFormat, Snapshot, SnapshotPixelFormat};
 use serde::{Deserialize, Serialize};
 use servo_config::pref;
 use webrender_api::ImageKey;
@@ -507,14 +507,14 @@ impl WebGLRenderingContext {
 
     pub(crate) fn get_current_unpack_state(
         &self,
-        premultiplied: bool,
+        premultiplied: Alpha,
     ) -> (Option<AlphaTreatment>, YAxisTreatment) {
         let settings = self.texture_unpacking_settings.get();
         let dest_premultiplied = settings.contains(TextureUnpacking::PREMULTIPLY_ALPHA);
 
         let alpha_treatment = match (premultiplied, dest_premultiplied) {
-            (true, false) => Some(AlphaTreatment::Unmultiply),
+            (Alpha::Premultiplied, false) => Some(AlphaTreatment::Unmultiply),
-            (false, true) => Some(AlphaTreatment::Premultiply),
+            (Alpha::NotPremultiplied, true) => Some(AlphaTreatment::Premultiply),
             _ => None,
         };
 
@@ -626,7 +626,8 @@ impl WebGLRenderingContext {
                 )
             },
             TexImageSource::ImageData(image_data) => {
-                let (alpha_treatment, y_axis_treatment) = self.get_current_unpack_state(false);
+                let (alpha_treatment, y_axis_treatment) =
+                    self.get_current_unpack_state(Alpha::NotPremultiplied);
 
                 TexPixels::new(
                     image_data.to_shared_memory(),
@@ -665,7 +666,7 @@ impl WebGLRenderingContext {
                 };
 
                 let (alpha_treatment, y_axis_treatment) =
-                    self.get_current_unpack_state(snapshot.alpha_mode().is_premultiplied());
+                    self.get_current_unpack_state(Alpha::NotPremultiplied);
 
                 TexPixels::new(
                     snapshot.to_ipc_shared_memory(),
@@ -695,7 +696,7 @@ impl WebGLRenderingContext {
                 };
 
                 let (alpha_treatment, y_axis_treatment) =
-                    self.get_current_unpack_state(snapshot.alpha_mode().is_premultiplied());
+                    self.get_current_unpack_state(snapshot.alpha_mode().alpha());
 
                 TexPixels::new(
                     snapshot.to_ipc_shared_memory(),
@@ -722,7 +723,7 @@ impl WebGLRenderingContext {
                 };
 
                 let (alpha_treatment, y_axis_treatment) =
-                    self.get_current_unpack_state(snapshot.alpha_mode().is_premultiplied());
+                    self.get_current_unpack_state(snapshot.alpha_mode().alpha());
 
                 TexPixels::new(
                     snapshot.to_ipc_shared_memory(),
@@ -4539,7 +4540,8 @@ impl WebGLRenderingContextMethods<crate::DomTypeHolder> for WebGLRenderingContex
 
         let size = Size2D::new(width, height);
 
-        let (alpha_treatment, y_axis_treatment) = self.get_current_unpack_state(false);
+        let (alpha_treatment, y_axis_treatment) =
+            self.get_current_unpack_state(Alpha::NotPremultiplied);
 
         self.tex_image_2d(
             &texture,
@@ -4716,7 +4718,8 @@ impl WebGLRenderingContextMethods<crate::DomTypeHolder> for WebGLRenderingContex
             };
         }
 
-        let (alpha_treatment, y_axis_treatment) = self.get_current_unpack_state(false);
+        let (alpha_treatment, y_axis_treatment) =
+            self.get_current_unpack_state(Alpha::NotPremultiplied);
 
         self.tex_sub_image_2d(
             texture,