1use alloc::boxed::Box;
23use alloc::collections::VecDeque;
24use alloc::sync::Arc;
25use alloc::vec;
26use alloc::vec::Vec;
27use core::fmt;
28
29use zencodec::decode::{
30 AnimationFrame, AnimationFrameDecoder, DecodeRowSink, OutputInfo, SinkError,
31};
32use zencodec::encode::{AnimationFrameEncoder, AnimationOutputTiming, BoxedError, EncodeOutput};
33use zencodec::time::{FrameDuration, Timestamp, TimingError, duration_to_ticks};
34use zencodec::{ImageFormat, ImageInfo, ImageSequence, UnsupportedOperation};
35use zencodec_io::{DriveFailure, SliceCursor};
36use zenpixels::{PixelBuffer, PixelDescriptor, PixelSlice, PixelSliceMut};
37
38use crate::container::{DemuxEvent, DemuxStage, Demuxer, MuxHint, Muxer};
39use crate::frame::{
40 ChromaSiting, Frame, FrameGeometry, FrameInterpretation, FrameTiming, MapError, OwnedPlane,
41 OwnedPlanes, PlaneData, PlaneRole, PlaneSamples, Presentation, SampleEncoding, Subsampling,
42};
43use crate::session::{
44 CodecConfig, ProtocolViolation, PumpError, SessionError, VideoDecoder, VideoEncoder,
45 drain_decoder, drain_encoder, send_frame_until_accepted, send_packet_until_accepted,
46};
47use crate::track::{TrackInfo, TrackKind};
48
49pub trait RenderFrame {
52 fn canvas_descriptor(&self, track: &TrackInfo) -> Result<PixelDescriptor, BoxedError>;
58
59 fn render(&mut self, frame: &Frame, canvas: &mut PixelSliceMut<'_>) -> Result<(), BoxedError>;
66}
67
68pub trait CaptureFrame {
71 fn capture(
78 &mut self,
79 pixels: &PixelSlice<'_>,
80 config: &CodecConfig,
81 timing: FrameTiming,
82 ) -> Result<Frame, BoxedError>;
83}
84
85#[derive(Clone, Debug, Default)]
88pub struct Gray8 {
89 interpretation: Option<Arc<FrameInterpretation>>,
90}
91
92#[derive(Debug)]
93struct ConvertError(&'static str);
94
95impl fmt::Display for ConvertError {
96 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97 f.write_str(self.0)
98 }
99}
100
101impl core::error::Error for ConvertError {}
102
103fn convert_error(msg: &'static str) -> BoxedError {
104 Box::new(ConvertError(msg))
105}
106
107impl RenderFrame for Gray8 {
108 fn canvas_descriptor(&self, track: &TrackInfo) -> Result<PixelDescriptor, BoxedError> {
109 if track.video_info().is_none() {
110 return Err(convert_error("Gray8 renders video tracks only"));
111 }
112 Ok(PixelDescriptor::GRAY8_SRGB)
113 }
114
115 fn render(&mut self, frame: &Frame, canvas: &mut PixelSliceMut<'_>) -> Result<(), BoxedError> {
116 let interp = frame.interpretation();
117 if interp.subsampling != Subsampling::Mono || interp.encoding != SampleEncoding::U8 {
118 return Err(convert_error("Gray8 renders 8-bit monochrome frames only"));
119 }
120 let mapped = frame.map()?;
121 let plane = mapped
122 .plane(0)
123 .ok_or_else(|| convert_error("frame has no plane"))?;
124 if plane.width() != canvas.width() as usize || plane.height() != canvas.rows() as usize {
125 return Err(convert_error("frame size differs from the canvas"));
126 }
127 for y in 0..plane.height() {
128 let PlaneSamples::U8(row) = plane.row(y).expect("in range") else {
129 return Err(convert_error("8-bit frame with 16-bit storage"));
130 };
131 canvas.row_mut(y as u32)[..row.len()].copy_from_slice(row);
132 }
133 Ok(())
134 }
135}
136
137impl CaptureFrame for Gray8 {
138 fn capture(
139 &mut self,
140 pixels: &PixelSlice<'_>,
141 config: &CodecConfig,
142 timing: FrameTiming,
143 ) -> Result<Frame, BoxedError> {
144 if pixels.descriptor() != PixelDescriptor::GRAY8_SRGB {
145 return Err(Box::new(UnsupportedOperation::PixelFormat));
146 }
147 let (w, h) = (pixels.width(), pixels.rows());
148 if config
149 .video
150 .is_some_and(|v| (v.coded_width, v.coded_height) != (w, h))
151 {
152 return Err(convert_error("frame size differs from the encoder's"));
153 }
154 let interpretation = match &self.interpretation {
155 Some(i) if i.geometry == FrameGeometry::full(w, h) => Arc::clone(i),
156 _ => {
157 let i = Arc::new(FrameInterpretation::new(
158 FrameGeometry::full(w, h),
159 Subsampling::Mono,
160 ChromaSiting::Unknown,
161 SampleEncoding::U8,
162 vec![PlaneRole::Luma],
163 ));
164 self.interpretation = Some(Arc::clone(&i));
165 i
166 }
167 };
168 let (wu, hu) = (w as usize, h as usize);
169 let mut data = Vec::new();
170 data.try_reserve_exact(wu * hu)
171 .map_err(|_| convert_error("frame does not fit in memory"))?;
172 for y in 0..h {
173 data.extend_from_slice(&pixels.row(y)[..wu]);
174 }
175 let planes = OwnedPlanes::new(
176 vec![OwnedPlane {
177 data: PlaneData::U8(data),
178 width: wu,
179 height: hu,
180 stride_bytes: wu,
181 }],
182 SampleEncoding::U8,
183 )?;
184 Ok(Frame::new(
185 Arc::new(planes),
186 interpretation,
187 timing,
188 Presentation::default(),
189 ))
190 }
191}
192
193#[derive(Debug)]
196#[non_exhaustive]
197pub enum AdaptError<A, B> {
198 Demux(DriveFailure<SessionError<A>, core::convert::Infallible>),
200 Decode(SessionError<B>),
202 Encode(SessionError<A>),
204 Mux(SessionError<B>),
206 Time(TimingError),
208 Map(MapError),
210 Convert(BoxedError),
212 Sink(SinkError),
214 Unsupported(UnsupportedOperation),
216 MissingTimestamp,
219 Shape,
221 FrameCountMismatch,
223 PatchesNotAllowed,
225}
226
227impl<A: fmt::Display, B: fmt::Display> fmt::Display for AdaptError<A, B> {
228 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
229 match self {
230 Self::Demux(e) => write!(f, "demux: {e}"),
231 Self::Decode(e) => write!(f, "decode: {e}"),
232 Self::Encode(e) => write!(f, "encode: {e}"),
233 Self::Mux(e) => write!(f, "mux: {e}"),
234 Self::Time(e) => write!(f, "time: {e}"),
235 Self::Map(e) => write!(f, "map: {e}"),
236 Self::Convert(e) => write!(f, "convert: {e}"),
237 Self::Sink(e) => write!(f, "sink: {e}"),
238 Self::Unsupported(op) => write!(f, "unsupported: {op}"),
239 Self::MissingTimestamp => f.write_str("video frame has no presentation timestamp"),
240 Self::Shape => {
241 f.write_str("no video track, or frame geometry differs from the animation canvas")
242 }
243 Self::FrameCountMismatch => {
244 f.write_str("finished after a different number of frames than declared")
245 }
246 Self::PatchesNotAllowed => {
247 f.write_str("muxer produced patches for a forward-only destination")
248 }
249 }
250 }
251}
252
253impl<A, B> core::error::Error for AdaptError<A, B>
254where
255 A: core::error::Error + 'static,
256 B: core::error::Error + 'static,
257{
258 fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
259 match self {
260 Self::Demux(e) => Some(e),
261 Self::Decode(e) | Self::Mux(e) => Some(e),
262 Self::Encode(e) => Some(e),
263 Self::Time(e) => Some(e),
264 Self::Map(e) => Some(e),
265 Self::Convert(e) => Some(&**e),
266 Self::Sink(e) => Some(&**e),
267 Self::Unsupported(e) => Some(e),
268 _ => None,
269 }
270 }
271}
272
273type DecErr<D, C> = AdaptError<<D as Demuxer>::Error, <C as VideoDecoder>::Error>;
274
275pub struct VideoAnimationDecoder<D, C, R> {
278 bytes: Arc<[u8]>,
279 cursor: SliceCursor,
280 dem: D,
281 dec: C,
282 render: R,
283 track: Option<u32>,
284 pending: VecDeque<Frame>,
285 eos: bool,
286 final_duration: FrameDuration,
287 canvas: PixelBuffer,
288 info: ImageInfo,
289 index: u32,
290 frames_decoded: u64,
291}
292
293impl<D: Demuxer, C: VideoDecoder, R: RenderFrame> VideoAnimationDecoder<D, C, R> {
294 pub fn new(
299 bytes: Arc<[u8]>,
300 dem: D,
301 dec: C,
302 render: R,
303 final_duration: FrameDuration,
304 window: usize,
305 ) -> Self {
306 Self {
307 bytes,
308 cursor: SliceCursor::new(window),
309 dem,
310 dec,
311 render,
312 track: None,
313 pending: VecDeque::new(),
314 eos: false,
315 final_duration,
316 canvas: PixelBuffer::new(1, 1, PixelDescriptor::GRAY8_SRGB),
317 info: ImageInfo::new(1, 1, ImageFormat::Unknown),
318 index: 0,
319 frames_decoded: 0,
320 }
321 }
322
323 pub fn frames_decoded(&self) -> u64 {
326 self.frames_decoded
327 }
328
329 pub fn demuxer(&self) -> &D {
331 &self.dem
332 }
333
334 fn on_headers(&mut self, pulse: Option<&dyn zencodec::Pulse>) -> Result<(), DecErr<D, C>> {
335 let tracks = self.dem.tracks();
336 let track = tracks
337 .iter()
338 .find(|t| t.kind() == TrackKind::Video)
339 .ok_or(AdaptError::Shape)?;
340 let video = track.video_info().ok_or(AdaptError::Shape)?;
341 let desc = self
342 .render
343 .canvas_descriptor(track)
344 .map_err(AdaptError::Convert)?;
345 self.canvas = PixelBuffer::try_new(video.coded_width, video.coded_height, desc)
346 .map_err(|_| AdaptError::Shape)?;
347 self.info = ImageInfo::new(video.coded_width, video.coded_height, ImageFormat::Unknown)
348 .with_sequence(ImageSequence::Animation {
349 frame_count: track.declared_packets().and_then(|n| u32::try_from(n).ok()),
350 loop_count: Some(1),
351 random_access: false,
352 });
353 self.track = Some(track.index());
354 self.dec
355 .configure(&CodecConfig::from(track), pulse)
356 .map_err(AdaptError::Decode)
357 }
358
359 fn fill(
360 &mut self,
361 min: usize,
362 pulse: Option<&dyn zencodec::Pulse>,
363 ) -> Result<(), DecErr<D, C>> {
364 while self.pending.len() < min && !self.eos {
365 let ev = self
366 .cursor
367 .next(&self.bytes, &mut DemuxStage(&mut self.dem), pulse)
368 .map_err(AdaptError::Demux)?
369 .into_inner();
370 match ev {
371 DemuxEvent::Headers => self.on_headers(pulse)?,
372 DemuxEvent::ConfigChanged { track, .. } if Some(track) == self.track => {
373 let t = self
374 .dem
375 .tracks()
376 .iter()
377 .find(|t| t.index() == track)
378 .ok_or(AdaptError::Shape)?;
379 self.dec
380 .configure(&CodecConfig::from(t), pulse)
381 .map_err(AdaptError::Decode)?;
382 }
383 DemuxEvent::Packet(mut p) if Some(p.track()) == self.track => {
384 let Self {
385 dec,
386 pending,
387 frames_decoded,
388 ..
389 } = self;
390 send_packet_until_accepted(dec, &mut p, pulse, |f| {
391 *frames_decoded += 1;
392 pending.push_back(f);
393 Ok::<(), SessionError<C::Error>>(())
394 })
395 .map_err(|e| match e {
396 PumpError::Upstream(e) | PumpError::Downstream(e) => AdaptError::Decode(e),
397 })?;
398 }
399 DemuxEvent::Done => {
400 let Self {
401 dec,
402 pending,
403 frames_decoded,
404 ..
405 } = self;
406 drain_decoder(dec, pulse, |f| {
407 *frames_decoded += 1;
408 pending.push_back(f);
409 Ok::<(), SessionError<C::Error>>(())
410 })
411 .map_err(|e| match e {
412 PumpError::Upstream(e) | PumpError::Downstream(e) => AdaptError::Decode(e),
413 })?;
414 self.eos = true;
415 }
416 _ => {}
418 }
419 }
420 Ok(())
421 }
422}
423
424impl<D, C, R> AnimationFrameDecoder for VideoAnimationDecoder<D, C, R>
425where
426 D: Demuxer + Send,
427 C: VideoDecoder + Send,
428 R: RenderFrame + Send,
429{
430 type Error = DecErr<D, C>;
431
432 fn wrap_sink_error(err: SinkError) -> Self::Error {
433 AdaptError::Sink(err)
434 }
435
436 fn info(&self) -> &ImageInfo {
437 &self.info
438 }
439
440 fn render_next_frame(
441 &mut self,
442 pulse: Option<&dyn zencodec::Pulse>,
443 ) -> Result<Option<AnimationFrame<'_>>, Self::Error> {
444 self.fill(2, pulse)?;
446 let Some(frame) = self.pending.pop_front() else {
447 return Ok(None);
448 };
449 let pts = frame.timing().pts.ok_or(AdaptError::MissingTimestamp)?;
450 let duration = match self.pending.front() {
451 Some(next) => {
452 let next_pts = next.timing().pts.ok_or(AdaptError::MissingTimestamp)?;
453 pts.duration_until(next_pts).map_err(AdaptError::Time)?
454 }
455 None => self.final_duration,
456 };
457 self.render
458 .render(&frame, &mut self.canvas.as_slice_mut())
459 .map_err(AdaptError::Convert)?;
460 let index = self.index;
461 self.index += 1;
462 Ok(Some(
463 AnimationFrame::with_duration(self.canvas.as_slice(), duration, index)
464 .with_pts(Some(pts)),
465 ))
466 }
467
468 fn render_next_frame_to_sink(
469 &mut self,
470 pulse: Option<&dyn zencodec::Pulse>,
471 sink: &mut dyn DecodeRowSink,
472 ) -> Result<Option<OutputInfo>, Self::Error> {
473 zencodec::helpers::copy_frame_to_sink(self, pulse, sink)
474 }
475}
476
477type EncErr<E, M> = AdaptError<<E as VideoEncoder>::Error, <M as Muxer>::Error>;
478
479pub struct AnimationVideoEncoder<E, M, P> {
485 enc: E,
486 mux: M,
487 capture: P,
488 next_pts: Timestamp,
489 begun: bool,
490 pushed: u32,
491 declared: Option<u32>,
492 lag_frames: u32,
493}
494
495impl<E: VideoEncoder, M: Muxer, P: CaptureFrame> AnimationVideoEncoder<E, M, P> {
496 pub fn new(enc: E, mux: M, capture: P, lag_frames: u32) -> Result<Self, EncErr<E, M>> {
505 let cfg = enc.output_config().ok_or(AdaptError::Shape)?;
506 let next_pts = Timestamp::new(0, cfg.time_base);
507 Ok(Self {
508 enc,
509 mux,
510 capture,
511 next_pts,
512 begun: false,
513 pushed: 0,
514 declared: None,
515 lag_frames,
516 })
517 }
518
519 fn begin_inner(
520 &mut self,
521 width: u32,
522 height: u32,
523 frame_count: Option<u32>,
524 ) -> Result<AnimationOutputTiming, EncErr<E, M>> {
525 if self.begun {
526 return Err(AdaptError::Mux(SessionError::Protocol(
527 ProtocolViolation::BeginTwice,
528 )));
529 }
530 let cfg = self.enc.output_config().ok_or(AdaptError::Shape)?;
531 let video = cfg.video.ok_or(AdaptError::Shape)?;
532 if (video.coded_width, video.coded_height) != (width, height) {
533 return Err(AdaptError::Shape);
534 }
535 let track = TrackInfo::video(0, cfg.codec.clone(), cfg.time_base, video)
536 .build()
537 .map_err(|_| AdaptError::Shape)?;
538 let timing = self
539 .mux
540 .begin(&[track], MuxHint::forward_only())
541 .map_err(AdaptError::Mux)?;
542 self.begun = true;
543 self.declared = frame_count;
544 Ok(match frame_count {
545 Some(n) if self.lag_frames >= n => AnimationOutputTiming::AT_FINISH,
546 _ if timing.streams() => AnimationOutputTiming::streaming(self.lag_frames),
547 _ => AnimationOutputTiming::AT_FINISH,
548 })
549 }
550}
551
552impl<E, M, P> AnimationFrameEncoder for AnimationVideoEncoder<E, M, P>
553where
554 E: VideoEncoder + Send,
555 M: Muxer + Send,
556 P: CaptureFrame + Send,
557{
558 type Error = EncErr<E, M>;
559
560 fn reject(op: UnsupportedOperation) -> Self::Error {
561 AdaptError::Unsupported(op)
562 }
563
564 fn begin(
565 &mut self,
566 width: u32,
567 height: u32,
568 _descriptor: PixelDescriptor,
569 frame_count: Option<u32>,
570 ) -> Result<AnimationOutputTiming, Self::Error> {
571 self.begin_inner(width, height, frame_count)
572 }
573
574 fn pending_output(&self) -> &[u8] {
575 self.mux.pending_output()
576 }
577
578 fn consume_output(&mut self, n: usize) {
579 self.mux.consume_output(n)
580 }
581
582 fn take_output(&mut self, into: &mut Vec<u8>) {
583 self.mux.take_output(into)
584 }
585
586 fn push_frame(
587 &mut self,
588 pixels: PixelSlice<'_>,
589 duration_ms: u32,
590 pulse: Option<&dyn zencodec::Pulse>,
591 ) -> Result<(), Self::Error> {
592 self.push_frame_exact(pixels, FrameDuration::from_millis(duration_ms), pulse)
593 }
594
595 fn push_frame_exact(
596 &mut self,
597 pixels: PixelSlice<'_>,
598 duration: FrameDuration,
599 pulse: Option<&dyn zencodec::Pulse>,
600 ) -> Result<(), Self::Error> {
601 if self.declared.is_some_and(|n| self.pushed >= n) {
602 return Err(AdaptError::FrameCountMismatch);
603 }
604 let ticks = duration_to_ticks(duration, self.next_pts.base())
605 .map_err(|_| AdaptError::Unsupported(UnsupportedOperation::AnimationTiming))?;
606 let next = self
607 .next_pts
608 .checked_add(duration)
609 .map_err(AdaptError::Time)?;
610 if !self.begun {
611 self.begin_inner(pixels.width(), pixels.rows(), None)?;
612 }
613 let cfg = self.enc.output_config().ok_or(AdaptError::Shape)?.clone();
614 let frame = self
615 .capture
616 .capture(
617 &pixels,
618 &cfg,
619 FrameTiming::at(self.next_pts).with_duration_ticks(ticks),
620 )
621 .map_err(AdaptError::Convert)?;
622 let mux = &mut self.mux;
623 send_frame_until_accepted(&mut self.enc, &frame, pulse, |p| mux.write_packet(p, pulse))
624 .map_err(|e| match e {
625 PumpError::Upstream(e) => AdaptError::Encode(e),
626 PumpError::Downstream(e) => AdaptError::Mux(e),
627 })?;
628 self.next_pts = next;
629 self.pushed += 1;
630 Ok(())
631 }
632
633 fn finish(mut self, pulse: Option<&dyn zencodec::Pulse>) -> Result<EncodeOutput, Self::Error> {
634 if !self.begun {
635 return Err(AdaptError::Mux(SessionError::Protocol(
636 ProtocolViolation::NoBegin,
637 )));
638 }
639 if self.declared.is_some_and(|n| n != self.pushed) {
640 return Err(AdaptError::FrameCountMismatch);
641 }
642 let mux = &mut self.mux;
643 drain_encoder(&mut self.enc, pulse, |p| mux.write_packet(p, pulse)).map_err(
644 |e| match e {
645 PumpError::Upstream(e) => AdaptError::Encode(e),
646 PumpError::Downstream(e) => AdaptError::Mux(e),
647 },
648 )?;
649 let fin = self.mux.finish(pulse).map_err(AdaptError::Mux)?;
650 if !fin.patches.is_empty() {
651 return Err(AdaptError::PatchesNotAllowed);
652 }
653 Ok(EncodeOutput::new(fin.tail, ImageFormat::Unknown))
654 }
655}