Cryptomatte: define CryptoAux, which the previous commit used but never added (#473)

crypto_aux_image_with now takes CryptoAux::{ColorsAndMatte, MatteOnly}
instead of a bool, and CryptoAux::for_display holds the rule for which
displays read the ID colours (Colors without Matte Only). IdColorMemo uses
get_mut instead of indexing, and the CPU effect zips the pixels with the
aux index map instead of indexing it. No behaviour change.

crates/effects/src/channel3d.rs

@@ -408,14 +408,19 @@ struct IdColorMemo(Vec<(u32, [f32; 3])>);
 impl IdColorMemo {
     #[inline]
     fn get(&mut self, rank: usize, id: f32) -> [f32; 3] {
-        if self.0.len() <= rank {
-            self.0.resize(rank + 1, (id.to_bits(), id_color(id)));
+        let bits = id.to_bits();
+        match self.0.get_mut(rank) {
+            Some(e) if e.0 == bits => e.1,
+            Some(e) => {
+                *e = (bits, id_color(id));
+                e.1
+            }
+            None => {
+                let c = id_color(id);
+                self.0.resize(rank + 1, (bits, c));
+                c
+            }
         }
-        let e = &mut self.0[rank];
-        if e.0 != id.to_bits() {
-            *e = (id.to_bits(), id_color(id));
-        }
-        e.1
     }
 }
 
@@ -437,16 +442,37 @@ fn crypto_colors(planes: &RankPlanes, i: usize, memo: &mut IdColorMemo) -> [f32;
 /// selection's coverage (A), computed in parallel with the rank planes resolved once. The GPU
 /// effect uploads this image and samples it per layer pixel.
 pub fn crypto_aux_image(a: &AuxChannels, layer: &str, sel: &[u32]) -> Image {
-    crypto_aux_image_with(a, layer, sel, true)
+    crypto_aux_image_with(a, layer, sel, CryptoAux::ColorsAndMatte)
 }
 
-/// [`crypto_aux_image`]; with `colors` false the RGB stays 0 and only the matte (A) is computed,
-/// for the displays that never read the ID colours.
-pub fn crypto_aux_image_with(a: &AuxChannels, layer: &str, sel: &[u32], colors: bool) -> Image {
+/// What [`crypto_aux_image_with`] computes per aux pixel.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum CryptoAux {
+    /// The coverage-weighted ID colours (RGB) and the selection's coverage (A).
+    ColorsAndMatte,
+    /// Only the selection's coverage (A); the RGB stays 0.
+    MatteOnly,
+}
+
+impl CryptoAux {
+    /// What the Cryptomatte effect reads for `display` and Matte Only: the ID colours are only
+    /// read by the Colors display (0) without Matte Only; Matte, Preview and Matte Only need the
+    /// matte alone.
+    pub fn for_display(display: u32, matte_only: bool) -> Self {
+        if display == 0 && !matte_only { Self::ColorsAndMatte } else { Self::MatteOnly }
+    }
+}
+
+/// [`crypto_aux_image`]; with [`CryptoAux::MatteOnly`] the RGB stays 0 and only the matte (A) is
+/// computed, for the displays that never read the ID colours.
+pub fn crypto_aux_image_with(a: &AuxChannels, layer: &str, sel: &[u32], what: CryptoAux) -> Image {
     let planes = a.crypto_rank_planes(layer);
     let mut img = Image::new(a.width, a.height);
     img.data.par_iter_mut().enumerate().for_each_init(IdColorMemo::default, |memo, (i, px)| {
-        let c = if colors { crypto_colors(&planes, i, memo) } else { [0.0; 3] };
+        let c = match what {
+            CryptoAux::ColorsAndMatte => crypto_colors(&planes, i, memo),
+            CryptoAux::MatteOnly => [0.0; 3],
+        };
         *px = [c[0], c[1], c[2], crypto_matte(&planes, sel, i)];
     });
     img
@@ -466,8 +492,8 @@ fn cryptomatte(ctx: &EffectCtx, mut b: Buf) -> Buf {
     let display = ctx.params.e("display");
     let matte_only = ctx.params.b("matteOnly");
     let planes = a.crypto_rank_planes(&layer);
-    b.img.data.par_iter_mut().enumerate().for_each_init(IdColorMemo::default, |memo, (i, px)| {
-        let m = idx[i].map_or(0.0, |k| crypto_matte(&planes, &sel, k));
+    b.img.data.par_iter_mut().zip(idx.par_iter()).for_each_init(IdColorMemo::default, |memo, (px, &k)| {
+        let m = k.map_or(0.0, |k| crypto_matte(&planes, &sel, k));
         if matte_only {
             *px = [m, m, m, m];
             return;
@@ -475,7 +501,7 @@ fn cryptomatte(ctx: &EffectCtx, mut b: Buf) -> Buf {
         *px = match display {
             // Colors: coverage-weighted ID colours (the spec's preview), selection brighter.
             0 => {
-                let c = idx[i].map_or([0.0; 3], |k| crypto_colors(&planes, k, memo));
+                let c = k.map_or([0.0; 3], |k| crypto_colors(&planes, k, memo));
                 let c = c.map(|v| v * (1.0 - 0.5 * m) + 0.5 * m);
                 [c[0], c[1], c[2], 1.0]
             }
@@ -884,13 +910,25 @@ mod tests {
         let a = edge_aux(true);
         let sel = edge_selection();
         let full = crypto_aux_image(&a, "CryptoObject", &sel);
-        let matte = crypto_aux_image_with(&a, "CryptoObject", &sel, false);
+        let matte = crypto_aux_image_with(&a, "CryptoObject", &sel, CryptoAux::MatteOnly);
         for (f, m) in full.data.iter().zip(&matte.data) {
             assert_eq!(f[3].to_bits(), m[3].to_bits());
             assert_eq!(m[..3], [0.0; 3]);
         }
     }
 
+    /// Only the Colors display without Matte Only reads the ID colours.
+    #[test]
+    fn crypto_aux_for_display_reads_colors_only_for_the_colors_display() {
+        assert_eq!(CryptoAux::for_display(0, false), CryptoAux::ColorsAndMatte);
+        assert_eq!(CryptoAux::for_display(0, true), CryptoAux::MatteOnly);
+        for display in [1, 2, 3, u32::MAX] {
+            for matte_only in [false, true] {
+                assert_eq!(CryptoAux::for_display(display, matte_only), CryptoAux::MatteOnly, "display {display}");
+            }
+        }
+    }
+
     /// The CPU effect on a 20×10 layer over the 6×2 edge aux: pixels inside the aux area match
     /// the lookup by name bit for bit in every display mode, pixels outside it (no aux index)
     /// get an empty matte and no ID colour.