1use rayon::prelude::*;
2
3use super::color::{ColorSpec, Levels};
4use super::plane::{Sample, check, samples_mut};
5use crate::media::Rgba;
6use crate::media::{Error, Frame, PixelFormat, Result};
7
8const BITS: u32 = 20;
9const MASK: u64 = (1 << BITS) - 1;
10const SLOTS: usize = 1024;
11const BLOCKS: usize = 256;
12
13#[derive(Debug, Clone, Default)]
14struct Band<T> {
15 top: usize,
16 spans: Vec<(usize, usize)>,
17 cells: Vec<T>,
18}
19
20#[derive(Debug, Clone, Default)]
21struct Layer<T> {
22 left: i64,
23 top: i64,
24 width: usize,
25 bands: Vec<Band<T>>,
26}
27
28impl<T: Copy> Band<T> {
29 fn new(top: usize, width: usize, cells: Vec<T>, alpha: impl Fn(T) -> u64) -> Self {
30 let spans = cells
31 .chunks_exact(width.max(1))
32 .map(|row| match row.iter().position(|&c| alpha(c) > 0) {
33 Some(first) => {
34 let last = row.iter().rposition(|&c| alpha(c) > 0).unwrap_or(first);
35 (first, last + 1)
36 }
37 None => (0, 0),
38 })
39 .collect();
40 Band { top, spans, cells }
41 }
42}
43
44impl<T: Copy> Layer<T> {
45 #[allow(clippy::too_many_arguments)]
46 fn draw<S: Sample>(
47 &self,
48 dst: &mut [S],
49 stride: usize,
50 dw: usize,
51 dh: usize,
52 x: i64,
53 y: i64,
54 pick: impl Fn(T) -> (u32, u32),
55 ) {
56 let (x, y) = (x + self.left, y + self.top);
57 for band in &self.bands {
58 for (r, &(a, b)) in band.spans.iter().enumerate() {
59 let ty = y + (band.top + r) as i64;
60 if ty < 0 || ty >= dh as i64 {
61 continue;
62 }
63 let from = (a as i64).max(-x) as usize;
64 let to = (b as i64).min(dw as i64 - x);
65 if to <= from as i64 {
66 continue;
67 }
68 let to = to as usize;
69 let at = ty as usize * stride + (x + from as i64) as usize;
70 let base = r * self.width;
71 let (Some(out), Some(cells)) = (
72 dst.get_mut(at..at + (to - from)),
73 band.cells.get(base + from..base + to),
74 ) else {
75 continue;
76 };
77 for (o, &c) in out.iter_mut().zip(cells) {
78 let (al, v) = pick(c);
79 *o = S::put(((o.get() as u32 * (256 - al) + v + 128) >> 8) as i32);
80 }
81 }
82 }
83 }
84}
85
86struct Convert<'a> {
87 data: &'a [u8],
88 width: usize,
89 matrix: [[f64; 3]; 3],
90 levels: Levels,
91 top: f64,
92 alpha: [f64; 256],
93}
94
95struct Cache {
96 keys: [u32; SLOTS],
97 values: [[f64; 3]; SLOTS],
98}
99
100struct Blocks {
101 keys: [u32; BLOCKS],
102 values: [(u32, u64); BLOCKS],
103}
104
105impl Convert<'_> {
106 fn word(&self, x: usize, y: usize) -> u32 {
107 let i = (y * self.width + x) * 4;
108 self.data
109 .get(i..i + 4)
110 .and_then(|p| p.try_into().ok())
111 .map_or(0, u32::from_le_bytes)
112 }
113
114 fn pixel(&self, cache: &mut Cache, word: u32) -> (u32, [f64; 3]) {
115 let [r, g, b, a] = word.to_le_bytes();
116 if a == 0 {
117 return (0, [0.0; 3]);
118 }
119 let key = word & 0xff_ffff;
120 let slot = (key.wrapping_mul(0x9e37_79b1) >> 22) as usize % SLOTS;
121 if cache.keys[slot] != key {
122 let c = [r as f64 / 255.0, g as f64 / 255.0, b as f64 / 255.0];
123 let m = &self.matrix;
124 let dot = |k: usize| m[k][0] * c[0] + m[k][1] * c[1] + m[k][2] * c[2];
125 let l = &self.levels;
126 cache.keys[slot] = key;
127 cache.values[slot] = [
128 (l.y_off + l.y_range * dot(0)).clamp(0.0, self.top),
129 l.c_off + l.c_range * dot(1),
130 l.c_off + l.c_range * dot(2),
131 ];
132 }
133 let [y, u, v] = cache.values[slot];
134 let a = self.alpha[a as usize];
135 let luma = (a as u32) << BITS | round(y * a) as u32;
136 (luma, [a, a * u, a * v])
137 }
138
139 fn chroma(&self, q: [[f64; 3]; 4]) -> u64 {
140 let mut s = [0.0f64; 3];
141 for q in q {
142 s.iter_mut().zip(q).for_each(|(a, b)| *a += b);
143 }
144 let a = round(s[0] / 4.0) as u64;
145 if a == 0 {
146 return 0;
147 }
148 let max = self.top * 256.0;
149 let u = round(s[1] / 4.0).clamp(0.0, max) as u64;
150 let v = round(s[2] / 4.0).clamp(0.0, max) as u64;
151 a << (2 * BITS) | u << BITS | v
152 }
153}
154
155fn round(x: f64) -> f64 {
156 let t = x as i64 as f64;
157 if x - t >= 0.5 { t + 1.0 } else { t }
158}
159
160fn painted(row: &[u8]) -> Option<(usize, usize)> {
161 let words = row.chunks_exact(8);
162 let tail = words.remainder().get(3).copied().unwrap_or(0);
163 let mask = words.fold(0u64, |m, c| m | c.try_into().map_or(0, u64::from_le_bytes));
164 if mask & 0xff00_0000_ff00_0000 == 0 && tail == 0 {
165 return None;
166 }
167 let first = row.chunks_exact(4).position(|p| p[3] > 0)?;
168 let last = row.chunks_exact(4).rposition(|p| p[3] > 0)?;
169 Some((first, last + 1))
170}
171
172#[derive(Debug, Clone)]
173pub struct Overlay {
174 format: PixelFormat,
175 width: u32,
176 height: u32,
177 luma: Layer<u32>,
178 chroma: Layer<u64>,
179}
180
181impl Overlay {
182 pub fn new(image: &Rgba, spec: ColorSpec, format: PixelFormat) -> Result<Self> {
183 let w = image.width as usize;
184 let h = image.height as usize;
185 super::plane::sized(image.width, image.height)?;
186 if image.data.len() < w * h * 4 {
187 return Err(Error::Invalid(
188 "overlay data is shorter than its size".into(),
189 ));
190 }
191 let mut out = Overlay {
192 format,
193 width: image.width,
194 height: image.height,
195 luma: Layer::default(),
196 chroma: Layer::default(),
197 };
198 if w == 0 || h == 0 {
199 return Ok(out);
200 }
201 let data = &image.data[..w * h * 4];
202 let rows: Vec<Option<(usize, usize)>> = data
203 .par_chunks(w * 4)
204 .with_min_len(64)
205 .map(painted)
206 .collect();
207 let (Some(y0), Some(y1)) = (
208 rows.iter().position(Option::is_some),
209 rows.iter().rposition(Option::is_some),
210 ) else {
211 return Ok(out);
212 };
213 let (x0, x1) = rows
214 .iter()
215 .flatten()
216 .fold((w, 0), |(a, b), &(f, l)| (a.min(f), b.max(l)));
217 let (lx0, ly0) = (x0 & !1, y0 & !1);
218 let (lx1, ly1) = (((x1 + 1) & !1).min(w), ((y1 + 2) & !1).min(h));
219 let (lw, lh) = (lx1 - lx0, ly1 - ly0);
220 let (cw, ch) = (lw.div_ceil(2), lh.div_ceil(2));
221 let convert = Convert {
222 data,
223 width: w,
224 matrix: spec.matrix.forward(),
225 levels: spec.levels(format.bit_depth()),
226 top: ((1u32 << format.bit_depth()) - 1) as f64,
227 alpha: std::array::from_fn(|a| (a as f64 * 256.0 / 255.0).round()),
228 };
229 let k = super::color::rows_job(ch).max(8);
230 let (luma, chroma): (Vec<_>, Vec<_>) = (0..ch.div_ceil(k))
231 .into_par_iter()
232 .map(|n| {
233 let (c0, c1) = (n * k, ((n + 1) * k).min(ch));
234 let (l0, l1) = (2 * c0, (2 * c1).min(lh));
235 let mut yc = vec![0u32; lw * (l1 - l0)];
236 let mut cc = vec![0u64; cw * (c1 - c0)];
237 let mut cache = Cache {
238 keys: [u32::MAX; SLOTS],
239 values: [[0.0; 3]; SLOTS],
240 };
241 let mut blocks = Blocks {
242 keys: [0; BLOCKS],
243 values: [(0, 0); BLOCKS],
244 };
245 for j in 0..c1 - c0 {
246 let r0 = ly0 + l0 + 2 * j;
247 let r1 = (r0 + 1).min(h - 1);
248 let Some((xa, xb)) = [rows.get(r0), rows.get(r1)]
249 .into_iter()
250 .flatten()
251 .flatten()
252 .fold(None, |s: Option<(usize, usize)>, &(f, l)| {
253 Some(s.map_or((f, l), |(a, b)| (a.min(f), b.max(l))))
254 })
255 else {
256 continue;
257 };
258 for bx in xa / 2..xb.div_ceil(2) {
259 let (p0, p1) = (2 * bx, (2 * bx + 1).min(w - 1));
260 let words = [
261 convert.word(p0, r0),
262 convert.word(p1, r0),
263 convert.word(p0, r1),
264 convert.word(p1, r1),
265 ];
266 if words.iter().fold(0, |m, w| m | w) >> 24 == 0 {
267 continue;
268 }
269 let (y, c) = if words.iter().all(|&w| w == words[0]) {
270 let slot = (words[0].wrapping_mul(0x9e37_79b1) >> 24) as usize % BLOCKS;
271 if blocks.keys[slot] != words[0] {
272 let q = convert.pixel(&mut cache, words[0]);
273 blocks.keys[slot] = words[0];
274 blocks.values[slot] = (q.0, convert.chroma([q.1; 4]));
275 }
276 let (y, c) = blocks.values[slot];
277 ([y; 4], c)
278 } else {
279 let q = words.map(|w| convert.pixel(&mut cache, w));
280 (q.map(|q| q.0), convert.chroma(q.map(|q| q.1)))
281 };
282 let row0 = 2 * j * lw;
283 let row1 = row0 + lw;
284 let at = [
285 (row0 + p0 - lx0, true),
286 (row0 + p1 - lx0, p1 != p0),
287 (row1 + p0 - lx0, r1 != r0),
288 (row1 + p1 - lx0, r1 != r0 && p1 != p0),
289 ];
290 for ((i, real), y) in at.into_iter().zip(y) {
291 if real && let Some(o) = yc.get_mut(i) {
292 *o = y;
293 }
294 }
295 if let Some(o) = cc.get_mut(j * cw + bx - lx0 / 2) {
296 *o = c;
297 }
298 }
299 }
300 (
301 Band::new(l0, lw, yc, |c| (c >> BITS) as u64),
302 Band::new(c0, cw, cc, |c| c >> (2 * BITS)),
303 )
304 })
305 .unzip();
306 out.luma = Layer {
307 left: lx0 as i64,
308 top: ly0 as i64,
309 width: lw,
310 bands: luma,
311 };
312 out.chroma = Layer {
313 left: (lx0 / 2) as i64,
314 top: (ly0 / 2) as i64,
315 width: cw,
316 bands: chroma,
317 };
318 Ok(out)
319 }
320
321 pub fn size(&self) -> (u32, u32) {
322 (self.width, self.height)
323 }
324
325 pub fn draw(&self, frame: &mut Frame, x: i32, y: i32) -> Result<()> {
326 check(frame)?;
327 if frame.format != self.format {
328 return Err(Error::Invalid(format!(
329 "overlay made for {:?}, frame is {:?}",
330 self.format, frame.format
331 )));
332 }
333 let size = (frame.width as usize, frame.height as usize);
334 let at = (x as i64, y as i64);
335 for (k, plane) in frame.planes.iter_mut().enumerate().take(3) {
336 match self.format {
337 PixelFormat::Yuv420p8 => {
338 let stride = plane.stride;
339 self.paint(k, samples_mut::<u8>(plane)?, stride, size, at);
340 }
341 PixelFormat::Yuv420p10 => {
342 let stride = plane.stride / 2;
343 self.paint(k, samples_mut::<u16>(plane)?, stride, size, at);
344 }
345 }
346 }
347 Ok(())
348 }
349
350 fn paint<S: Sample>(
351 &self,
352 k: usize,
353 dst: &mut [S],
354 stride: usize,
355 (fw, fh): (usize, usize),
356 (x, y): (i64, i64),
357 ) {
358 let (cw, ch) = (fw.div_ceil(2), fh.div_ceil(2));
359 let (cx, cy) = (x.div_euclid(2), y.div_euclid(2));
360 let mask = MASK as u32;
361 match k {
362 0 => self
363 .luma
364 .draw(dst, stride, fw, fh, x, y, |c| (c >> BITS, c & mask)),
365 1 => self.chroma.draw(dst, stride, cw, ch, cx, cy, |c| {
366 ((c >> (2 * BITS)) as u32, ((c >> BITS) & MASK) as u32)
367 }),
368 _ => self.chroma.draw(dst, stride, cw, ch, cx, cy, |c| {
369 ((c >> (2 * BITS)) as u32, (c & MASK) as u32)
370 }),
371 }
372 }
373}