Export alpha channel into PDF 🧊

7 files changed, 65 insertions, 25 deletions

diff --git a/src/export/pdf.rs b/src/export/pdf.rs
index fcd451d1..14bc11c3 100644
--- a/src/export/pdf.rs
+++ b/src/export/pdf.rs
@@ -5,7 +5,7 @@ use std::collections::HashMap;
 use std::hash::Hash;
 
 use fontdock::FaceId;
-use image::{DynamicImage, GenericImageView};
+use image::{DynamicImage, GenericImageView, Rgba};
 use pdf_writer::{
     CidFontType, ColorSpace, Content, FontFlags, Name, PdfWriter, Rect, Ref, Str,
     SystemInfo, UnicodeCmap,
@@ -43,17 +43,24 @@ impl<'a> PdfExporter<'a> {
 
         let mut fonts = Remapper::new();
         let mut images = Remapper::new();
+        let mut alpha_masks = 0;
 
         for layout in layouts {
             for (_, element) in &layout.elements {
                 match element {
                     LayoutElement::Text(shaped) => fonts.insert(shaped.face),
-                    LayoutElement::Image(image) => images.insert(image.resource),
+                    LayoutElement::Image(image) => {
+                        let buf = env.resources.get_loaded::<DynamicImage>(image.res);
+                        if buf.color().has_alpha() {
+                            alpha_masks += 1;
+                        }
+                        images.insert(image.res);
+                    }
                 }
             }
         }
 
-        let refs = Refs::new(layouts.len(), fonts.len(), images.len());
+        let refs = Refs::new(layouts.len(), fonts.len(), images.len(), alpha_masks);
 
         Self {
             writer,
@@ -154,7 +161,7 @@ impl<'a> PdfExporter<'a> {
 
         for (pos, element) in &page.elements {
             if let LayoutElement::Image(image) = element {
-                let name = format!("Im{}", self.images.map(image.resource));
+                let name = format!("Im{}", self.images.map(image.res));
                 let size = image.size;
                 let x = pos.x.to_pt() as f32;
                 let y = (page.size.height - pos.y - size.height).to_pt() as f32;
@@ -281,15 +288,40 @@ impl<'a> PdfExporter<'a> {
     }
 
     fn write_images(&mut self) {
+        let mut mask = 0;
+
         for (id, resource) in self.refs.images().zip(self.images.layout_indices()) {
-            let image = self.env.resources.get_loaded::<DynamicImage>(resource);
-            let data = image.to_rgb8().into_raw();
-            self.writer
-                .image_stream(id, &data)
-                .width(image.width() as i32)
-                .height(image.height() as i32)
-                .color_space(ColorSpace::DeviceRGB)
-                .bits_per_component(8);
+            let buf = self.env.resources.get_loaded::<DynamicImage>(resource);
+            let data = buf.to_rgb8().into_raw();
+
+            let mut image = self.writer.image_stream(id, &data);
+            image.width(buf.width() as i32);
+            image.height(buf.height() as i32);
+            image.color_space(ColorSpace::DeviceRGB);
+            image.bits_per_component(8);
+
+            // Add a second gray-scale image containing the alpha values if this
+            // is image has an alpha channel.
+            if buf.color().has_alpha() {
+                let mask_id = self.refs.alpha_mask(mask);
+
+                image.s_mask(mask_id);
+                drop(image);
+
+                let mut samples = vec![];
+                for (_, _, Rgba([_, _, _, a])) in buf.pixels() {
+                    samples.push(a);
+                }
+
+                self.writer
+                    .image_stream(mask_id, &samples)
+                    .width(buf.width() as i32)
+                    .height(buf.height() as i32)
+                    .color_space(ColorSpace::DeviceGray)
+                    .bits_per_component(8);
+
+                mask += 1;
+            }
         }
     }
 }
@@ -304,6 +336,7 @@ struct Refs {
     contents_start: i32,
     fonts_start: i32,
     images_start: i32,
+    alpha_masks_start: i32,
     end: i32,
 }
 
@@ -318,14 +351,15 @@ struct FontRefs {
 impl Refs {
     const OBJECTS_PER_FONT: usize = 5;
 
-    fn new(layouts: usize, fonts: usize, images: usize) -> Self {
+    fn new(layouts: usize, fonts: usize, images: usize, alpha_masks: usize) -> Self {
         let catalog = 1;
         let page_tree = catalog + 1;
         let pages_start = page_tree + 1;
         let contents_start = pages_start + layouts as i32;
         let fonts_start = contents_start + layouts as i32;
         let images_start = fonts_start + (Self::OBJECTS_PER_FONT * fonts) as i32;
-        let end = images_start + images as i32;
+        let alpha_masks_start = images_start + images as i32;
+        let end = alpha_masks_start + alpha_masks as i32;
 
         Self {
             catalog: Ref::new(catalog),
@@ -334,6 +368,7 @@ impl Refs {
             contents_start,
             fonts_start,
             images_start,
+            alpha_masks_start,
             end,
         }
     }
@@ -361,6 +396,10 @@ impl Refs {
     fn images(&self) -> impl Iterator<Item = Ref> {
         (self.images_start .. self.end).map(Ref::new)
     }
+
+    fn alpha_mask(&self, i: usize) -> Ref {
+        Ref::new(self.alpha_masks_start + i as i32)
+    }
 }
 
 /// Used to assign new, consecutive PDF-internal indices to things.
diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index 1f7d6c9f..44a2c2fa 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -183,7 +183,7 @@ pub enum LayoutElement {
 #[derive(Debug, Clone, PartialEq)]
 pub struct ImageElement {
     /// The image.
-    pub resource: ResourceId,
+    pub res: ResourceId,
     /// The document size of the image.
     pub size: Size,
 }
diff --git a/src/library/insert.rs b/src/library/insert.rs
index b2cdc255..c196ce9e 100644
--- a/src/library/insert.rs
+++ b/src/library/insert.rs
@@ -27,11 +27,11 @@ pub fn image(mut args: Args, ctx: &mut EvalContext) -> Value {
                 .and_then(|reader| reader.decode().ok())
         });
 
-        if let Some((resource, buf)) = loaded {
+        if let Some((res, buf)) = loaded {
             let dimensions = buf.dimensions();
             drop(env);
             ctx.push(Image {
-                resource,
+                res,
                 dimensions,
                 width,
                 height,
@@ -50,7 +50,7 @@ pub fn image(mut args: Args, ctx: &mut EvalContext) -> Value {
 #[derive(Debug, Clone, PartialEq)]
 struct Image {
     /// The resource id of the image file.
-    resource: ResourceId,
+    res: ResourceId,
     /// The pixel dimensions of the image.
     dimensions: (u32, u32),
     /// The fixed width, if any.
@@ -87,7 +87,7 @@ impl Layout for Image {
         let mut boxed = BoxLayout::new(size);
         boxed.push(
             Point::ZERO,
-            LayoutElement::Image(ImageElement { resource: self.resource, size }),
+            LayoutElement::Image(ImageElement { res: self.res, size }),
         );
 
         Layouted::Layout(boxed, self.align)
diff --git a/tests/ref/image.png b/tests/ref/image.png
index 41fa6a9b..37349600 100644
--- a/tests/ref/image.png
+++ b/tests/ref/image.png
diff --git a/tests/res/tiger-alpha.png b/tests/res/tiger-alpha.png
new file mode 100644
index 00000000..178b4b65
--- /dev/null
+++ b/tests/res/tiger-alpha.png
diff --git a/tests/typ/image.typ b/tests/typ/image.typ
index ee82024f..0dd22950 100644
--- a/tests/typ/image.typ
+++ b/tests/typ/image.typ
@@ -1,14 +1,15 @@
-[page: width=10cm, height=10cm, margins=1cm]
+[page: width=5cm, height=5cm, margins=0.25cm]
 
 [image: "res/tiger.jpg"]
 
 [pagebreak]
 
 # Tiger
-[image: "res/tiger.jpg", width=3cm]
-[image: "res/tiger.jpg", height=3cm]
+[image: "res/tiger.jpg", width=2cm]
+[image: "res/tiger-alpha.png", width=1cm]
+[image: "res/tiger-alpha.png", height=2cm]
 
 [pagebreak]
 
-[align: center]
-[image: "res/tiger.jpg", width=6cm, height=6cm]
+[align: center, bottom]
+[image: "res/tiger.jpg", width=2cm, height=3.5cm]
diff --git a/tests/typeset.rs b/tests/typeset.rs
index 82d801a2..cc146fe4 100644
--- a/tests/typeset.rs
+++ b/tests/typeset.rs
@@ -248,7 +248,7 @@ fn draw_text(canvas: &mut Canvas, pos: Point, env: &Env, shaped: &Shaped) {
 }
 
 fn draw_image(canvas: &mut Canvas, pos: Point, env: &Env, image: &ImageElement) {
-    let buf = env.resources.get_loaded::<DynamicImage>(image.resource);
+    let buf = env.resources.get_loaded::<DynamicImage>(image.res);
 
     let mut pixmap = Pixmap::new(buf.width(), buf.height()).unwrap();
     for ((_, _, src), dest) in buf.pixels().zip(pixmap.pixels_mut()) {