1use std::panic::AssertUnwindSafe;
2use std::path::Path;
3use std::sync::atomic::{AtomicUsize, Ordering};
4use std::sync::{Arc, Mutex};
5
6use crate::mkv;
7
8use super::demux::{Demux, mkv_head};
9use super::native;
10use super::picture::Color;
11use super::probe::keyframes;
12use super::route::{Caps, Prefs, pick_decode};
13use super::video::{Source, source};
14use crate::media::{Codec, Error, Result, Rgba, Route};
15
16pub fn first_frame(path: &Path, height: u32, prefs: &Prefs) -> Result<Rgba> {
17 frames_at(path, &[0.0], height, prefs)?
18 .pop()
19 .ok_or_else(|| Error::Invalid("no frame".to_string()))
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
23pub enum Seek {
24 #[default]
25 Keyframe,
26 Exact,
27}
28
29pub fn frames_at(path: &Path, times: &[f64], height: u32, prefs: &Prefs) -> Result<Vec<Rgba>> {
30 frames_at_seek(path, times, height, Seek::Keyframe, prefs)
31}
32
33pub fn frames_at_seek(
34 path: &Path,
35 times: &[f64],
36 height: u32,
37 seek: Seek,
38 prefs: &Prefs,
39) -> Result<Vec<Rgba>> {
40 if let Ok(Some(r)) = mkv_head(path)
41 && let Some(route) = mkv_route(&r, prefs)
42 {
43 return mkv_stills(&r, times, height, route, seek);
44 }
45 match seek {
46 Seek::Keyframe => keyframe_stills(path, times, height, prefs),
47 Seek::Exact => exact_stills(path, times, height, prefs),
48 }
49}
50
51fn keyframe_stills(path: &Path, times: &[f64], height: u32, prefs: &Prefs) -> Result<Vec<Rgba>> {
52 let src = source(path, prefs)?;
53 let keys = if times.iter().all(|t| *t <= 0.0) {
54 vec![0.0]
55 } else {
56 keyframes(path, prefs)?
57 };
58 let snap = |t: f64| {
59 keys.iter()
60 .copied()
61 .take_while(|k| *k <= t + 1e-6)
62 .last()
63 .unwrap_or(0.0)
64 };
65 let times: Vec<f64> = times.iter().map(|t| snap(*t)).collect();
66 let workers = times.len().min(16);
67 let next = AtomicUsize::new(0);
68 let out: Mutex<Vec<Option<Result<Rgba>>>> =
69 Mutex::new((0..times.len()).map(|_| None).collect());
70 std::thread::scope(|s| {
71 for w in 0..workers {
72 let (src, next, out, times) = (&src, &next, &out, ×);
73 s.spawn(move || {
74 let route = if w == 0 { src.route } else { Route::Cpu };
75 let mut dec = match (src.route, src.mp4()) {
76 (Route::Cpu | Route::Gpu, Some(r)) => {
77 native::still_decoder(r, route).ok().map(|d| (r, d))
78 }
79 _ => None,
80 };
81 loop {
82 let i = next.fetch_add(1, Ordering::Relaxed);
83 let Some(&t) = times.get(i) else { break };
84 let repeat = times.get(..i).is_some_and(|b| b.contains(&t));
85 let got =
86 std::panic::catch_unwind(AssertUnwindSafe(|| match (repeat, &mut dec) {
87 (true, _) => None,
88 (false, Some((r, d))) => Some(native::still(r, d, t, height)),
89 (false, None) => Some(still(src, t, height)),
90 }));
91 let got = got.unwrap_or_else(|_| {
92 dec = None;
93 Some(Err(Error::Backend(format!(
94 "decoding the still at {t:.3} s failed"
95 ))))
96 });
97 if let (Some(got), Ok(mut o)) = (got, out.lock()) {
98 o[i] = Some(got);
99 }
100 }
101 });
102 }
103 });
104 let done = out
105 .into_inner()
106 .map_err(|_| Error::Backend("a still worker failed".to_string()))?;
107 let mut stills: Vec<Rgba> = Vec::with_capacity(times.len());
108 for (i, r) in done.into_iter().enumerate() {
109 match r {
110 Some(r) => stills.push(r?),
111 None => {
112 let same = times[..i].iter().position(|o| *o == times[i]);
113 match same.and_then(|j| stills.get(j)) {
114 Some(s) => stills.push(s.clone()),
115 None => return Err(Error::Backend("a still worker failed".to_string())),
116 }
117 }
118 }
119 }
120 Ok(stills)
121}
122
123fn mkv_route(r: &mkv::Reader, prefs: &Prefs) -> Option<Route> {
124 let codec = r.tracks.get(r.video()?)?.codec()?;
125 if !matches!(codec, Codec::H264 | Codec::Hevc | Codec::Vp9) {
126 return None;
127 }
128 let picked = pick_decode(codec, prefs, &Caps::detect(prefs)).ok()?;
129 (picked.route != Route::Ffmpeg).then_some(picked.route)
130}
131
132fn failed() -> Error {
133 Error::Backend("a still worker failed".to_string())
134}
135
136fn spread<W>(
137 jobs: usize,
138 route: Route,
139 make: impl Fn(Route) -> W + Sync,
140 run: impl Fn(&mut W, usize) -> Result<Vec<Rgba>> + Sync,
141) -> Result<Vec<Rgba>> {
142 let next = AtomicUsize::new(0);
143 let out: Mutex<Vec<Option<Result<Vec<Rgba>>>>> = Mutex::new((0..jobs).map(|_| None).collect());
144 std::thread::scope(|s| {
145 for w in 0..jobs.min(16) {
146 let (next, out, make, run) = (&next, &out, &make, &run);
147 s.spawn(move || {
148 let mut state = make(if w == 0 { route } else { Route::Cpu });
149 loop {
150 let i = next.fetch_add(1, Ordering::Relaxed);
151 if i >= jobs {
152 break;
153 }
154 let got = std::panic::catch_unwind(AssertUnwindSafe(|| run(&mut state, i)))
155 .unwrap_or_else(|_| {
156 Err(Error::Backend("decoding a still failed".to_string()))
157 });
158 if got.is_err() {
159 state = make(Route::Cpu);
160 }
161 if let Ok(mut o) = out.lock()
162 && let Some(slot) = o.get_mut(i)
163 {
164 *slot = Some(got);
165 }
166 }
167 });
168 }
169 });
170 let done = out.into_inner().map_err(|_| failed())?;
171 let mut stills = Vec::new();
172 for r in done {
173 stills.extend(r.ok_or_else(failed)??);
174 }
175 Ok(stills)
176}
177
178fn gather(wanted: &[f64], unique: &[f64], stills: &[Rgba]) -> Result<Vec<Rgba>> {
179 wanted
180 .iter()
181 .map(|at| {
182 unique
183 .iter()
184 .position(|u| u == at)
185 .and_then(|j| stills.get(j).cloned())
186 .ok_or_else(failed)
187 })
188 .collect()
189}
190
191fn share(jobs: usize) -> usize {
192 (native::threads() / jobs.clamp(1, 16)).max(1)
193}
194
195fn sorted(times: &[f64]) -> Vec<f64> {
196 let mut unique: Vec<f64> = times.iter().map(|t| t.max(0.0)).collect();
197 unique.sort_by(f64::total_cmp);
198 unique.dedup();
199 unique
200}
201
202fn batches<K: PartialEq>(unique: &[f64], key: impl Fn(f64) -> K) -> Vec<Vec<f64>> {
203 let mut jobs: Vec<(K, Vec<f64>)> = Vec::new();
204 for &at in unique {
205 let k = key(at);
206 match jobs.last_mut() {
207 Some((last, batch)) if *last == k => batch.push(at),
208 _ => jobs.push((k, vec![at])),
209 }
210 }
211 jobs.into_iter().map(|(_, b)| b).collect()
212}
213
214fn tagged(d: &Demux, track: usize) -> Color {
215 let (matrix, full_range) = d
216 .stream(track)
217 .ok()
218 .and_then(|s| s.color)
219 .unwrap_or((2, false));
220 Color { matrix, full_range }
221}
222
223fn mkv_stills(
224 r: &Arc<mkv::Reader>,
225 times: &[f64],
226 height: u32,
227 route: Route,
228 seek: Seek,
229) -> Result<Vec<Rgba>> {
230 let d = Demux::Mkv(r.clone());
231 let v = d.video()?;
232 let t = r
233 .tracks
234 .get(v)
235 .ok_or_else(|| Error::Invalid("no video track".to_string()))?;
236 let rotation = t.rotation();
237 let color = tagged(&d, v);
238 let start = d.start();
239 let shift = d.shift(v);
240 let mut keys: Vec<f64> = r
241 .cues
242 .iter()
243 .filter(|c| c.track == t.number)
244 .map(|c| (t.seconds(c.time as i64) - start).max(0.0))
245 .collect();
246 keys.sort_by(f64::total_cmp);
247 keys.dedup();
248 let snap = |at: f64| {
249 if at <= 0.0 {
250 0.0
251 } else if keys.is_empty() {
252 at
253 } else {
254 keys.iter()
255 .copied()
256 .take_while(|k| *k <= at + 1e-6)
257 .last()
258 .unwrap_or(0.0)
259 }
260 };
261 let (wanted, unique, jobs) = match seek {
262 Seek::Keyframe => {
263 let wanted: Vec<f64> = times.iter().map(|at| snap(*at)).collect();
264 let mut unique: Vec<f64> = Vec::with_capacity(wanted.len());
265 for at in &wanted {
266 if !unique.contains(at) {
267 unique.push(*at);
268 }
269 }
270 let jobs = unique.iter().map(|u| vec![*u]).collect();
271 (wanted, unique, jobs)
272 }
273 Seek::Exact => {
274 let wanted: Vec<f64> = times.iter().map(|at| at.max(0.0)).collect();
275 let unique = sorted(times);
276 let jobs = if keys.is_empty() {
277 unique.iter().map(|u| vec![*u]).collect()
278 } else {
279 batches(&unique, |at| snap(at).to_bits())
280 };
281 (wanted, unique, jobs)
282 }
283 };
284 let threads = match seek {
285 Seek::Keyframe => 1,
286 Seek::Exact => share(jobs.len()),
287 };
288 let stills = spread(
289 jobs.len(),
290 route,
291 |rt| {
292 let dec = native::video_decoder(&d, rt, threads).map(|mut dec| {
293 if seek == Seek::Keyframe {
294 dec.stills();
295 }
296 dec
297 });
298 (dec, mkv::Cursor::new(r))
299 },
300 |(dec, cur), j| {
301 let dec = dec.as_mut().map_err(|e| Error::Backend(e.to_string()))?;
302 let batch = jobs.get(j).ok_or_else(failed)?;
303 let first = batch.first().copied().ok_or_else(failed)?;
304 match seek {
305 Seek::Keyframe => Ok(vec![native::still_mkv(r, cur, dec, first, height)?]),
306 Seek::Exact => {
307 cur.seek(r, v, first + start)?;
308 let targets: Vec<i64> = batch.iter().map(|at| t.ticks(*at)).collect();
309 let mut src = native::Packets::Mkv {
310 r,
311 cur,
312 track: v,
313 shift,
314 };
315 native::exact(&mut src, dec, &targets, rotation, height, color)
316 }
317 }
318 },
319 )?;
320 gather(&wanted, &unique, &stills)
321}
322
323fn exact_stills(path: &Path, times: &[f64], height: u32, prefs: &Prefs) -> Result<Vec<Rgba>> {
324 let src = source(path, prefs)?;
325 let wanted: Vec<f64> = times.iter().map(|at| at.max(0.0)).collect();
326 let unique = sorted(times);
327 let reader = match (src.route, src.mp4()) {
328 (Route::Ffmpeg, _) => {
329 let stills = unique
330 .iter()
331 .map(|at| ffmpeg_still(&src, *at, height, true))
332 .collect::<Result<Vec<_>>>()?;
333 return gather(&wanted, &unique, &stills);
334 }
335 (_, Some(r)) => r.clone(),
336 (_, None) => return Err(src.native()),
337 };
338 let d = Demux::Mp4(reader.clone());
339 let v = d.video()?;
340 let rotation = reader.tracks.get(v).map_or(0, |t| t.rotation());
341 let color = tagged(&d, v);
342 let seeks: Vec<(f64, usize, i64)> = unique
343 .iter()
344 .map(|at| d.seek(v, *at).map(|(key, start)| (*at, key, start)))
345 .collect::<Result<_>>()?;
346 let jobs = batches(&unique, |at| {
347 seeks.iter().find(|s| s.0 == at).map_or(0, |s| s.1)
348 });
349 let threads = share(jobs.len());
350 let stills = spread(
351 jobs.len(),
352 src.route,
353 |rt| native::video_decoder(&d, rt, threads),
354 |dec, j| {
355 let dec = dec.as_mut().map_err(|e| Error::Backend(e.to_string()))?;
356 let batch = jobs.get(j).ok_or_else(failed)?;
357 let plan: Vec<(usize, i64)> = batch
358 .iter()
359 .filter_map(|at| seeks.iter().find(|s| s.0 == *at).map(|s| (s.1, s.2)))
360 .collect();
361 let key = plan.first().map(|p| p.0).ok_or_else(failed)?;
362 let targets: Vec<i64> = plan.iter().map(|p| p.1).collect();
363 let mut packets = native::Packets::Indexed {
364 d: &d,
365 track: v,
366 next: key,
367 };
368 native::exact(&mut packets, dec, &targets, rotation, height, color)
369 },
370 )?;
371 gather(&wanted, &unique, &stills)
372}
373
374fn still(src: &Source, key: f64, height: u32) -> Result<Rgba> {
375 ffmpeg_still(src, key, height, false)
376}
377
378fn ffmpeg_still(src: &Source, key: f64, height: u32, exact: bool) -> Result<Rgba> {
379 if src.route != Route::Ffmpeg {
380 return Err(src.native());
381 }
382 let f = src.ffmpeg()?;
383 let mut args = f.decoder_args(src.ffmpeg_codec.as_deref());
384 args.extend(["-threads".into(), "2".into()]);
385 if !exact {
386 args.extend(["-skip_frame".into(), "nokey".into()]);
387 }
388 if key > 0.0 {
389 let at = if exact { key } else { (key - 0.0005).max(0.0) };
390 args.extend(["-ss".into(), format!("{at:.6}")]);
391 }
392 args.extend([
393 "-i".into(),
394 src.path.display().to_string(),
395 "-map".into(),
396 "0:v:0".into(),
397 "-frames:v".into(),
398 "1".into(),
399 "-vf".into(),
400 format!("scale=-2:{}:flags=bilinear", height.max(2)),
401 "-f".into(),
402 "rawvideo".into(),
403 "-pix_fmt".into(),
404 "rgba".into(),
405 "-".into(),
406 ]);
407 let data = f.run_args(&args)?;
408 let h = height.max(2);
409 let w = (data.len() / 4 / h as usize) as u32;
410 if w == 0 || data.len() != w as usize * h as usize * 4 {
411 return Err(Error::Backend(format!(
412 "ffmpeg gave {} bytes for a still {h} high",
413 data.len()
414 )));
415 }
416 Ok(Rgba {
417 width: w,
418 height: h,
419 data,
420 })
421}