From b36906e8e9b22a1321de11ad9b6668a800acac90 Mon Sep 17 00:00:00 2001
From: Kostya Shishkov <kostya.shishkov@gmail.com>
Date: Sun, 4 Jun 2023 15:22:49 +0200
Subject: [PATCH] raw video "encoder"

---
 nihav-commonfmt/Cargo.toml                |   3 +-
 nihav-commonfmt/src/codecs/mod.rs         |   4 +
 nihav-commonfmt/src/codecs/rawvideoenc.rs | 100 ++++++++++++++++++++++
 3 files changed, 106 insertions(+), 1 deletion(-)
 create mode 100644 nihav-commonfmt/src/codecs/rawvideoenc.rs

diff --git a/nihav-commonfmt/Cargo.toml b/nihav-commonfmt/Cargo.toml
index 6e7d020..adea42f 100644
--- a/nihav-commonfmt/Cargo.toml
+++ b/nihav-commonfmt/Cargo.toml
@@ -52,8 +52,9 @@ decoder_aac = ["decoders"]
 
 all_encoders = ["all_video_encoders", "all_audio_encoders"]
 
-all_video_encoders = ["encoder_cinepak", "encoder_zmbv"]
+all_video_encoders = ["encoder_cinepak", "encoder_rawvideo", "encoder_zmbv"]
 encoder_cinepak = ["encoders"]
+encoder_rawvideo = ["encoders"]
 encoder_zmbv = ["encoders"]
 
 all_audio_encoders = ["encoder_pcm"]
diff --git a/nihav-commonfmt/src/codecs/mod.rs b/nihav-commonfmt/src/codecs/mod.rs
index f69fc64..4c6569e 100644
--- a/nihav-commonfmt/src/codecs/mod.rs
+++ b/nihav-commonfmt/src/codecs/mod.rs
@@ -77,6 +77,8 @@ pub fn generic_register_all_decoders(rd: &mut RegisteredDecoders) {
 
 #[cfg(feature="encoder_cinepak")]
 mod cinepakenc;
+#[cfg(feature="encoder_rawvideo")]
+mod rawvideoenc;
 #[cfg(feature="encoder_zmbv")]
 mod zmbvenc;
 
@@ -84,6 +86,8 @@ mod zmbvenc;
 const ENCODERS: &[EncoderInfo] = &[
 #[cfg(feature="encoder_cinepak")]
     EncoderInfo { name: "cinepak", get_encoder: cinepakenc::get_encoder },
+#[cfg(feature="encoder_rawvideo")]
+    EncoderInfo { name: "rawvideo", get_encoder: rawvideoenc::get_encoder },
 #[cfg(feature="encoder_zmbv")]
     EncoderInfo { name: "zmbv", get_encoder: zmbvenc::get_encoder },
 
diff --git a/nihav-commonfmt/src/codecs/rawvideoenc.rs b/nihav-commonfmt/src/codecs/rawvideoenc.rs
new file mode 100644
index 0000000..8847520
--- /dev/null
+++ b/nihav-commonfmt/src/codecs/rawvideoenc.rs
@@ -0,0 +1,100 @@
+use nihav_core::codecs::*;
+
+struct RawEncoder {
+    stream: Option<NAStreamRef>,
+    pkt:    Option<NAPacket>,
+}
+
+impl RawEncoder {
+    fn new() -> Self {
+        Self {
+            stream: None,
+            pkt:    None,
+        }
+    }
+}
+
+impl NAEncoder for RawEncoder {
+    fn negotiate_format(&self, encinfo: &EncodeParameters) -> EncoderResult<EncodeParameters> {
+        match encinfo.format {
+            NACodecTypeInfo::None => {
+                Ok(EncodeParameters {
+                    format: NACodecTypeInfo::Video(NAVideoInfo::new(0, 0, false, YUV420_FORMAT)),
+                    ..Default::default()
+                })
+            },
+            NACodecTypeInfo::Video(_) => Ok(*encinfo),
+            NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
+        }
+    }
+    fn get_capabilities(&self) -> u64 { ENC_CAPS_CBR }
+    fn init(&mut self, stream_id: u32, encinfo: EncodeParameters) -> EncoderResult<NAStreamRef> {
+        match encinfo.format {
+            NACodecTypeInfo::None => Err(EncoderError::FormatError),
+            NACodecTypeInfo::Audio(_) => Err(EncoderError::FormatError),
+            NACodecTypeInfo::Video(_) => {
+                let info = NACodecInfo::new("rawvideo", encinfo.format, None);
+                let mut stream = NAStream::new(StreamType::Video, stream_id, info, encinfo.tb_num, encinfo.tb_den, 0);
+                stream.set_num(stream_id as usize);
+                let stream = stream.into_ref();
+                self.stream = Some(stream.clone());
+                Ok(stream)
+            }
+        }
+    }
+    fn encode(&mut self, frm: &NAFrame) -> EncoderResult<()> {
+        let buf = frm.get_buffer();
+        let mut dbuf;
+        match buf {
+            NABufferType::Video(ref vbuf) => {
+                let vinfo = vbuf.get_info();
+                if !vinfo.format.model.is_yuv() || !vinfo.format.is_unpacked() {
+                    return Err(EncoderError::NotImplemented);
+                }
+
+                let src = vbuf.get_data();
+                dbuf = Vec::with_capacity(src.len());
+                for (comp, cinfo) in vinfo.format.comp_info.iter().enumerate() {
+                    if cinfo.is_none() {
+                        continue;
+                    }
+                    let (width, height) = vbuf.get_dimensions(comp);
+                    let off = vbuf.get_offset(comp);
+                    let stride = vbuf.get_stride(comp);
+
+                    for line in src[off..].chunks(stride).take(height) {
+                        dbuf.extend_from_slice(&line[..width]);
+                    }
+                }
+            },
+            NABufferType::VideoPacked(ref _vbuf) => return Err(EncoderError::NotImplemented),
+            NABufferType::Video16(ref _vbuf) => return Err(EncoderError::NotImplemented),
+            NABufferType::Video32(ref _vbuf) => return Err(EncoderError::NotImplemented),
+            NABufferType::None => {
+                self.pkt = None;
+                return Ok(());
+            },
+            _ => return Err(EncoderError::FormatError),
+        };
+        self.pkt = Some(NAPacket::new(self.stream.clone().unwrap(), frm.ts, true, dbuf));
+        Ok(())
+    }
+    fn get_packet(&mut self) -> EncoderResult<Option<NAPacket>> {
+        let mut npkt = None;
+        std::mem::swap(&mut self.pkt, &mut npkt);
+        Ok(npkt)
+    }
+    fn flush(&mut self) -> EncoderResult<()> {
+        Ok(())
+    }
+}
+
+impl NAOptionHandler for RawEncoder {
+    fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
+    fn set_options(&mut self, _options: &[NAOption]) { }
+    fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
+}
+
+pub fn get_encoder() -> Box<dyn NAEncoder + Send> {
+    Box::new(RawEncoder::new())
+}
-- 
2.30.2