Cryptomatte: read an ID's coverage from the kept probe (#473)

Session::crypto_coverage returns the coverage of one Cryptomatte ID over the layer, read from
the probe Session::crypto_probe_job keeps: at most max_side samples on the long side, each at
the centre of its cell, or None until the probe is built. It is what the Cryptomatte
eyedropper's hover highlight will draw.

Tests cover the coverage of an ID from the kept probe and the cost of a highlight with six
ranks.

crates/engine/src/cmd.rs

@@ -18,8 +18,8 @@ pub mod content_fill;
 mod create;
 mod edit;
 mod effect;
-pub use effect::CryptoProbeJob;
-pub(crate) use effect::{CryptoProbeSlot, crypto_probe_job, crypto_probe_ready, keep_crypto_probe, kept_crypto_probe};
+pub use effect::{CryptoCoverage, CryptoProbeJob};
+pub(crate) use effect::{CryptoProbeSlot, crypto_coverage, crypto_probe_job, crypto_probe_ready, keep_crypto_probe, kept_crypto_probe};
 pub mod essential;
 pub mod expr_tools;
 mod file;

crates/engine/src/cmd/effect.rs

@@ -500,6 +500,39 @@ pub(crate) fn crypto_probe_job(s: &Session, p: &Value) -> Result<Option<CryptoPr
     })))
 }
 
+/// [`Session::crypto_coverage`].
+pub(crate) fn crypto_coverage(s: &Session, p: &Value, hash: u32, max_side: usize) -> Option<CryptoCoverage> {
+    let (key, ..) = probe_key(s, p, "effect.cryptomatteAt").ok()?;
+    let probe = cached_probe(s, key)?;
+    let b = probe.grid.as_ref().ok()?;
+    let (aw, ah) = (b.aux.width as usize, b.aux.height as usize);
+    if aw == 0 || ah == 0 {
+        return None;
+    }
+    let step = aw.max(ah).div_ceil(max_side.max(1)).max(1);
+    let (w, h) = (aw.div_ceil(step), ah.div_ceil(step));
+    let planes = b.aux.crypto_rank_planes(&b.layer);
+    let sample = |j: usize| {
+        planes.iter().filter(|(id, _)| id.get(j).is_some_and(|v| v.to_bits() == hash)).map(|(_, c)| c.get(j).map_or(0.0, |c| c.max(0.0))).sum::<f32>().min(1.0)
+    };
+    // Each sample at the centre of its step × step cell; the last cells are cut by the layer.
+    let at = |x: usize, y: usize| sample((y * step + step / 2).min(ah - 1) * aw + (x * step + step / 2).min(aw - 1));
+    let coverage = (0..w * h).map(|i| at(i % w, i / w)).collect();
+    let (size, uv) = ([aw as f64 / b.aux.scale, ah as f64 / b.aux.scale], [aw as f64 / (w * step) as f64, ah as f64 / (h * step) as f64]);
+    Some(CryptoCoverage { width: w, height: h, coverage, size, uv })
+}
+
+/// An ID's coverage over a layer, for the eyedropper's hover highlight: `width` × `height`
+/// samples (row-major, 0..1). The layer, `size` effect-space pixels from its origin, shows the
+/// texture from (0, 0) to `uv`.
+pub struct CryptoCoverage {
+    pub width: usize,
+    pub height: usize,
+    pub coverage: Vec<f32>,
+    pub size: [f64; 2],
+    pub uv: [f64; 2],
+}
+
 /// What a pick finds at a point.
 enum Hit {
     /// `probeOnly` and the probe isn't built yet.

crates/engine/src/session/render.rs

@@ -36,6 +36,13 @@ impl Session {
         commands::crypto_probe_job(self, p)
     }
 
+    /// The coverage of Cryptomatte ID `hash` (as `effect.cryptomatteAt` reports it, parsed) over
+    /// the layer of params `p`, at most `max_side` samples on the long side, from the kept probe:
+    /// `None` until [`Session::crypto_probe_job`] has built it.
+    pub fn crypto_coverage(&self, p: &Value, hash: u32, max_side: usize) -> Option<commands::CryptoCoverage> {
+        commands::crypto_coverage(self, p, hash, max_side)
+    }
+
     /// Let go of the kept Cryptomatte probe (and its aux channels), and cancel a build in
     /// flight: the eyedropper was disarmed. Edits, footage changes and Open / New do it too.
     pub fn drop_crypto_probe(&self) {

crates/engine/src/tests/crypto_pick.rs

@@ -431,7 +431,13 @@ fn hover_latency() {
                 format!("{:.3}", t0.elapsed().as_secs_f64() * 1e3)
             })
             .collect();
-        eprintln!("HOVER {name}: first {} ms, then {} ms", ms[0], ms[1..].join(" / "));
+        // The highlight of the hovered ID (built once per ID, at most 1024 samples wide).
+        let q = json!({"effect": "Cryptomatte", "x": 100.0, "y": 500.0});
+        let hash = u32::from_str_radix(s.execute("effect.cryptomatteAt", q.clone()).unwrap()["hash"].as_str().unwrap(), 16).unwrap();
+        let t0 = std::time::Instant::now();
+        let c = s.crypto_coverage(&q, hash, 1024).unwrap();
+        let cov = t0.elapsed().as_secs_f64() * 1e3;
+        eprintln!("HOVER {name}: first {} ms, then {} ms; highlight {}x{} {cov:.2} ms", ms[0], ms[1..].join(" / "), c.width, c.height);
     }
 }
 
@@ -794,3 +800,99 @@ fn undoing_picks_keeps_the_probe_and_the_step_says_pick() {
     s.execute_checked("edit.undo", json!({})).unwrap();
     assert!(!s.crypto_probe_ready(&q).unwrap(), "an undone opacity change needs a new probe");
 }
+
+/// The hover highlight: an ID's coverage from the kept probe, one sample at the centre of each
+/// cell of aux pixels, over the layer's extent (the last cells cut by it).
+#[test]
+fn coverage_of_an_id_from_the_kept_probe() {
+    let s = session();
+    let q = json!({"effect": 1, "x": 0.5, "y": 1.5});
+    let left = crypto_hash("left");
+    assert!(s.crypto_coverage(&q, left, 64).is_none(), "no coverage before the probe");
+    s.crypto_probe_job(&q).unwrap().unwrap()();
+    let c = s.crypto_coverage(&q, left, 64).unwrap();
+    assert_eq!((c.width, c.height, c.size, c.uv), (16, 8, [8.0, 4.0], [1.0, 1.0]), "the aux pixels over the layer");
+    assert!(c.coverage.iter().all(|v| (0.0..=1.0).contains(v)) && c.coverage.iter().any(|v| *v > 0.5) && c.coverage.contains(&0.0), "{:?}", c.coverage);
+    assert_eq!(c.coverage[0], 0.6, "left's coverage at the first aux pixel");
+    assert!(s.crypto_coverage(&q, crypto_hash("nobody"), 64).unwrap().coverage.iter().all(|v| *v == 0.0));
+    assert_eq!(s.crypto_coverage(&q, left, 2).map(|c| (c.width, c.height)), Some((2, 1)));
+    // An odd size: 2 × 2 cells, the last ones half outside the layer; samples at cell centres.
+    let mut s = cells_session(1919, 1079, false);
+    let q = json!({"effect": "Cryptomatte", "x": 40.0, "y": 10.0});
+    assert_eq!(s.execute("effect.cryptomatteAt", q.clone()).unwrap()["name"], "c1_0");
+    let c = s.crypto_coverage(&q, crypto_hash("c1_0"), 1024).unwrap();
+    assert_eq!((c.width, c.height, c.size), (960, 540, [1919.0, 1079.0]));
+    assert_eq!(c.uv, [1919.0 / 1920.0, 1079.0 / 1080.0]);
+    // Sample 18 covers pixels 36-37: its centre (37) is in c1_0, its corner (36) in c0_0.
+    assert_eq!((c.coverage[17], c.coverage[18], c.coverage[37]), (0.0, 1.0, 0.0));
+}
+
+struct Ranked(u32, u32, std::sync::OnceLock<Arc<AuxChannels>>);
+impl FootageSource for Ranked {
+    fn frame(&self, _: ItemId, f: &Footage, _: Tick) -> Option<Arc<Image>> {
+        Some(Arc::new(Image::new(f.width, f.height)))
+    }
+    fn aux(&self, _: ItemId, _: &Footage, _: Tick) -> Option<Arc<AuxChannels>> {
+        Some(
+            self.2
+                .get_or_init(|| {
+                    let (w, h) = (self.0 as usize, self.1 as usize);
+                    let mut aux = AuxChannels::new(w as u32, h as u32, 1.0);
+                    for part in 0..3 {
+                        for (k, c) in ["R", "G", "B", "A"].iter().enumerate() {
+                            let rank = part * 2 + k / 2;
+                            let v: Vec<f32> = (0..w * h)
+                                .map(|i| if k % 2 == 0 { crypto_float(crypto_hash(&format!("r{rank}_{}", Cells::name(i % w, i / w)))) } else { 1.0 / 6.0 })
+                                .collect();
+                            aux.channels.push((format!("CryptoObject{part:02}.{c}"), v));
+                        }
+                    }
+                    Arc::new(aux)
+                })
+                .clone(),
+        )
+    }
+}
+
+fn rss_mb() -> f64 {
+    let s = std::fs::read_to_string("/proc/self/status").unwrap_or_default();
+    s.lines().find(|l| l.starts_with("VmRSS")).and_then(|l| l.split_whitespace().nth(1)).and_then(|v| v.parse::<f64>().ok()).unwrap_or(0.0) / 1024.0
+}
+
+/// Highlight cost with 6 ranks (Cryptomatte's default 3 parts) at 1080p, 4K and 8K, and the memory the
+/// probe pins. Run with `--ignored --nocapture`; it measures, it doesn't assert.
+#[test]
+#[ignore]
+fn highlight_cost_with_six_ranks() {
+    for (w, h) in [(1920u32, 1080u32), (3840, 2160), (7680, 4320)] {
+        let r0 = rss_mb();
+        let mut s = Session { footage: Arc::new(Ranked(w, h, Default::default())), importer: Some(Arc::new(CellsProbe(w, h))), ..Default::default() };
+        let item = s.execute("file.import", json!({"paths": ["/renders/ranked.exr"]})).unwrap()["items"][0].clone();
+        s.execute("comp.new", json!({"name": "C", "width": w, "height": h, "frameRate": 24, "duration": 2})).unwrap();
+        s.execute("layer.addItem", json!({"item": item})).unwrap();
+        s.execute_checked("effect.apply", json!({"effect": "Cryptomatte"})).unwrap();
+        let q = json!({"effect": "Cryptomatte", "x": 100.0, "y": 500.0});
+        let t0 = std::time::Instant::now();
+        let r = s.execute("effect.cryptomatteAt", q.clone()).unwrap();
+        let build = t0.elapsed().as_secs_f64() * 1e3;
+        let hash = u32::from_str_radix(r["hash"].as_str().unwrap(), 16).unwrap();
+        let mut ms = vec![];
+        for _ in 0..5 {
+            let t0 = std::time::Instant::now();
+            let c = s.crypto_coverage(&q, hash, 1024).unwrap();
+            ms.push(t0.elapsed().as_secs_f64() * 1e3);
+            assert!(c.coverage.iter().any(|v| *v > 0.0));
+        }
+        ms.sort_by(f64::total_cmp);
+        let aux = s.footage.aux(ItemId(0), &Footage::default(), Tick::ZERO).unwrap();
+        let bytes: usize = aux.channels.iter().map(|(_, v)| v.len() * 4).sum();
+        eprintln!(
+            "HIGHLIGHT 6 ranks {w}x{h}: aux {:.0} MB ({} planes), probe build {build:.1} ms, highlight median {:.2} ms (min {:.2}), RSS +{:.0} MB",
+            bytes as f64 / 1e6,
+            aux.channels.len(),
+            ms[2],
+            ms[0],
+            rss_mb() - r0
+        );
+    }
+}