1use crate::hevc::pel::Pel;
2use crate::hevc::picture::{Frame, FrameRef, PlanePool};
3use crate::hevc::slicedec::{PicState, Pool, Scaling, SliceDecoder, SliceJob, SliceRun};
4use crate::hevc::{ColorInfo, Hvcc};
5use crate::media::{Frame as Image, Packet, PixelFormat, Plane, TimeBase};
6use crate::syntax::bits::{annexb_nals, length_prefixed_nals, unescape_rbsp, unescape_rbsp_marked};
7use crate::syntax::error::{Error, Result, ensure, invalid};
8use crate::syntax::hevc::dpb::{Dpb, Output, OutputMeta, PocState, RefPicSet, build_ref_lists};
9use crate::syntax::hevc::params::{Layout, Pps, Sps, Vps};
10use crate::syntax::hevc::slice::{NalHeader, SliceHeader, nal_type};
11use std::collections::VecDeque;
12use std::sync::{Arc, Mutex, PoisonError};
13
14type LayoutEntry = (Arc<Pps>, Arc<Sps>, Arc<Layout>, Scaling);
15
16struct PendingPic {
17 frame: FrameRef,
18 sps: Arc<Sps>,
19 pps: Arc<Pps>,
20 layout: Arc<Layout>,
21 last_sh: SliceHeader,
22 first_sh: SliceHeader,
23 rps: RefPicSet<FrameRef>,
24 poc: i32,
25 output: bool,
26 meta: Arc<OutputMeta>,
27 slices: Vec<SliceJob>,
28}
29
30#[derive(Default)]
31struct Shared {
32 error: Mutex<Option<Error>>,
33 pool8: Pool<u8>,
34 pool16: Pool<u16>,
35 rows: PlanePool,
36 images: Mutex<Vec<Vec<u8>>>,
37}
38
39impl Shared {
40 fn image_buffers(&self) -> [Vec<u8>; 3] {
41 let mut p = self.images.lock().unwrap_or_else(PoisonError::into_inner);
42 let mut b = [
43 p.pop().unwrap_or_default(),
44 p.pop().unwrap_or_default(),
45 p.pop().unwrap_or_default(),
46 ];
47 b.sort_by_key(|v| std::cmp::Reverse(v.capacity()));
48 b
49 }
50}
51
52pub struct Decoder {
53 vpss: Vec<Option<Arc<Vps>>>,
54 spss: Vec<Option<Arc<Sps>>>,
55 ppss: Vec<Option<Arc<Pps>>>,
56 nal_length_size: Option<usize>,
57 dpb: Dpb<FrameRef>,
58 poc_state: PocState,
59 pending: Option<PendingPic>,
60 next_id: u32,
61 active_sps: Option<Arc<Sps>>,
62 first_picture: bool,
63 after_eos: bool,
64 skip_rasl: bool,
65 layout_cache: Vec<LayoutEntry>,
66 out_queue: VecDeque<Output<FrameRef>>,
67 shared: Arc<Shared>,
68 pool: Option<rayon::ThreadPool>,
69 in_flight: VecDeque<FrameRef>,
70 max_in_flight: usize,
71 draft: bool,
72 time_base: TimeBase,
73 threads: usize,
74}
75
76impl Default for Decoder {
77 fn default() -> Self {
78 Self::new()
79 }
80}
81
82impl Decoder {
83 pub fn new() -> Self {
84 Self::with_threads(crate::media::default_threads())
85 }
86
87 pub fn with_threads(threads: usize) -> Self {
88 let pool = if threads > 1 {
89 rayon::ThreadPoolBuilder::new()
90 .num_threads(threads)
91 .thread_name(|i| format!("hevc-{i}"))
92 .build()
93 .ok()
94 } else {
95 None
96 };
97 Decoder {
98 vpss: vec![None; 16],
99 spss: vec![None; 16],
100 ppss: vec![None; 64],
101 nal_length_size: None,
102 dpb: Dpb::default(),
103 poc_state: PocState::default(),
104 pending: None,
105 next_id: 1,
106 active_sps: None,
107 first_picture: true,
108 after_eos: false,
109 skip_rasl: false,
110 layout_cache: Vec::new(),
111 out_queue: VecDeque::new(),
112 shared: Arc::new(Shared::default()),
113 pool,
114 in_flight: VecDeque::new(),
115 max_in_flight: threads.max(1) + 2,
116 threads: threads.max(1),
117 draft: false,
118 time_base: TimeBase::new(1, 90000),
119 }
120 }
121
122 pub fn from_hvcc(hvcc: &[u8]) -> crate::media::Result<Self> {
123 let mut d = Decoder::new();
124 d.configure_hvcc(hvcc)?;
125 Ok(d)
126 }
127
128 pub fn configure_hvcc(&mut self, hvcc: &[u8]) -> crate::media::Result<()> {
129 let config = Hvcc::parse(hvcc)?;
130 for nal in config.vps.iter().chain(&config.sps).chain(&config.pps) {
131 self.handle_nal(nal, 0).map_err(|e| e.media("hevc"))?;
132 }
133 self.nal_length_size = Some(config.nal_length_size);
134 Ok(())
135 }
136
137 pub fn set_time_base(&mut self, time_base: TimeBase) {
138 self.time_base = time_base;
139 }
140
141 pub fn color(&self) -> Option<ColorInfo> {
142 let vui = self.active_sps.as_ref()?.vui.clone().unwrap_or_default();
143 Some(ColorInfo {
144 full_range: vui.full_range,
145 primaries: vui.colour_primaries,
146 transfer: vui.transfer_characteristics,
147 matrix: vui.matrix_coefficients,
148 sar: vui.sar,
149 })
150 }
151
152 pub fn recycle(&self, frame: Image) {
153 let mut p = self
154 .shared
155 .images
156 .lock()
157 .unwrap_or_else(PoisonError::into_inner);
158 for plane in frame.planes {
159 if p.len() < 48 {
160 p.push(plane.data);
161 }
162 }
163 }
164
165 pub fn decode(&mut self, packet: &Packet) -> crate::media::Result<Vec<Image>> {
166 self.decode_au(&packet.data, packet.pts)
167 .map_err(|e| e.media("hevc"))
168 }
169
170 pub fn flush(&mut self) -> crate::media::Result<Vec<Image>> {
171 let pics = self.flush_pictures();
172 if let Some(e) = self.take_error() {
173 return Err(e.media("hevc"));
174 }
175 Ok(pics)
176 }
177
178 pub(crate) fn decode_au(&mut self, data: &[u8], pts: i64) -> Result<Vec<Image>> {
179 let nals = match self.nal_length_size {
180 Some(n) => length_prefixed_nals(data, n)?,
181 None => annexb_nals(data),
182 };
183 let mut result = Ok(());
184 for nal in nals {
185 if let Err(e) = self.handle_nal(nal, pts) {
186 result = Err(e);
187 break;
188 }
189 }
190 self.submit_pending();
191 if let Some(e) = self.take_error() {
192 result = result.and(Err(e));
193 }
194 result?;
195 let max_queue = self.max_in_flight * 2;
196 let mut pics = Vec::new();
197 while let Some(o) = self.out_queue.front() {
198 if o.frame.image.get().is_none() && self.out_queue.len() <= max_queue {
199 break;
200 }
201 let Some(o) = self.out_queue.pop_front() else {
202 break;
203 };
204 pics.push(make_picture(&o, self.time_base));
205 }
206 Ok(pics)
207 }
208
209 pub(crate) fn flush_pictures(&mut self) -> Vec<Image> {
210 self.submit_pending();
211 let mut outs = Vec::new();
212 self.dpb.flush(&mut outs);
213 self.out_queue.extend(outs);
214 let time_base = self.time_base;
215 let pics = self
216 .out_queue
217 .drain(..)
218 .map(|o| make_picture(&o, time_base))
219 .collect();
220 while let Some(f) = self.in_flight.pop_front() {
221 f.wait_complete();
222 }
223 self.first_picture = true;
224 pics
225 }
226
227 pub fn set_draft(&mut self, on: bool) {
228 self.draft = on;
229 }
230
231 pub(crate) fn take_error(&mut self) -> Option<Error> {
232 self.shared.error.lock().ok().and_then(|mut e| e.take())
233 }
234
235 fn handle_nal(&mut self, nal: &[u8], pts: i64) -> Result<()> {
236 if nal.len() < 2 {
237 return Ok(());
238 }
239 let hdr = NalHeader::parse(nal)?;
240 if hdr.layer_id != 0 {
241 return Ok(());
242 }
243 match hdr.nal_type {
244 nal_type::VPS => {
245 let v = Vps::parse(&unescape_rbsp(&nal[2..]))?;
246 let id = v.id as usize;
247 self.vpss[id] = Some(Arc::new(v));
248 }
249 nal_type::SPS => {
250 let s = Sps::parse(&unescape_rbsp(&nal[2..]))?;
251 let id = s.id as usize;
252 self.spss[id] = Some(Arc::new(s));
253 }
254 nal_type::PPS => {
255 let p = Pps::parse(&unescape_rbsp(&nal[2..]))?;
256 let id = p.id as usize;
257 self.ppss[id] = Some(Arc::new(p));
258 }
259 nal_type::EOS | nal_type::EOB => {
260 self.submit_pending();
261 self.after_eos = true;
262 }
263 0..=9 | 16..=21 => self.handle_slice(nal, hdr, pts)?,
264 _ => {}
265 }
266 Ok(())
267 }
268
269 fn layout_for(&mut self, pps: &Arc<Pps>, sps: &Arc<Sps>) -> Result<(Arc<Layout>, Scaling)> {
270 if let Some((_, _, l, s)) = self
271 .layout_cache
272 .iter()
273 .find(|(p, s, _, _)| Arc::ptr_eq(p, pps) && Arc::ptr_eq(s, sps))
274 {
275 return Ok((l.clone(), s.clone()));
276 }
277 let layout = Arc::new(Layout::new(sps, pps)?);
278 let scaling = if sps.scaling_list_enabled {
279 let list = pps
280 .scaling_list
281 .as_ref()
282 .or(sps.scaling_list.as_ref())
283 .cloned()
284 .unwrap_or_else(crate::syntax::hevc::params::ScalingList::default_lists);
285 let mut v = Vec::with_capacity(24);
286 for size in 0..4 {
287 for m in 0..6 {
288 v.push(list.factor(size, m));
289 }
290 }
291 Some(Arc::new(v))
292 } else {
293 None
294 };
295 if self.layout_cache.len() > 8 {
296 self.layout_cache.remove(0);
297 }
298 self.layout_cache
299 .push((pps.clone(), sps.clone(), layout.clone(), scaling.clone()));
300 Ok((layout, scaling))
301 }
302
303 fn handle_slice(&mut self, nal: &[u8], hdr: NalHeader, pts: i64) -> Result<()> {
304 let (rbsp, escapes) = unescape_rbsp_marked(&nal[2..]);
305 let first = rbsp.first().is_some_and(|b| b & 0x80 != 0);
306 if first {
307 self.submit_pending();
308 }
309 if hdr.is_rasl() && (self.skip_rasl || self.first_picture) {
310 return Ok(());
311 }
312 let ppss = &self.ppss;
313 let spss = &self.spss;
314 let prev = self.pending.as_ref().map(|p| &p.last_sh);
315 let (sh, pps, sps) = SliceHeader::parse(
316 &rbsp,
317 hdr,
318 |id| {
319 let pps = ppss[id as usize]
320 .as_ref()
321 .ok_or_else(|| Error::MissingParameterSet(format!("PPS {id}")))?;
322 let sps = spss[pps.sps_id as usize]
323 .as_ref()
324 .ok_or_else(|| Error::MissingParameterSet(format!("SPS {}", pps.sps_id)))?;
325 Ok((pps.clone(), sps.clone()))
326 },
327 prev,
328 )?;
329 sps.check_supported()?;
330 pps.check_supported()?;
331 if first {
332 self.start_picture(&sh, &sps, &pps, pts)?;
333 } else {
334 let Some(p) = &self.pending else {
335 return invalid("slice segment without the first slice of its picture");
336 };
337 ensure!(Arc::ptr_eq(&p.pps, &pps), "PPS changed within a picture");
338 }
339 let Some(pending) = self.pending.as_mut() else {
340 return invalid("slice segment without the first slice of its picture");
341 };
342 let refs = build_ref_lists(&sh, &pending.rps)?;
343 if !sh.dependent {
344 pending.last_sh = sh.clone();
345 }
346 pending.slices.push(SliceJob {
347 sh,
348 pps,
349 sps,
350 rbsp,
351 escapes,
352 refs,
353 });
354 Ok(())
355 }
356
357 fn start_picture(
358 &mut self,
359 sh: &SliceHeader,
360 sps: &Arc<Sps>,
361 pps: &Arc<Pps>,
362 pts: i64,
363 ) -> Result<()> {
364 ensure!(
365 sh.first_slice_segment_in_pic && !sh.dependent,
366 "invalid first slice segment"
367 );
368 let hdr = sh.nal;
369 let changed = match &self.active_sps {
370 None => true,
371 Some(a) => {
372 a.width != sps.width
373 || a.height != sps.height
374 || a.bit_depth_luma != sps.bit_depth_luma
375 || a.bit_depth_chroma != sps.bit_depth_chroma
376 || a.log2_ctb != sps.log2_ctb
377 }
378 };
379 let mut outs = Vec::new();
380 if changed && self.active_sps.is_some() {
381 self.dpb.flush(&mut outs);
382 }
383 self.active_sps = Some(sps.clone());
384 self.dpb.max_dec_pic_buffering = sps.max_dec_pic_buffering as usize;
385 self.dpb.max_num_reorder = sps.max_num_reorder as usize;
386 self.dpb.max_latency = if sps.max_latency_increase_plus1 != 0 {
387 sps.max_num_reorder + sps.max_latency_increase_plus1 - 1
388 } else {
389 0
390 };
391 let irap_no_rasl =
392 hdr.is_irap() && (hdr.is_idr() || hdr.is_bla() || self.first_picture || self.after_eos);
393 if hdr.is_irap() {
394 self.skip_rasl = irap_no_rasl;
395 }
396 let max_lsb = sps.max_poc_lsb();
397 let poc = self.poc_state.compute(sh, max_lsb, irap_no_rasl);
398 let (w, h) = (sps.width as usize, sps.height as usize);
399 let template = Frame::new(
400 0,
401 0,
402 w,
403 h,
404 sps.log2_ctb,
405 sps.bit_depth_luma,
406 sps.bit_depth_chroma,
407 );
408 let next_id = &mut self.next_id;
409 let mut make_missing = |p: i32| -> FrameRef {
410 let id = *next_id;
411 *next_id += 1;
412 Arc::new(Frame::gray(id, p, &template))
413 };
414 let rps = self
415 .dpb
416 .apply_rps(sh, poc, max_lsb, irap_no_rasl, &mut make_missing)?;
417 if irap_no_rasl && !self.first_picture {
418 if sh.no_output_of_prior_pics && hdr.is_idr() {
419 self.dpb.clear();
420 } else {
421 self.dpb.flush(&mut outs);
422 }
423 } else {
424 self.dpb.bump_before_decode(&mut outs);
425 }
426 self.out_queue.extend(outs);
427 let (layout, scaling) = self.layout_for(pps, sps)?;
428 let _ = scaling;
429 let id = self.next_id;
430 self.next_id += 1;
431 let frame = Arc::new(
432 Frame::new(
433 id,
434 poc,
435 w,
436 h,
437 sps.log2_ctb,
438 sps.bit_depth_luma,
439 sps.bit_depth_chroma,
440 )
441 .pooled(self.shared.rows.clone()),
442 );
443 let output = sh.pic_output && !(hdr.is_rasl() && self.skip_rasl);
444 let draft = self.draft
445 && hdr.is_sub_layer_non_ref()
446 && hdr.temporal_id as u32 >= sps.max_sub_layers_minus1;
447 let meta = Arc::new(output_meta(sps, pts, draft));
448 self.pending = Some(PendingPic {
449 frame,
450 sps: sps.clone(),
451 pps: pps.clone(),
452 layout,
453 last_sh: sh.clone(),
454 first_sh: sh.clone(),
455 rps,
456 poc,
457 output,
458 meta,
459 slices: Vec::new(),
460 });
461 self.first_picture = false;
462 self.after_eos = false;
463 Ok(())
464 }
465
466 fn submit_pending(&mut self) {
467 let Some(p) = self.pending.take() else { return };
468 let PendingPic {
469 frame,
470 sps,
471 pps,
472 layout,
473 first_sh,
474 poc,
475 output,
476 meta,
477 slices,
478 ..
479 } = p;
480 let scaling = self
481 .layout_cache
482 .iter()
483 .find(|(pp, ss, _, _)| Arc::ptr_eq(pp, &pps) && Arc::ptr_eq(ss, &sps))
484 .and_then(|(_, _, _, s)| s.clone());
485 let out = output.then(|| (meta.clone(), self.time_base));
486 self.dispatch(Job {
487 frame: frame.clone(),
488 sps,
489 pps,
490 layout,
491 scaling,
492 slices,
493 draft: meta.draft,
494 out,
495 helpers: if self.pool.is_some() {
496 self.threads.saturating_sub(1).min(3)
497 } else {
498 0
499 },
500 });
501 self.poc_state.update(&first_sh, poc);
502 let mut outs = Vec::new();
503 self.dpb.insert(frame, poc, output, meta, &mut outs);
504 self.out_queue.extend(outs);
505 }
506
507 fn dispatch(&mut self, job: Job) {
508 let shared = self.shared.clone();
509 if let Some(pool) = &self.pool {
510 self.in_flight.retain(|f| !f.is_complete());
511 let rows = job.pps.entropy_coding_sync && !job.pps.tiles_enabled;
512 let limit = if rows {
513 self.threads / 4 + 1
514 } else {
515 self.max_in_flight
516 };
517 while self.in_flight.len() >= limit {
518 if let Some(f) = self.in_flight.pop_front() {
519 f.wait_complete();
520 }
521 }
522 self.in_flight.push_back(job.frame.clone());
523 pool.spawn_fifo(move || run_job(job, &shared));
524 return;
525 }
526 run_job(job, &shared);
527 }
528}
529
530struct Job {
531 frame: FrameRef,
532 sps: Arc<Sps>,
533 pps: Arc<Pps>,
534 layout: Arc<Layout>,
535 scaling: Scaling,
536 slices: Vec<SliceJob>,
537 draft: bool,
538 out: Option<(Arc<OutputMeta>, TimeBase)>,
539 helpers: usize,
540}
541
542fn run_job(job: Job, shared: &Shared) {
543 let Job {
544 frame,
545 sps,
546 pps,
547 layout,
548 scaling,
549 slices,
550 draft,
551 out,
552 helpers,
553 } = job;
554 let slices = Arc::new(slices);
555 if let Some((meta, _)) = &out {
556 let (cx, cy, cw, ch) = meta.crop;
557 let crop = (cx as usize, cy as usize, cw as usize, ch as usize);
558 frame.set_output(crop, shared.image_buffers());
559 }
560 let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
561 let input = Input {
562 frame: frame.clone(),
563 sps,
564 pps,
565 layout,
566 scaling,
567 slices,
568 draft,
569 helpers,
570 };
571 if input.sps.bit_depth_luma > 8 {
572 decode_picture::<u16>(input, shared.pool16.clone())
573 } else {
574 decode_picture::<u8>(input, shared.pool8.clone())
575 }
576 }));
577 let ready = result.is_ok();
578 let err = match result {
579 Ok(e) => e,
580 Err(_) => {
581 frame.fill_unpublished_gray();
582 Some(Error::Invalid(
583 "internal error while decoding a picture".into(),
584 ))
585 }
586 };
587 if let (Some(e), Ok(mut slot)) = (err, shared.error.lock()) {
588 slot.get_or_insert(e);
589 }
590 if let Some((meta, time_base)) = out {
591 let bufs = frame.take_output();
592 let image = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| match bufs {
593 Some(b) if ready => image(&meta, time_base, b),
594 b => convert(&frame, &meta, time_base, b.unwrap_or_default()),
595 }))
596 .ok();
597 let _ = frame.image.set(Mutex::new(image));
598 }
599}
600
601struct Input {
602 frame: FrameRef,
603 sps: Arc<Sps>,
604 pps: Arc<Pps>,
605 layout: Arc<Layout>,
606 scaling: Scaling,
607 slices: Arc<Vec<SliceJob>>,
608 draft: bool,
609 helpers: usize,
610}
611
612fn decode_picture<P: Pel>(input: Input, pool: Pool<P>) -> Option<Error> {
613 let Input {
614 frame,
615 sps,
616 pps,
617 layout,
618 scaling,
619 slices,
620 draft,
621 helpers,
622 } = input;
623 let pic = Arc::new(PicState::<P>::new(
624 frame,
625 &sps,
626 pps,
627 layout,
628 scaling,
629 &slices,
630 draft,
631 Some(pool),
632 ));
633 if helpers > 0
634 && let Some(run) = SliceRun::plan(&pic, &slices)
635 {
636 let run = Arc::new(run);
637 for _ in 0..helpers.min(run.slices() - 1) {
638 let (pic, slices, run) = (pic.clone(), slices.clone(), run.clone());
639 rayon::spawn(move || run.work(&pic, &slices));
640 }
641 run.work(&pic, &slices);
642 pic.signal.wait_until(|| run.done());
643 crate::hevc::filter::finish(&pic);
644 return run.take_error();
645 }
646 let mut first_err = None;
647 for (i, s) in slices.iter().enumerate() {
648 let plan = if helpers > 0 { s.row_plan(&pic) } else { None };
649 let Some(plan) = plan else {
650 if let Err(e) = SliceDecoder::decode(&pic, s, i) {
651 first_err.get_or_insert(e);
652 }
653 continue;
654 };
655 let plan = Arc::new(plan);
656 for _ in 0..helpers.min(plan.rows() - 1) {
657 let (pic, slices, plan) = (pic.clone(), slices.clone(), plan.clone());
658 rayon::spawn(move || {
659 if let Some(s) = slices.get(i) {
660 plan.work(&pic, s, i);
661 }
662 });
663 }
664 plan.work(&pic, s, i);
665 pic.signal.wait_until(|| plan.done());
666 if let Some(e) = plan.take_error() {
667 first_err.get_or_insert(e);
668 }
669 }
670 crate::hevc::filter::finish(&pic);
671 first_err
672}
673
674fn output_meta(sps: &Sps, pts: i64, draft: bool) -> OutputMeta {
675 OutputMeta {
676 pts,
677 crop: sps.crop_rect(),
678 bit_depth: sps.bit_depth_luma,
679 draft,
680 }
681}
682
683fn make_picture(o: &Output<FrameRef>, time_base: TimeBase) -> Image {
684 let ready = o.frame.image.wait().lock().ok().and_then(|mut s| s.take());
685 ready.unwrap_or_else(|| convert(&o.frame, &o.meta, time_base, Default::default()))
686}
687
688crate::media::isa::dispatch! {
689fn convert(frame: &Frame, meta: &OutputMeta, time_base: TimeBase, bufs: [Vec<u8>; 3]) -> Image {
690 let (cx, cy, cw, ch) = meta.crop;
691 let (cx, cy, cw, ch) = (cx as usize, cy as usize, cw as usize, ch as usize);
692 image(meta, time_base, frame.copy_cropped((cx, cy, cw, ch), bufs))
693}
694}
695
696fn image(meta: &OutputMeta, time_base: TimeBase, [y, u, v]: [Vec<u8>; 3]) -> Image {
697 let (cw, ch) = (meta.crop.2 as usize, meta.crop.3 as usize);
698 let wide = meta.bit_depth > 8;
699 let (ccw, cch) = (cw.div_ceil(2), ch.div_ceil(2));
700 let bps = if wide { 2 } else { 1 };
701 let plane = |data: Vec<u8>, w: usize, h: usize| Plane {
702 data,
703 stride: w * bps,
704 width: w,
705 height: h,
706 };
707 Image {
708 format: if wide {
709 PixelFormat::Yuv420p10
710 } else {
711 PixelFormat::Yuv420p8
712 },
713 width: cw as u32,
714 height: ch as u32,
715 pts: meta.pts,
716 time_base,
717 planes: [plane(y, cw, ch), plane(u, ccw, cch), plane(v, ccw, cch)],
718 }
719}