mpegaudio: make packetiser checks less verbose

[nihav.git] / nihav-commonfmt / src / demuxers / y4m.rs

use nihav_core::demuxers::*;
use std::str::FromStr;

struct Y4MDemuxer<'a> {
    src:        &'a mut ByteReader<'a>,
    width:      usize,
    height:     usize,
    frame_size: usize,
    hdr_size:   u64,
    fps_num:    u32,
    fps_den:    u32,
    frameno:    u64,
}

impl<'a> DemuxCore<'a> for Y4MDemuxer<'a> {
    fn open(&mut self, strmgr: &mut StreamManager, seek_index: &mut SeekIndex) -> DemuxerResult<()> {
        let format = self.parse_header()?;
        seek_index.mode = SeekIndexMode::Automatic;

        let vhdr = NAVideoInfo::new(self.width, self.height, false, format);
        let vci = NACodecTypeInfo::Video(vhdr);
        let vinfo = NACodecInfo::new("rawvideo", vci, None);
        if strmgr.add_stream(NAStream::new(StreamType::Video, 0, vinfo, self.fps_num, self.fps_den, 0)).is_none() {
            return Err(DemuxerError::MemoryError);
        }

        Ok(())
    }

    fn get_frame(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<NAPacket> {
        let mut marker = [0u8; 6];
        let res                         = self.src.read_buf(&mut marker);
        match res {
            Err(ByteIOError::EOF) => return Err(DemuxerError::EOF),
            Err(err) => return Err(err.into()),
            _ => {},
        };
        validate!(&marker == b"FRAME\n");
        let stream = strmgr.get_stream(0).unwrap();
        let (tb_num, tb_den) = stream.get_timebase();
        let ts = NATimeInfo::new(Some(self.frameno), None, None, tb_num, tb_den);
        let pkt = self.src.read_packet(stream, ts, true, self.frame_size)?;
        self.frameno += 1;
        Ok(pkt)
    }

    fn seek(&mut self, time: NATimePoint, _seek_index: &SeekIndex) -> DemuxerResult<()> {
        let new_fno = match time {
                NATimePoint::PTS(pts) => {
                    pts
                },
                NATimePoint::Milliseconds(ms) => {
                    if (self.fps_num == 0) || (self.fps_den == 0) {
                        return Err(DemuxerError::SeekError);
                    }
                    NATimeInfo::time_to_ts(ms, 1000, self.fps_num, self.fps_den)
                },
                NATimePoint::None => return Err(DemuxerError::SeekError),
            };
        let pos = self.hdr_size + new_fno * ((self.frame_size + 6) as u64);
        self.src.seek(SeekFrom::Start(pos))?;
        self.frameno = new_fno;

        Ok(())
    }
    fn get_duration(&self) -> u64 { 0 }
}

impl<'a> NAOptionHandler for Y4MDemuxer<'a> {
    fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
    fn set_options(&mut self, _options: &[NAOption]) { }
    fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
}

impl<'a> Y4MDemuxer<'a> {
    fn new(io: &'a mut ByteReader<'a>) -> Self {
        Self {
            src:        io,
            width:      0,
            height:     0,
            frame_size: 0,
            fps_num:    0,
            fps_den:    0,
            hdr_size:   0,
            frameno:    0,
        }
    }
    fn parse_header(&mut self) -> DemuxerResult<NAPixelFormaton> {
        let mut format = RGB24_FORMAT;

        let mut magic = [0u8; 10];
                                    self.src.read_buf(&mut magic)?;
        validate!(&magic == b"YUV4MPEG2 ");
        while let Ok((last, tok)) = read_token(&mut self.src) {
            let (id, val) = tok.split_at(1);
            validate!(id.len() == 1);
            match id.bytes().next().unwrap() {
                b'W' => {
                    if let Ok(w) = val.parse::<usize>() {
                        self.width = w;
                    }
                },
                b'H' => {
                    if let Ok(h) = val.parse::<usize>() {
                        self.height = h;
                    }
                },
                b'F' => {
                    if let Ok(fden) = val.parse::<u32>() {
                        self.fps_num = 1;
                        self.fps_den = fden;
                    } else {
                        let vals: Vec<&str> = val.split(':').collect();
                        if vals.len() == 2 {
                            if let Ok(fnum) = vals[1].parse::<u32>() {
                                self.fps_num = fnum;
                            }
                            if let Ok(fden) = vals[0].parse::<u32>() {
                                self.fps_den = fden;
                            }
                        }
                    }
                },
                b'C' => {
                    let fmt_str = val.as_bytes();
                    validate!(fmt_str.len() >= 3);
                    let mut pix_name: [u8; 7] = *b"yuv000p";
                    validate!(fmt_str[0] == b'4');
                    pix_name[3..6].copy_from_slice(&fmt_str[..3]);

                    if let Ok(fmt_name) = std::str::from_utf8(&pix_name) {
                        if let Ok(val) = NAPixelFormaton::from_str(fmt_name) {
                            format = val;
                        }
                    }
                    if format.model.is_yuv() {
                        format.model = ColorModel::YUV(YUVSubmodel::YCbCr);
                        if fmt_str.len() > 3 {
                            let (_, tail) = fmt_str.split_at(3);
                            if tail == b"jpeg" {
                                format.model = ColorModel::YUV(YUVSubmodel::YUVJ);
                            }
                        }
                    }
                },
                _ => {},
            };

            if last {
                break;
            }
        }
        validate!(self.width > 0 && self.height > 0 && format.model.is_yuv());
        self.frame_size = 0;
        for chromaton in format.comp_info.iter() {
            if let Some(ref chr) = chromaton {
                self.frame_size += chr.get_data_size(self.width, self.height);
            }
        }
        validate!(self.frame_size > 0);

        Ok(format)
    }
}

fn read_token(src: &mut ByteReader) -> DemuxerResult<(bool, String)> {
    let mut string = String::new();
    let ws;
    loop {
        let b                       = src.read_byte()?;
        match b {
            b' ' | b'\n' => { ws = b; break; },
            0..=0x7F => string.push(b as char),
            _ => return Err(DemuxerError::InvalidData),
        }
    }

    Ok((ws == b'\n', string))
}

pub struct Y4MDemuxerCreator { }

impl DemuxerCreator for Y4MDemuxerCreator {
    fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a> {
        Box::new(Y4MDemuxer::new(br))
    }
    fn get_name(&self) -> &'static str { "yuv4mpeg" }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::fs::File;

    #[test]
    fn test_y4m_demux() {
        // sample: self-created with avconv
        let mut file = File::open("assets/Misc/test.y4m").unwrap();
        let mut fr = FileReader::new_read(&mut file);
        let mut br = ByteReader::new(&mut fr);
        let mut dmx = Y4MDemuxer::new(&mut br);
        let mut sm = StreamManager::new();
        let mut si = SeekIndex::new();
        dmx.open(&mut sm, &mut si).unwrap();

        loop {
            let pktres = dmx.get_frame(&mut sm);
            if let Err(e) = pktres {
                if e == DemuxerError::EOF { break; }
                panic!("error");
            }
            let pkt = pktres.unwrap();
            println!("Got {}", pkt);
        }
    }
}