[nihav.git] / nihav-realmedia / src / muxers / rmvb / videostream.rs

use std::collections::VecDeque;
use nihav_core::frame::*;
use nihav_core::muxers::*;
use super::RMStreamWriter;

static VIDEO_CODEC_REGISTRY: &[(&[u8;4], &str)] = &[
    (b"RV10", "realvideo1"),
    (b"RV20", "realvideo2"),
    (b"RVTR", "realvideo2"),
    (b"RV30", "realvideo3"),
    (b"RV40", "realvideo4"),
    (b"RV60", "realvideo6"),
    (b"CLV1", "clearvideo_rm"),
];

pub struct DummyStreamWriter {}
impl RMStreamWriter for DummyStreamWriter {
    fn write_header(&mut self, _bw: &mut ByteWriter, _stream: &NAStream) -> MuxerResult<()> {
        Ok(())
    }
    fn queue_packet(&mut self, _pkt: NAPacket, _ms: u32) -> bool {
        true
    }
    fn get_packet(&mut self) -> Option<(Vec<u8>, u32, bool)> {
        None
    }
    fn flush(&mut self) { }
    fn finish(&mut self, _bw: &mut ByteWriter) -> MuxerResult<()> {
        Ok(())
    }
    fn set_pkt_size(&mut self, _pkt_size: usize) {}
}

#[derive(Clone,Copy)]
enum VideoDataType {
    Frame,
    Slice{pkt_no: u8, npkt: u8, full_size: u32, offset: u32},
}

impl VideoDataType {
    fn is_frame(self) -> bool { matches!(self, VideoDataType::Frame) }
}

fn val_to_size(val: u32) -> usize { if val < (1 << 14) { 2 } else { 4 } }

#[derive(Clone)]
struct VideoData {
    vtype:      VideoDataType,
    pts:        u32,
    seq_no:     u8,
    data:       Vec<u8>,
}

impl VideoData {
    fn get_pkt_len(&self) -> usize {
        let plen = self.data.len();
        let ts_size = val_to_size(self.pts);
        match self.vtype {
            VideoDataType::Frame => plen + val_to_size(plen as u32) + ts_size + 2,
            VideoDataType::Slice{pkt_no: _, npkt: _, full_size, offset} => plen + val_to_size(full_size) + val_to_size(offset) + 3,
        }
    }
}

struct VideoStreamWriter {
    fcc:        [u8; 4],
    seq_no:     u8,
    time:       u32,
    mi_time:    u32,
    pkt_size:   usize,
    queue:      VecDeque<VideoData>,
    flush:      bool,
}

impl RMStreamWriter for VideoStreamWriter {
    fn write_header(&mut self, bw: &mut ByteWriter, vstream: &NAStream) -> MuxerResult<()> {
        let info = vstream.get_info().get_properties().get_video_info().unwrap();
        let start = bw.tell();

        bw.write_u32be(0)?; // header size
        bw.write_buf(b"VIDO")?;
        bw.write_buf(&self.fcc)?;
        bw.write_u16be(info.width as u16)?;
        bw.write_u16be(info.height as u16)?;
        bw.write_u16be(12)?; // bpp
        bw.write_u16be(0)?; // aligned width
        bw.write_u16be(0)?; // aligned height
        let (tb_num, tb_den) = vstream.get_timebase();
        if tb_num != 0 && tb_den != 0 {
            bw.write_u16be((tb_den / tb_num) as u16)?;
            let mut fps_frac = tb_den % tb_num;
            let mut div = tb_num;
            while div >= 0x10000 {
                fps_frac >>= 1;
                div      >>= 1;
            }
            fps_frac = (fps_frac << 16) / div;
            bw.write_u16le(fps_frac as u16)?;
        } else {
            bw.write_u16be(0)?;
            bw.write_u16be(0)?;
        }

        if let Some(edata) = vstream.get_info().get_extradata() {
            bw.write_buf(&edata)?;
        }
        let end = bw.tell();
        bw.seek(SeekFrom::Start(start))?;
        bw.write_u32be((end - start) as u32)?;
        bw.seek(SeekFrom::Start(end))?;
        Ok(())
    }
    fn queue_packet(&mut self, pkt: NAPacket, ms: u32) -> bool {
        let tot_size = self.queue.iter().fold(0usize, |acc, q| acc + q.get_pkt_len());
        if tot_size > self.pkt_size {
            return false;
        }

        self.time = ms;
        if ms > 0 {
            self.mi_time = ms.max(self.mi_time + 1);
        }

        let src = pkt.get_buffer();
        let nslices = usize::from(src[0]) + 1;
        let hdr_size = nslices * 8 + 1;

        if nslices == 1 {
            self.queue.push_back(VideoData {
                    vtype:      VideoDataType::Frame,
                    pts:        self.mi_time,
                    seq_no:     self.seq_no,
                    data:       src[9..].to_vec(),
                });
        } else if src.len() > hdr_size {
            let mut slice_sizes = [0; 256];
            let mut slice_offs  = [0; 256];

            for (el, src) in slice_offs.iter_mut().zip(src[1..].chunks_exact(8)) {
                *el = read_u32be(&src[4..]).unwrap() as usize;
            }
            for (dst, offs) in slice_sizes[..nslices - 1].iter_mut().zip(slice_offs.windows(2)) {
                *dst = offs[1] - offs[0];
            }
            slice_sizes[nslices - 1] = src.len() - hdr_size - slice_offs[nslices - 1];

            let src = &src[hdr_size..];
            let full_size = src.len() as u32;
            let npkt = nslices as u8;
            for (pkt_no, (&offset, &size)) in slice_offs.iter().zip(slice_sizes.iter()).take(nslices).enumerate() {
                let vtype = VideoDataType::Slice{pkt_no: (pkt_no + 1) as u8, npkt, full_size, offset: offset as u32};
                self.queue.push_back(VideoData {
                        vtype,
                        pts:        self.mi_time,
                        seq_no:     self.seq_no,
                        data:       src[offset..][..size].to_vec(),
                    });
            }
        }

        self.seq_no = self.seq_no.wrapping_add(1);

        true
    }
    fn get_packet(&mut self) -> Option<(Vec<u8>, u32, bool)> {
        if self.queue.is_empty() {
            return None;
        }
        let tot_size = self.queue.iter().fold(0usize, |acc, q| acc + q.get_pkt_len());
        if tot_size < self.pkt_size && !self.flush {
            return None;
        }
        let mut pkt_buf = Vec::new();

        let first = self.queue.pop_front().unwrap();
        let is_first = match first.vtype {
                VideoDataType::Frame => true,
                VideoDataType::Slice{pkt_no, npkt: _, full_size: _, offset: _} => pkt_no == 1,
            };
        if self.queue.is_empty() || (first.get_pkt_len() + self.queue[0].get_pkt_len() + 4 > self.pkt_size) {
            match first.vtype {
                VideoDataType::Frame => {
                    pkt_buf.push(0x40); // 0x1 = whole frame
                    pkt_buf.push(first.seq_no);
                    pkt_buf.extend_from_slice(&first.data);
                },
                VideoDataType::Slice{pkt_no, npkt, full_size: _, offset: _} => {
                    let id = if pkt_no == npkt { 2 } else { 0 };
                    write_slice(&mut pkt_buf, id, &first);
                },
            };
        } else {
            let second = &self.queue[0];
            match (first.vtype.is_frame(), second.vtype.is_frame()) {
                (true, true) => {
                    write_multiple_frame(&mut pkt_buf, &first);
                    while !self.queue.is_empty() && self.queue[0].vtype.is_frame() && (pkt_buf.len() + self.queue[0].get_pkt_len() < self.pkt_size) {
                        let frm = self.queue.pop_front().unwrap();
                        write_multiple_frame(&mut pkt_buf, &frm);
                    }
                },
                (true, false) => {
                    pkt_buf.push(0x40); // 0x1 = whole frame
                    pkt_buf.push(first.seq_no);
                    pkt_buf.extend_from_slice(&first.data);
                },
                (false, true) => {
                    write_slice(&mut pkt_buf, 2, &first);
                    while !self.queue.is_empty() && self.queue[0].vtype.is_frame() && (pkt_buf.len() + self.queue[0].get_pkt_len() < self.pkt_size) {
                        let frm = self.queue.pop_front().unwrap();
                        write_multiple_frame(&mut pkt_buf, &frm);
                    }
                },
                (false, false) => {
                    if let VideoDataType::Slice{pkt_no, npkt, full_size: _, offset: _} = first.vtype {
                       let id = if pkt_no == npkt { 2 } else { 0 };
                       write_slice(&mut pkt_buf, id, &first);
                    } else {
                        unreachable!()
                    }
                },
            };
        }
        Some((pkt_buf, first.pts, is_first))
    }
    fn flush(&mut self) {
        self.flush = true;
    }
    fn finish(&mut self, _bw: &mut ByteWriter) -> MuxerResult<()> {
        Ok(())
    }
    fn set_pkt_size(&mut self, pkt_size: usize) {
        self.pkt_size = pkt_size;
    }
}

fn write_16_or_32(dst: &mut Vec<u8>, val: u32) {
    if val < (1 << 14) {
        dst.push((1 << 6) | ((val >> 8) as u8));
        dst.push(val as u8);
    } else {
        dst.push((val >> 24) as u8);
        dst.push((val >> 16) as u8);
        dst.push((val >>  8) as u8);
        dst.push( val        as u8);
    }
}

fn write_multiple_frame(dst: &mut Vec<u8>, frm: &VideoData) {
    dst.push(0xC0); // 0x3 = multiple frame
    write_16_or_32(dst, frm.data.len() as u32);
    write_16_or_32(dst, frm.pts);
    dst.push(frm.seq_no);
    dst.extend_from_slice(&frm.data);
}

fn write_slice(dst: &mut Vec<u8>, id: u8, src: &VideoData) {
    if let VideoDataType::Slice{pkt_no, npkt, full_size, offset} = src.vtype {
        dst.push((id << 6) | (npkt >> 1));
        dst.push((npkt << 7) | pkt_no);
        write_16_or_32(dst, full_size);
        if id == 0 {
            write_16_or_32(dst, offset);
        } else {
            write_16_or_32(dst, src.data.len() as u32);
        }
        dst.push(src.seq_no);
        dst.extend_from_slice(&src.data);
    } else {
        unreachable!()
    }
}

pub fn create_video_stream(stream: &NAStream, pkt_size: usize) -> MuxerResult<Box<dyn RMStreamWriter>> {
    let info = stream.get_info();
    let cname = info.get_name();

    for &(fcc, name) in VIDEO_CODEC_REGISTRY.iter() {
        if name == cname {
            return Ok(Box::new(VideoStreamWriter {
                    fcc:        *fcc,
                    seq_no:     0,
                    time:       0,
                    mi_time:    0,
                    pkt_size,
                    queue:      VecDeque::new(),
                    flush:      false,
                }));
        }
    }
    Err(MuxerError::UnsupportedFormat)
}
Commit	Line	Data
353373a3	1	use std::collections::VecDeque;
9dc1fb4b KS	2	use nihav_core::frame::*;
	3	use nihav_core::muxers::*;
	4	use super::RMStreamWriter;
	5
	6	static VIDEO_CODEC_REGISTRY: &[(&[u8;4], &str)] = &[
	7	(b"RV10", "realvideo1"),
	8	(b"RV20", "realvideo2"),
	9	(b"RVTR", "realvideo2"),
	10	(b"RV30", "realvideo3"),
	11	(b"RV40", "realvideo4"),
	12	(b"RV60", "realvideo6"),
	13	(b"CLV1", "clearvideo_rm"),
	14	];
	15
	16	pub struct DummyStreamWriter {}
	17	impl RMStreamWriter for DummyStreamWriter {
	18	fn write_header(&mut self, _bw: &mut ByteWriter, _stream: &NAStream) -> MuxerResult<()> {
	19	Ok(())
	20	}
c5963b17	21	fn queue_packet(&mut self, _pkt: NAPacket, _ms: u32) -> bool {
9dc1fb4b KS	22	true
9dc1fb4b KS	23	}
c5963b17	24	fn get_packet(&mut self) -> Option<(Vec<u8>, u32, bool)> {
9dc1fb4b KS	25	None
	26	}
	27	fn flush(&mut self) { }
	28	fn finish(&mut self, _bw: &mut ByteWriter) -> MuxerResult<()> {
	29	Ok(())
	30	}
353373a3 KS	31	fn set_pkt_size(&mut self, _pkt_size: usize) {}
	32	}
	33
	34	#[derive(Clone,Copy)]
	35	enum VideoDataType {
	36	Frame,
	37	Slice{pkt_no: u8, npkt: u8, full_size: u32, offset: u32},
	38	}
	39
	40	impl VideoDataType {
	41	fn is_frame(self) -> bool { matches!(self, VideoDataType::Frame) }
	42	}
	43
	44	fn val_to_size(val: u32) -> usize { if val < (1 << 14) { 2 } else { 4 } }
	45
	46	#[derive(Clone)]
	47	struct VideoData {
	48	vtype: VideoDataType,
	49	pts: u32,
	50	seq_no: u8,
	51	data: Vec<u8>,
	52	}
	53
	54	impl VideoData {
	55	fn get_pkt_len(&self) -> usize {
	56	let plen = self.data.len();
	57	let ts_size = val_to_size(self.pts);
	58	match self.vtype {
	59	VideoDataType::Frame => plen + val_to_size(plen as u32) + ts_size + 2,
	60	VideoDataType::Slice{pkt_no: _, npkt: _, full_size, offset} => plen + val_to_size(full_size) + val_to_size(offset) + 3,
	61	}
	62	}
9dc1fb4b KS	63	}
	64
	65	struct VideoStreamWriter {
	66	fcc: [u8; 4],
9dc1fb4b	67	seq_no: u8,
c5963b17	68	time: u32,
e0d229ca	69	mi_time: u32,
353373a3 KS	70	pkt_size: usize,
	71	queue: VecDeque<VideoData>,
	72	flush: bool,
9dc1fb4b KS	73	}
	74
	75	impl RMStreamWriter for VideoStreamWriter {
	76	fn write_header(&mut self, bw: &mut ByteWriter, vstream: &NAStream) -> MuxerResult<()> {
	77	let info = vstream.get_info().get_properties().get_video_info().unwrap();
	78	let start = bw.tell();
	79
	80	bw.write_u32be(0)?; // header size
	81	bw.write_buf(b"VIDO")?;
	82	bw.write_buf(&self.fcc)?;
	83	bw.write_u16be(info.width as u16)?;
	84	bw.write_u16be(info.height as u16)?;
	85	bw.write_u16be(12)?; // bpp
	86	bw.write_u16be(0)?; // aligned width
	87	bw.write_u16be(0)?; // aligned height
	88	let (tb_num, tb_den) = vstream.get_timebase();
	89	if tb_num != 0 && tb_den != 0 {
	90	bw.write_u16be((tb_den / tb_num) as u16)?;
	91	let mut fps_frac = tb_den % tb_num;
	92	let mut div = tb_num;
	93	while div >= 0x10000 {
	94	fps_frac >>= 1;
	95	div >>= 1;
	96	}
	97	fps_frac = (fps_frac << 16) / div;
	98	bw.write_u16le(fps_frac as u16)?;
	99	} else {
	100	bw.write_u16be(0)?;
	101	bw.write_u16be(0)?;
	102	}
	103
	104	if let Some(edata) = vstream.get_info().get_extradata() {
	105	bw.write_buf(&edata)?;
	106	}
	107	let end = bw.tell();
	108	bw.seek(SeekFrom::Start(start))?;
	109	bw.write_u32be((end - start) as u32)?;
	110	bw.seek(SeekFrom::Start(end))?;
	111	Ok(())
	112	}
c5963b17	113	fn queue_packet(&mut self, pkt: NAPacket, ms: u32) -> bool {
353373a3 KS	114	let tot_size = self.queue.iter().fold(0usize, \|acc, q\| acc + q.get_pkt_len());
	115	if tot_size > self.pkt_size {
	116	return false;
	117	}
	118
	119	self.time = ms;
	120	if ms > 0 {
	121	self.mi_time = ms.max(self.mi_time + 1);
	122	}
	123
	124	let src = pkt.get_buffer();
	125	let nslices = usize::from(src[0]) + 1;
	126	let hdr_size = nslices * 8 + 1;
	127
	128	if nslices == 1 {
	129	self.queue.push_back(VideoData {
	130	vtype: VideoDataType::Frame,
	131	pts: self.mi_time,
	132	seq_no: self.seq_no,
	133	data: src[9..].to_vec(),
	134	});
	135	} else if src.len() > hdr_size {
	136	let mut slice_sizes = [0; 256];
	137	let mut slice_offs = [0; 256];
	138
	139	for (el, src) in slice_offs.iter_mut().zip(src[1..].chunks_exact(8)) {
	140	*el = read_u32be(&src[4..]).unwrap() as usize;
	141	}
	142	for (dst, offs) in slice_sizes[..nslices - 1].iter_mut().zip(slice_offs.windows(2)) {
	143	*dst = offs[1] - offs[0];
	144	}
	145	slice_sizes[nslices - 1] = src.len() - hdr_size - slice_offs[nslices - 1];
	146
	147	let src = &src[hdr_size..];
	148	let full_size = src.len() as u32;
	149	let npkt = nslices as u8;
	150	for (pkt_no, (&offset, &size)) in slice_offs.iter().zip(slice_sizes.iter()).take(nslices).enumerate() {
	151	let vtype = VideoDataType::Slice{pkt_no: (pkt_no + 1) as u8, npkt, full_size, offset: offset as u32};
	152	self.queue.push_back(VideoData {
	153	vtype,
	154	pts: self.mi_time,
	155	seq_no: self.seq_no,
	156	data: src[offset..][..size].to_vec(),
	157	});
9dc1fb4b	158	}
9dc1fb4b	159	}
353373a3 KS	160
	161	self.seq_no = self.seq_no.wrapping_add(1);
	162
	163	true
9dc1fb4b	164	}
c5963b17	165	fn get_packet(&mut self) -> Option<(Vec<u8>, u32, bool)> {
353373a3 KS	166	if self.queue.is_empty() {
	167	return None;
	168	}
	169	let tot_size = self.queue.iter().fold(0usize, \|acc, q\| acc + q.get_pkt_len());
	170	if tot_size < self.pkt_size && !self.flush {
	171	return None;
	172	}
	173	let mut pkt_buf = Vec::new();
	174
	175	let first = self.queue.pop_front().unwrap();
	176	let is_first = match first.vtype {
	177	VideoDataType::Frame => true,
	178	VideoDataType::Slice{pkt_no, npkt: _, full_size: _, offset: _} => pkt_no == 1,
	179	};
	180	if self.queue.is_empty() \|\| (first.get_pkt_len() + self.queue[0].get_pkt_len() + 4 > self.pkt_size) {
	181	match first.vtype {
	182	VideoDataType::Frame => {
	183	pkt_buf.push(0x40); // 0x1 = whole frame
	184	pkt_buf.push(first.seq_no);
	185	pkt_buf.extend_from_slice(&first.data);
	186	},
	187	VideoDataType::Slice{pkt_no, npkt, full_size: _, offset: _} => {
	188	let id = if pkt_no == npkt { 2 } else { 0 };
	189	write_slice(&mut pkt_buf, id, &first);
	190	},
	191	};
	192	} else {
	193	let second = &self.queue[0];
	194	match (first.vtype.is_frame(), second.vtype.is_frame()) {
	195	(true, true) => {
	196	write_multiple_frame(&mut pkt_buf, &first);
	197	while !self.queue.is_empty() && self.queue[0].vtype.is_frame() && (pkt_buf.len() + self.queue[0].get_pkt_len() < self.pkt_size) {
	198	let frm = self.queue.pop_front().unwrap();
	199	write_multiple_frame(&mut pkt_buf, &frm);
	200	}
	201	},
	202	(true, false) => {
	203	pkt_buf.push(0x40); // 0x1 = whole frame
	204	pkt_buf.push(first.seq_no);
	205	pkt_buf.extend_from_slice(&first.data);
	206	},
	207	(false, true) => {
	208	write_slice(&mut pkt_buf, 2, &first);
	209	while !self.queue.is_empty() && self.queue[0].vtype.is_frame() && (pkt_buf.len() + self.queue[0].get_pkt_len() < self.pkt_size) {
	210	let frm = self.queue.pop_front().unwrap();
	211	write_multiple_frame(&mut pkt_buf, &frm);
9dc1fb4b	212	}
353373a3 KS	213	},
	214	(false, false) => {
	215	if let VideoDataType::Slice{pkt_no, npkt, full_size: _, offset: _} = first.vtype {
	216	let id = if pkt_no == npkt { 2 } else { 0 };
	217	write_slice(&mut pkt_buf, id, &first);
9dc1fb4b	218	} else {
353373a3	219	unreachable!()
9dc1fb4b	220	}
353373a3 KS	221	},
353373a3 KS	222	};
9dc1fb4b	223	}
353373a3 KS	224	Some((pkt_buf, first.pts, is_first))
	225	}
	226	fn flush(&mut self) {
	227	self.flush = true;
9dc1fb4b	228	}
9dc1fb4b KS	229	fn finish(&mut self, _bw: &mut ByteWriter) -> MuxerResult<()> {
	230	Ok(())
	231	}
353373a3 KS	232	fn set_pkt_size(&mut self, pkt_size: usize) {
	233	self.pkt_size = pkt_size;
	234	}
	235	}
	236
	237	fn write_16_or_32(dst: &mut Vec<u8>, val: u32) {
	238	if val < (1 << 14) {
	239	dst.push((1 << 6) \| ((val >> 8) as u8));
	240	dst.push(val as u8);
	241	} else {
	242	dst.push((val >> 24) as u8);
	243	dst.push((val >> 16) as u8);
	244	dst.push((val >> 8) as u8);
	245	dst.push( val as u8);
	246	}
	247	}
	248
	249	fn write_multiple_frame(dst: &mut Vec<u8>, frm: &VideoData) {
	250	dst.push(0xC0); // 0x3 = multiple frame
	251	write_16_or_32(dst, frm.data.len() as u32);
e6aaad5c	252	write_16_or_32(dst, frm.pts);
353373a3 KS	253	dst.push(frm.seq_no);
	254	dst.extend_from_slice(&frm.data);
	255	}
	256
	257	fn write_slice(dst: &mut Vec<u8>, id: u8, src: &VideoData) {
	258	if let VideoDataType::Slice{pkt_no, npkt, full_size, offset} = src.vtype {
	259	dst.push((id << 6) \| (npkt >> 1));
	260	dst.push((npkt << 7) \| pkt_no);
	261	write_16_or_32(dst, full_size);
	262	if id == 0 {
	263	write_16_or_32(dst, offset);
	264	} else {
	265	write_16_or_32(dst, src.data.len() as u32);
	266	}
	267	dst.push(src.seq_no);
	268	dst.extend_from_slice(&src.data);
	269	} else {
	270	unreachable!()
	271	}
9dc1fb4b KS	272	}
9dc1fb4b KS	273
353373a3	274	pub fn create_video_stream(stream: &NAStream, pkt_size: usize) -> MuxerResult<Box<dyn RMStreamWriter>> {
9dc1fb4b KS	275	let info = stream.get_info();
	276	let cname = info.get_name();
	277
	278	for &(fcc, name) in VIDEO_CODEC_REGISTRY.iter() {
	279	if name == cname {
	280	return Ok(Box::new(VideoStreamWriter {
	281	fcc: *fcc,
9dc1fb4b	282	seq_no: 0,
c5963b17	283	time: 0,
e0d229ca	284	mi_time: 0,
353373a3 KS	285	pkt_size,
	286	queue: VecDeque::new(),
	287	flush: false,
9dc1fb4b KS	288	}));
	289	}
	290	}
	291	Err(MuxerError::UnsupportedFormat)
	292	}