encoder: Make StatelessEncoder codec agnostic

examples/ccenc.rs

@@ -13,10 +13,11 @@ use argh::FromArgs;
 use cros_codecs::backend::vaapi::surface_pool::VaSurfacePool;
 use cros_codecs::codec::h264::parser::Profile;
 use cros_codecs::decoder::FramePool;
-use cros_codecs::encoder::Bitrate;
 use cros_codecs::encoder::stateless::h264::EncoderConfig;
-use cros_codecs::encoder::stateless::h264::StatelessEncoder;
+use cros_codecs::encoder::stateless::h264::H264;
+use cros_codecs::encoder::stateless::StatelessEncoder;
 use cros_codecs::encoder::stateless::StatelessVideoEncoder;
+use cros_codecs::encoder::Bitrate;
 use cros_codecs::encoder::FrameMetadata;
 use cros_codecs::BlockingMode;
 use cros_codecs::Fourcc;
@@ -157,7 +158,7 @@ fn main() {
 
     let display = libva::Display::open().unwrap();
     let fourcc = b"NV12".into();
-    let mut encoder = StatelessEncoder::new_vaapi(
+    let mut encoder = StatelessEncoder::<H264, _, _>::new_vaapi(
         Rc::clone(&display),
         config,
         fourcc,

src/backend/vaapi/encoder.rs

@@ -28,7 +28,7 @@ use crate::decoder::FramePool;
 use crate::encoder::stateless::BackendPromise;
 use crate::encoder::stateless::StatelessBackendError;
 use crate::encoder::stateless::StatelessBackendResult;
-use crate::encoder::stateless::StatelessVideoEncoderBackend;
+use crate::encoder::stateless::StatelessEncoderBackendImport;
 use crate::encoder::FrameMetadata;
 use crate::Fourcc;
 use crate::Resolution;
@@ -45,7 +45,6 @@ impl From<libva::VaError> for StatelessBackendError {
     }
 }
 
-
 pub struct Reconstructed(PooledVaSurface<()>);
 
 impl Reconstructed {
@@ -62,7 +61,7 @@ impl Reconstructed {
 pub struct VaapiBackend<M, H>
 where
     M: SurfaceMemoryDescriptor,
-    H: std::borrow::Borrow<Surface<M>>,
+    H: std::borrow::Borrow<Surface<M>> + 'static,
 {
     /// VA config.
     #[allow(dead_code)]
@@ -171,18 +170,16 @@ where
     }
 }
 
-impl<M, H> StatelessVideoEncoderBackend<H> for VaapiBackend<M, H>
+impl<M, H> StatelessEncoderBackendImport<H, H> for VaapiBackend<M, H>
 where
     M: SurfaceMemoryDescriptor,
-    H: std::borrow::Borrow<Surface<M>>,
+    H: std::borrow::Borrow<Surface<M>> + 'static,
 {
-    type Picture = H;
-
     fn import_picture(
         &mut self,
         _metadata: &FrameMetadata,
         handle: H,
-    ) -> StatelessBackendResult<Self::Picture> {
+    ) -> StatelessBackendResult<H> {
         Ok(handle)
     }
 }

src/encoder/stateless.rs

@@ -121,6 +121,35 @@ where
     }
 }
 
+/// Generic trait for stateless encoder backends
+pub trait StatelessVideoEncoderBackend<Codec: StatelessCodec> {
+    /// Backend's specific representation of the input frame, transformed with [`import_picture`].
+    /// Might be a wrapper of the input handle with additional backend specific data or a copy of
+    /// an input frame in internal backend's representation.
+    ///
+    /// [`import_picture`]: StatelessVideoEncoderBackend::import_picture
+    type Picture: 'static;
+
+    /// Backend's reconstructed frame handle.
+    type Reconstructed: 'static;
+
+    /// Backend's speficic [`BackendPromise`] for bitstream, a result of [`Request`] submission.
+    type CodedPromise: BackendPromise<Output = Vec<u8>>;
+
+    /// Backend's speficic [`BackendPromise`] for [`StatelessVideoEncoderBackend::Reconstructed`],
+    /// a result of [`Request`] submission.
+    type ReconPromise: BackendPromise<Output = Self::Reconstructed>;
+}
+
+pub trait StatelessEncoderBackendImport<Handle, Picture> {
+    /// Imports the input [`Handle`] from client and transforms into [`Picture`]
+    fn import_picture(
+        &mut self,
+        metadata: &FrameMetadata,
+        handle: Handle,
+    ) -> StatelessBackendResult<Picture>;
+}
+
 /// Predictor is responsible for yielding stream parameter sets and creating requests to backend.
 /// It accepts the frames and reconstructed frames and returns [`Request`]s for
 /// [`StatelessEncoderExecute`] to execute. For example [`Predictor`] may hold frames from
@@ -145,33 +174,43 @@ pub(super) trait Predictor<Picture, Reference, Request> {
     fn drain(&mut self) -> EncodeResult<Vec<Request>>;
 }
 
-/// Generic trait for stateless encoder backends
-pub trait StatelessVideoEncoderBackend<H> {
-    /// Backend's specific representation of the input frame, transformed with [`import_picture`].
-    /// Might be a wrapper of the input handle with additional backend specific data or a copy of
-    /// an input frame in internal backend's representation.
-    ///
-    /// [`import_picture`]: StatelessVideoEncoderBackend::import_picture
-    type Picture;
+/// Trait helping contain all codec specific and backend specific types
+pub trait StatelessCodecSpecific<Codec, Backend>
+where
+    Codec: StatelessCodec,
+    Backend: StatelessVideoEncoderBackend<Codec>,
+{
+    /// Codec specific representation of frame reference wrapping a backend reference type
+    /// containing a codec specific frame metadata
+    type Reference;
 
-    /// Imports the input handle from client and transforms into [`Picture`]
-    ///
-    /// [`Picture`]: StatelessVideoEncoderBackend::Picture
-    fn import_picture(
-        &mut self,
-        metadata: &FrameMetadata,
-        handle: H,
-    ) -> StatelessBackendResult<Self::Picture>;
+    /// A request type that will be delivered to codec specific stateless encoder backend
+    type Request;
+
+    /// Codec specific [`BackendPromise`] for [`CodedBitstreamBuffer`] wrapping a backend specific
+    /// [`StatelessVideoEncoderBackend::CodedPromise`]
+    type CodedPromise: BackendPromise<Output = CodedBitstreamBuffer>;
+
+    /// Codec specific [`BackendPromise`] for [`StatelessCodecSpecific::Reference`] wrapping a
+    /// backend speficic [`StatelessVideoEncoderBackend::ReconPromise`]
+    type ReferencePromise: BackendPromise<Output = Self::Reference>;
+}
+
+pub trait StatelessCodec: Sized {
+    /// Codec and backend specific types container
+    type Specific<Backend>: StatelessCodecSpecific<Self, Backend>
+    where
+        Backend: StatelessVideoEncoderBackend<Self>;
 }
 
 /// Stateless video encoder interface.
-pub trait StatelessVideoEncoder<H> {
+pub trait StatelessVideoEncoder<Handle> {
     /// Enqueues the frame for encoding. The implementation will drop the handle after it is no
     /// longer be needed. The encoder is not required to immediately start processing the frame
     /// and yield output bitstream. It is allowed to hold frames until certain conditions are met
     /// eg. for specified prediction structures or referencing in order to further optimize
     /// the compression rate of the bitstream.
-    fn encode(&mut self, meta: FrameMetadata, handle: H) -> Result<(), EncodeError>;
+    fn encode(&mut self, meta: FrameMetadata, handle: Handle) -> Result<(), EncodeError>;
 
     /// Drains the encoder. This means that encoder is required to finish processing of all the
     /// frames in the internal queue and yield output bitstream by the end of the call. The output
@@ -186,7 +225,7 @@ pub trait StatelessVideoEncoder<H> {
 
     /// Polls on the encoder for the available output bitstream with compressed frames that where
     /// submitted with [`encode`].
-    ///
+    //kk/
     /// The call may also trigger a further processing aside of returning output. Therefore it
     /// *recommended* that this function is called frequently.
     ///
@@ -217,3 +256,169 @@ where
 
     Ok(())
 }
+
+/// Helper aliases for codec and backend specific types
+type Picture<C, B> = <B as StatelessVideoEncoderBackend<C>>::Picture;
+
+type CodecSpecific<C, B> = <C as StatelessCodec>::Specific<B>;
+
+type Reference<C, B> = <CodecSpecific<C, B> as StatelessCodecSpecific<C, B>>::Reference;
+
+type Request<C, B> = <CodecSpecific<C, B> as StatelessCodecSpecific<C, B>>::Request;
+
+type CodedPromise<C, B> = <CodecSpecific<C, B> as StatelessCodecSpecific<C, B>>::CodedPromise;
+
+type ReferencePromise<C, B> =
+    <CodecSpecific<C, B> as StatelessCodecSpecific<C, B>>::ReferencePromise;
+
+type BoxPredictor<C, B> = Box<dyn Predictor<Picture<C, B>, Reference<C, B>, Request<C, B>>>;
+
+pub struct StatelessEncoder<Codec, Handle, Backend>
+where
+    Backend: StatelessVideoEncoderBackend<Codec>,
+    Codec: StatelessCodec,
+{
+    /// Pending frame output promise queue
+    output_queue: OutputQueue<CodedPromise<Codec, Backend>>,
+
+    /// Pending reconstructed pictures promise queue
+    recon_queue: OutputQueue<ReferencePromise<Codec, Backend>>,
+
+    /// [`Predictor`] instance responsible for the encoder decision making
+    predictor: BoxPredictor<Codec, Backend>,
+
+    // predictor: Box<dyn Predictor<B::Picture, B::Reference>>,
+    coded_queue: VecDeque<CodedBitstreamBuffer>,
+
+    /// Number of the currently held frames by the predictor
+    predictor_frame_count: usize,
+
+    /// [`StatelessVP9EncoderBackend`] instance to delegate [`BackendRequest`] to
+    backend: Backend,
+
+    _phantom: std::marker::PhantomData<Handle>,
+}
+
+/// A bridge trait between [`StatelessEncoder`] and codec specific backend trait (eg.
+/// [`h264::StatelessH264EncoderBackend`] or [`vp9::StatelessVP9EncoderBackend`]).
+/// Accepts [`Request`] and is responsible for adding resutling [`BackendPromise`] to
+/// [`StatelessEncoder`] internal queues and  decrementing the internal predictor frame counter if
+/// the backend moved the frame outside predictor ownership.
+pub trait StatelessEncoderExecute<Codec, Handle, Backend>
+where
+    Backend: StatelessVideoEncoderBackend<Codec>,
+    Codec: StatelessCodec,
+{
+    fn execute(&mut self, request: Request<Codec, Backend>) -> EncodeResult<()>;
+}
+
+impl<Codec, Handle, Backend> StatelessEncoder<Codec, Handle, Backend>
+where
+    Codec: StatelessCodec,
+    Backend: StatelessVideoEncoderBackend<Codec>,
+    Self: StatelessEncoderExecute<Codec, Handle, Backend>,
+{
+    fn new(
+        backend: Backend,
+        mode: BlockingMode,
+        predictor: BoxPredictor<Codec, Backend>,
+    ) -> EncodeResult<Self> {
+        Ok(Self {
+            backend,
+            predictor,
+            predictor_frame_count: 0,
+            coded_queue: Default::default(),
+            output_queue: OutputQueue::new(mode),
+            recon_queue: OutputQueue::new(mode),
+            _phantom: Default::default(),
+        })
+    }
+
+    fn poll_pending(&mut self, mode: BlockingMode) -> EncodeResult<()> {
+        // Poll the output queue once and then continue polling while new promise is submitted
+        while let Some(coded) = self.output_queue.poll(mode)? {
+            self.coded_queue.push_back(coded);
+        }
+
+        while let Some(recon) = self.recon_queue.poll(mode)? {
+            let requests = self.predictor.reconstructed(recon)?;
+            if requests.is_empty() {
+                // No promise was submitted, therefore break
+                break;
+            }
+
+            for request in requests {
+                self.execute(request)?;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl<Codec, Handle, Backend> StatelessVideoEncoder<Handle>
+    for StatelessEncoder<Codec, Handle, Backend>
+where
+    Codec: StatelessCodec,
+    Backend: StatelessVideoEncoderBackend<Codec>,
+    Backend: StatelessEncoderBackendImport<Handle, Backend::Picture>,
+    Self: StatelessEncoderExecute<Codec, Handle, Backend>,
+{
+    fn encode(&mut self, metadata: FrameMetadata, handle: Handle) -> EncodeResult<()> {
+        log::trace!(
+            "encode: timestamp={} layout={:?}",
+            metadata.timestamp,
+            metadata.layout
+        );
+
+        // Import `handle` to backends representation
+        let backend_pic = self.backend.import_picture(&metadata, handle)?;
+
+        // Increase the number of frames that predictor holds, before handing one to it
+        self.predictor_frame_count += 1;
+
+        // Ask predictor to decide on the next move and execute it
+        let requests = self.predictor.new_frame(backend_pic, metadata)?;
+        for request in requests {
+            self.execute(request)?;
+        }
+
+        Ok(())
+    }
+
+    fn drain(&mut self) -> EncodeResult<()> {
+        log::trace!("currently predictor holds {}", self.predictor_frame_count);
+
+        // Drain the predictor
+        while self.predictor_frame_count > 0 || !self.recon_queue.is_empty() {
+            if self.output_queue.is_empty() && self.recon_queue.is_empty() {
+                // The OutputQueue is empty and predictor holds frames, force it to yield a request
+                // to empty it's internal queue.
+                let requests = self.predictor.drain()?;
+                if requests.is_empty() {
+                    log::error!("failed to drain predictor, no request was returned");
+                    return Err(EncodeError::InvalidInternalState);
+                }
+
+                for request in requests {
+                    self.execute(request)?;
+                }
+            }
+
+            self.poll_pending(BlockingMode::Blocking)?;
+        }
+
+        // There are still some requests being processed. Continue on polling them.
+        while !self.output_queue.is_empty() {
+            self.poll_pending(BlockingMode::Blocking)?;
+        }
+
+        Ok(())
+    }
+
+    fn poll(&mut self) -> EncodeResult<Option<CodedBitstreamBuffer>> {
+        // Poll on output queue without blocking and try to dueue from coded queue
+        self.poll_pending(BlockingMode::NonBlocking)?;
+        Ok(self.coded_queue.pop_front())
+    }
+}