1use std::ops::{Deref, DerefMut};
6
7use euclid::default::{Rect, Size2D};
8use image::codecs::jpeg::JpegEncoder;
9use image::codecs::png::PngEncoder;
10use image::codecs::webp::WebPEncoder;
11use image::{ExtendedColorType, GenericImageView, ImageEncoder, ImageError, Rgb};
12use ipc_channel::ipc::IpcSharedMemory;
13use malloc_size_of_derive::MallocSizeOf;
14use serde::{Deserialize, Serialize};
15
16use crate::{EncodedImageType, Multiply, rgba8_get_rect, transform_inplace};
17
18#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, MallocSizeOf, PartialEq, Serialize)]
19pub enum SnapshotPixelFormat {
20 #[default]
21 RGBA,
22 BGRA,
23}
24
25#[derive(Clone, Copy, Debug, Deserialize, Eq, MallocSizeOf, PartialEq, Serialize)]
26pub enum Alpha {
27 Premultiplied,
28 NotPremultiplied,
29 DontCare,
32}
33
34impl Alpha {
35 pub const fn from_premultiplied(is_premultiplied: bool) -> Self {
36 if is_premultiplied {
37 Self::Premultiplied
38 } else {
39 Self::NotPremultiplied
40 }
41 }
42
43 pub const fn needs_alpha_multiplication(&self) -> bool {
44 match self {
45 Alpha::Premultiplied => false,
46 Alpha::NotPremultiplied => true,
47 Alpha::DontCare => false,
48 }
49 }
50}
51
52#[derive(Clone, Copy, Debug, Deserialize, Eq, MallocSizeOf, PartialEq, Serialize)]
53pub enum SnapshotAlphaMode {
54 Opaque,
56 AsOpaque { premultiplied: bool },
58 Transparent { premultiplied: bool },
60}
61
62impl Default for SnapshotAlphaMode {
63 fn default() -> Self {
64 Self::Transparent {
65 premultiplied: true,
66 }
67 }
68}
69
70impl SnapshotAlphaMode {
71 pub const fn alpha(&self) -> Alpha {
72 match self {
73 SnapshotAlphaMode::Opaque => Alpha::DontCare,
74 SnapshotAlphaMode::AsOpaque { premultiplied } => {
75 Alpha::from_premultiplied(*premultiplied)
76 },
77 SnapshotAlphaMode::Transparent { premultiplied } => {
78 Alpha::from_premultiplied(*premultiplied)
79 },
80 }
81 }
82}
83
84#[derive(Clone, Debug, Deserialize, MallocSizeOf, Serialize)]
85pub enum SnapshotData {
86 Owned(Vec<u8>),
89}
90
91impl Deref for SnapshotData {
92 type Target = [u8];
93
94 fn deref(&self) -> &Self::Target {
95 match &self {
96 SnapshotData::Owned(items) => items,
98 }
99 }
100}
101
102impl DerefMut for SnapshotData {
103 fn deref_mut(&mut self) -> &mut Self::Target {
104 match self {
105 SnapshotData::Owned(items) => items,
107 }
108 }
109}
110
111pub type IpcSnapshot = Snapshot<IpcSharedMemory>;
112
113#[derive(Clone, Debug, Deserialize, MallocSizeOf, Serialize)]
121pub struct Snapshot<T = SnapshotData> {
122 size: Size2D<u32>,
123 data: T,
125 format: SnapshotPixelFormat,
127 alpha_mode: SnapshotAlphaMode,
129}
130
131impl<T> Snapshot<T> {
132 pub const fn size(&self) -> Size2D<u32> {
133 self.size
134 }
135
136 pub const fn format(&self) -> SnapshotPixelFormat {
137 self.format
138 }
139
140 pub const fn alpha_mode(&self) -> SnapshotAlphaMode {
141 self.alpha_mode
142 }
143}
144
145impl Snapshot<SnapshotData> {
146 pub fn empty() -> Self {
147 Self {
148 size: Size2D::zero(),
149 data: SnapshotData::Owned(vec![]),
150 format: SnapshotPixelFormat::RGBA,
151 alpha_mode: SnapshotAlphaMode::Transparent {
152 premultiplied: true,
153 },
154 }
155 }
156
157 pub fn cleared(size: Size2D<u32>) -> Self {
159 Self {
160 size,
161 data: SnapshotData::Owned(vec![0; size.area() as usize * 4]),
162 format: SnapshotPixelFormat::RGBA,
163 alpha_mode: SnapshotAlphaMode::Transparent {
164 premultiplied: true,
165 },
166 }
167 }
168
169 pub fn from_vec(
170 size: Size2D<u32>,
171 format: SnapshotPixelFormat,
172 alpha_mode: SnapshotAlphaMode,
173 data: Vec<u8>,
174 ) -> Self {
175 Self {
176 size,
177 data: SnapshotData::Owned(data),
178 format,
179 alpha_mode,
180 }
181 }
182
183 pub fn get_rect(&self, rect: Rect<u32>) -> Self {
184 let data = rgba8_get_rect(self.as_raw_bytes(), self.size(), rect).to_vec();
185 Self::from_vec(rect.size, self.format, self.alpha_mode, data)
186 }
187
188 pub fn transform(
212 &mut self,
213 target_alpha_mode: SnapshotAlphaMode,
214 target_format: SnapshotPixelFormat,
215 ) {
216 let swap_rb = target_format != self.format;
217 let multiply = match (self.alpha_mode, target_alpha_mode) {
218 (SnapshotAlphaMode::Opaque, _) => Multiply::None,
219 (alpha_mode, SnapshotAlphaMode::Opaque) => {
220 if alpha_mode.alpha() == Alpha::Premultiplied {
221 Multiply::UnMultiply
222 } else {
223 Multiply::None
224 }
225 },
226 (
227 SnapshotAlphaMode::Transparent { premultiplied } |
228 SnapshotAlphaMode::AsOpaque { premultiplied },
229 SnapshotAlphaMode::Transparent {
230 premultiplied: target_premultiplied,
231 } |
232 SnapshotAlphaMode::AsOpaque {
233 premultiplied: target_premultiplied,
234 },
235 ) => {
236 if premultiplied == target_premultiplied {
237 Multiply::None
238 } else if target_premultiplied {
239 Multiply::PreMultiply
240 } else {
241 Multiply::UnMultiply
242 }
243 },
244 };
245 let clear_alpha = !matches!(self.alpha_mode, SnapshotAlphaMode::Opaque) &&
246 matches!(target_alpha_mode, SnapshotAlphaMode::Opaque);
247 transform_inplace(self.data.deref_mut(), multiply, swap_rb, clear_alpha);
248 self.alpha_mode = target_alpha_mode;
249 self.format = target_format;
250 }
251
252 pub fn as_raw_bytes(&self) -> &[u8] {
253 &self.data
254 }
255
256 pub fn as_raw_bytes_mut(&mut self) -> &mut [u8] {
257 &mut self.data
258 }
259
260 pub fn as_bytes(
261 &mut self,
262 target_alpha_mode: Option<SnapshotAlphaMode>,
263 target_format: Option<SnapshotPixelFormat>,
264 ) -> (&mut [u8], SnapshotAlphaMode, SnapshotPixelFormat) {
265 let target_alpha_mode = target_alpha_mode.unwrap_or(self.alpha_mode);
266 let target_format = target_format.unwrap_or(self.format);
267 self.transform(target_alpha_mode, target_format);
268 (&mut self.data, target_alpha_mode, target_format)
269 }
270
271 pub fn to_vec(
272 mut self,
273 target_alpha_mode: Option<SnapshotAlphaMode>,
274 target_format: Option<SnapshotPixelFormat>,
275 ) -> (Vec<u8>, SnapshotAlphaMode, SnapshotPixelFormat) {
276 let target_alpha_mode = target_alpha_mode.unwrap_or(self.alpha_mode);
277 let target_format = target_format.unwrap_or(self.format);
278 self.transform(target_alpha_mode, target_format);
279 let SnapshotData::Owned(data) = self.data;
280 (data, target_alpha_mode, target_format)
281 }
282
283 pub fn as_ipc(self) -> Snapshot<IpcSharedMemory> {
284 let Snapshot {
285 size,
286 data,
287 format,
288 alpha_mode,
289 } = self;
290 let data = match data {
291 SnapshotData::Owned(items) => IpcSharedMemory::from_bytes(&items),
293 };
294 Snapshot {
295 size,
296 data,
297 format,
298 alpha_mode,
299 }
300 }
301
302 pub fn encode_for_mime_type<W: std::io::Write>(
303 &mut self,
304 image_type: &EncodedImageType,
305 quality: Option<f64>,
306 encoder: &mut W,
307 ) -> Result<(), ImageError> {
308 let width = self.size.width;
309 let height = self.size.height;
310
311 let (data, _, _) = self.as_bytes(
312 if *image_type == EncodedImageType::Jpeg {
313 Some(SnapshotAlphaMode::AsOpaque {
314 premultiplied: true,
315 })
316 } else {
317 Some(SnapshotAlphaMode::Transparent {
318 premultiplied: false,
319 })
320 },
321 Some(SnapshotPixelFormat::RGBA),
322 );
323
324 match image_type {
325 EncodedImageType::Png => {
326 PngEncoder::new(encoder).write_image(data, width, height, ExtendedColorType::Rgba8)
329 },
330 EncodedImageType::Jpeg => {
331 let mut jpeg_encoder = if let Some(quality) = quality {
332 if (0.0..=1.0).contains(&quality) {
335 JpegEncoder::new_with_quality(
336 encoder,
337 (quality * 100.0).round().clamp(1.0, 100.0) as u8,
338 )
339 } else {
340 JpegEncoder::new(encoder)
341 }
342 } else {
343 JpegEncoder::new(encoder)
344 };
345
346 struct RgbaDataForJpegEncoder<'a> {
349 width: u32,
350 height: u32,
351 data: &'a [u8],
352 }
353
354 impl<'a> GenericImageView for RgbaDataForJpegEncoder<'a> {
355 type Pixel = Rgb<u8>;
356
357 fn dimensions(&self) -> (u32, u32) {
358 (self.width, self.height)
359 }
360
361 fn get_pixel(&self, x: u32, y: u32) -> Self::Pixel {
362 let offset = (self.width * y + x) as usize * 4;
363 Rgb([
364 self.data[offset],
365 self.data[offset + 1],
366 self.data[offset + 2],
367 ])
368 }
369 }
370
371 let image = RgbaDataForJpegEncoder {
372 width,
373 height,
374 data,
375 };
376
377 jpeg_encoder.encode_image(&image)
378 },
379 EncodedImageType::Webp => {
380 WebPEncoder::new_lossless(encoder).write_image(
382 data,
383 width,
384 height,
385 ExtendedColorType::Rgba8,
386 )
387 },
388 }
389 }
390}
391
392impl Snapshot<IpcSharedMemory> {
393 pub fn to_owned(self) -> Snapshot<SnapshotData> {
415 let Snapshot {
416 size,
417 data,
418 format,
419 alpha_mode,
420 } = self;
421 Snapshot {
422 size,
423 data: SnapshotData::Owned(data.to_vec()),
424 format,
425 alpha_mode,
426 }
427 }
428
429 pub fn data(&self) -> &[u8] {
430 &self.data
431 }
432
433 pub fn to_ipc_shared_memory(self) -> IpcSharedMemory {
434 self.data
435 }
436}