[nihav.git] / nihav-commonfmt / src / muxers / avi.rs

use nihav_core::muxers::*;
use nihav_registry::register::*;

#[derive(Clone,Copy)]
struct IdxEntry {
    stream:     u32,
    stype:      StreamType,
    key:        bool,
    pos:        u32,
    len:        u32,
}

#[derive(Clone,Copy)]
struct AVIStream {
    strh_pos:   u64,
    nframes:    u32,
    is_video:   bool,
    max_size:   u32,
    pal_change: bool,
}

struct AVIMuxer<'a> {
    bw:             &'a mut ByteWriter<'a>,
    index:          Vec<IdxEntry>,
    video_str:      Option<usize>,
    video_id:       u32,
    data_pos:       u64,
    stream_info:    Vec<AVIStream>,
    pal_pos:        Vec<u32>,
}

impl<'a> AVIMuxer<'a> {
    fn new(bw: &'a mut ByteWriter<'a>) -> Self {
        Self {
            bw,
            index:          Vec::new(),
            video_str:      None,
            video_id:       0,
            data_pos:       0,
            stream_info:    Vec::with_capacity(2),
            pal_pos:        Vec::with_capacity(2),
        }
    }
}

fn patch_size(bw: &mut ByteWriter, pos: u64) -> MuxerResult<()> {
    let size = bw.tell() - pos;
    bw.seek(SeekFrom::Current(-((size + 4) as i64)))?;
    bw.write_u32le(size as u32)?;
    bw.seek(SeekFrom::End(0))?;
    Ok(())
}

fn write_chunk_hdr(bw: &mut ByteWriter, stype: StreamType, str_no: u32) -> MuxerResult<()> {
    bw.write_byte(b'0' + ((str_no / 10) as u8))?;
    bw.write_byte(b'0' + ((str_no % 10) as u8))?;
    match stype {
        StreamType::Video => { bw.write_buf(b"dc")?; },
        StreamType::Audio => { bw.write_buf(b"wb")?; },
        StreamType::Subtitles => { bw.write_buf(b"tx")?; },
        _ => return Err(MuxerError::UnsupportedFormat),
    };
    Ok(())
}

impl<'a> MuxCore<'a> for AVIMuxer<'a> {
    #[allow(clippy::unreadable_literal)]
    #[allow(clippy::cast_lossless)]
    fn create(&mut self, strmgr: &StreamManager) -> MuxerResult<()> {
        if strmgr.get_num_streams() == 0 {
            return Err(MuxerError::InvalidArgument);
        }
        if strmgr.get_num_streams() > 99 {
            return Err(MuxerError::UnsupportedFormat);
        }
        for (str_no, strm) in strmgr.iter().enumerate() {
            if strm.get_media_type() == StreamType::Video {
                self.video_str = Some(str_no);
                self.video_id  = strm.id;
                break;
            }
        }
        let (vinfo, tb_num, tb_den) = if let Some(str_id) = self.video_str {
                let vstr = strmgr.get_stream(str_id).unwrap();
                (vstr.get_info(), vstr.tb_num, vstr.tb_den)
            } else {
                (NACodecInfo::new_dummy(), 0, 1)
            };
        let hdrl_pos = self.bw.tell() + 20;
        self.bw.write_buf(b"RIFF\0\0\0\0AVI LIST\0\0\0\0hdrlavih")?;
        self.bw.write_u32le(56)?; // avih size
        let ms_per_frame = NATimeInfo::ts_to_time(1, 1000000, tb_num, tb_den);
        self.bw.write_u32le(ms_per_frame as u32)?;
        self.bw.write_u32le(0)?; // max transfer rate
        self.bw.write_u32le(0)?; // padding granularity
        self.bw.write_u32le(0)?; // flags
        self.bw.write_u32le(0)?; // total frames
        self.bw.write_u32le(0)?; // initial frames
        self.bw.write_u32le(strmgr.get_num_streams() as u32)?;
        self.bw.write_u32le(0)?; // suggested buffer size
        if let NACodecTypeInfo::Video(ref vinfo) = vinfo.get_properties() {
            self.bw.write_u32le(vinfo.width as u32)?;
            self.bw.write_u32le(vinfo.height as u32)?;
        } else {
            self.bw.write_u32le(0)?;
            self.bw.write_u32le(0)?;
        }
        self.bw.write_u32le(0)?; // reserved
        self.bw.write_u32le(0)?; // reserved
        self.bw.write_u32le(0)?; // reserved
        self.bw.write_u32le(0)?; // reserved

        self.pal_pos.clear();
        self.pal_pos.resize(strmgr.get_num_streams(), 0);
        for (strno, strm) in strmgr.iter().enumerate() {
            let strl_pos = self.bw.tell() + 8;
            self.bw.write_buf(b"LIST\0\0\0\0strlstrh")?;
            self.bw.write_u32le(56)?; // strh size

            match strm.get_media_type() {
                StreamType::Video => {
                    self.bw.write_buf(b"vids")?;
                    let fcc = find_avi_fourcc(strm.get_info().get_name());
                    if fcc.is_none() {
                        return Err(MuxerError::UnsupportedFormat);
                    }
                    self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
                    let vinfo = strm.get_info().get_properties().get_video_info().unwrap();
                    if vinfo.width >= (1 << 16) || vinfo.height >= (1 << 16) {
                        return Err(MuxerError::UnsupportedFormat);
                    }
                },
                StreamType::Audio => {
                    self.bw.write_buf(b"auds")?;
                    self.bw.write_u32le(0)?;
                },
                StreamType::Subtitles => {
                    self.bw.write_buf(b"txts")?;
                    self.bw.write_u32le(0)?;
                },
                _ => return Err(MuxerError::UnsupportedFormat),
            };
            self.stream_info.push(AVIStream {
                    strh_pos:   self.bw.tell(),
                    is_video:   strm.get_media_type() == StreamType::Video,
                    nframes:    0,
                    max_size:   0,
                    pal_change: false,
                });

            self.bw.write_u32le(0)?; // flags
            self.bw.write_u16le(0)?; // priority
            self.bw.write_u16le(0)?; // language
            self.bw.write_u32le(0)?; // initial frames
            self.bw.write_u32le(strm.tb_num)?;
            self.bw.write_u32le(strm.tb_den)?;
            self.bw.write_u32le(0)?; // start
            self.bw.write_u32le(0)?; // length
            self.bw.write_u32le(0)?; // suggested buffer size
            self.bw.write_u32le(0)?; // quality
            self.bw.write_u32le(0)?; // sample_size
            self.bw.write_u16le(0)?; // x
            self.bw.write_u16le(0)?; // y
            self.bw.write_u16le(0)?; // w
            self.bw.write_u16le(0)?; // h

            self.bw.write_buf(b"strf")?;
            self.bw.write_u32le(0)?;
            let strf_pos = self.bw.tell();
            match strm.get_media_type() {
                StreamType::Video => {
                    let vinfo = strm.get_info().get_properties().get_video_info().unwrap();
                    let hdr_pos = self.bw.tell();
                    self.bw.write_u32le(0)?;
                    self.bw.write_u32le(vinfo.width as u32)?;
                    self.bw.write_u32le(vinfo.height as u32)?;
                    if !vinfo.format.palette {
                        self.bw.write_u16le(vinfo.format.components as u16)?;
                        self.bw.write_u16le(vinfo.format.get_total_depth() as u16)?;
                    } else {
                        self.bw.write_u16le(1)?;
                        self.bw.write_u16le(8)?;
                    }
                    let fcc = find_avi_fourcc(strm.get_info().get_name());
                    if fcc.is_none() {
                        return Err(MuxerError::UnsupportedFormat);
                    }
                    self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
                    self.bw.write_u32le(0)?; // image size
                    self.bw.write_u32le(0)?; // x dpi
                    self.bw.write_u32le(0)?; // y dpi
                    if vinfo.format.palette {
                        self.bw.write_u32le(256)?; // total colors
                        self.bw.write_u32le(0)?; // important colors
                        self.pal_pos[strno] = self.bw.tell() as u32;
                        for _ in 0..256 {
                            self.bw.write_u32le(0)?;
                        }
                    } else {
                        self.bw.write_u32le(0)?; // total colors
                        self.bw.write_u32le(0)?; // important colors
                    }
                    if let Some(ref edata) = strm.get_info().get_extradata() {
                        self.bw.write_buf(edata.as_slice())?;
                    }
                    let bisize = self.bw.tell() - hdr_pos;
                    self.bw.seek(SeekFrom::Current(-(bisize as i64)))?;
                    self.bw.write_u32le(bisize as u32)?;
                    self.bw.seek(SeekFrom::End(0))?;
                },
                StreamType::Audio => {
                    let ainfo = strm.get_info().get_properties().get_audio_info().unwrap();
                    let twocc = find_wav_twocc(strm.get_info().get_name());
                    if twocc.is_none() {
                        return Err(MuxerError::UnsupportedFormat);
                    }
                    self.bw.write_u16le(twocc.unwrap_or(0))?;
                    self.bw.write_u16le(ainfo.channels as u16)?;
                    self.bw.write_u32le(ainfo.sample_rate)?;
                    self.bw.write_u32le(0)?; // avg bytes per second
                    self.bw.write_u16le(ainfo.block_len as u16)?;
                    self.bw.write_u16le(ainfo.format.bits as u16)?;
                    if let Some(ref edata) = strm.get_info().get_extradata() {
                        self.bw.write_buf(edata.as_slice())?;
                    }
                },
                StreamType::Subtitles => {
                    if let Some(ref edata) = strm.get_info().get_extradata() {
                        self.bw.write_buf(edata.as_slice())?;
                    }
                },
                _ => unreachable!(),
            };
            patch_size(self.bw, strf_pos)?;
            patch_size(self.bw, strl_pos)?;
        }
        patch_size(self.bw, hdrl_pos)?;

        self.data_pos = self.bw.tell() + 8;
        self.bw.write_buf(b"LIST\0\0\0\0movi")?;

        Ok(())
    }
    fn mux_frame(&mut self, _strmgr: &StreamManager, pkt: NAPacket) -> MuxerResult<()> {
        if self.data_pos == 0 {
            return Err(MuxerError::NotCreated);
        }
        let stream = pkt.get_stream();
        let str_num = stream.get_num();
        if str_num > 99 || str_num >= self.stream_info.len() {
            return Err(MuxerError::UnsupportedFormat);
        }

        let chunk_len = pkt.get_buffer().len() as u32;

        if self.pal_pos[str_num] != 0 {
            for sdata in pkt.side_data.iter() {
                if let NASideData::Palette(_, ref pal) = sdata {
                    let cur_pos = self.bw.tell();
                    self.bw.seek(SeekFrom::Start(u64::from(self.pal_pos[str_num])))?;
                    for quad in pal.chunks(4) {
                        self.bw.write_byte(quad[2])?;
                        self.bw.write_byte(quad[1])?;
                        self.bw.write_byte(quad[0])?;
                        self.bw.write_byte(0)?;
                    }
                    self.bw.seek(SeekFrom::Start(cur_pos))?;
                    self.pal_pos[str_num] = 0;
                    break;
                }
            }
        } else {
            for sdata in pkt.side_data.iter() {
                if let NASideData::Palette(true, ref pal) = sdata {
                    //todo search for changed region
                    let start_clr = 0usize;
                    let end_clr = 256usize;
                    if start_clr < end_clr {
                        let chunk_len = ((end_clr - start_clr) as u32) * 4 + 4;
                        self.bw.write_byte(b'0' + ((str_num / 10) as u8))?;
                        self.bw.write_byte(b'0' + ((str_num % 10) as u8))?;
                        self.bw.write_buf(b"pc")?;
                        self.bw.write_u32le(chunk_len)?;
                        self.bw.write_byte(start_clr as u8)?;
                        self.bw.write_byte((end_clr - start_clr) as u8)?;
                        self.bw.write_u16le(0)?; //flags
                        self.bw.write_buf(&pal[start_clr * 4..end_clr * 4])?;
                        self.stream_info[str_num].pal_change = true;
                    }
                }
            }
        }

        self.stream_info[str_num].nframes += 1;
        self.stream_info[str_num].max_size = self.stream_info[str_num].max_size.max(chunk_len);
        self.index.push(IdxEntry {
                stream: str_num as u32,
                stype:  stream.get_media_type(),
                key:    pkt.keyframe,
                pos:    self.bw.tell() as u32,
                len:    chunk_len });
        write_chunk_hdr(self.bw, stream.get_media_type(), str_num as u32)?;
        self.bw.write_u32le(chunk_len)?;
        self.bw.write_buf(pkt.get_buffer().as_slice())?;
        if (self.bw.tell() & 1) != 0 {
            self.bw.write_byte(0)?;
        }
        Ok(())
    }
    fn flush(&mut self) -> MuxerResult<()> {
        Ok(())
    }
    fn end(&mut self) -> MuxerResult<()> {
        patch_size(self.bw, self.data_pos)?;
        if !self.index.is_empty() {
            self.bw.write_buf(b"idx1")?;
            self.bw.write_u32le((self.index.len() * 16) as u32)?;
            for item in self.index.iter() {
                write_chunk_hdr(self.bw, item.stype, item.stream)?;
                self.bw.write_u32le(if item.key { 0x10 } else { 0 })?;
                self.bw.write_u32le(item.pos)?;
                self.bw.write_u32le(item.len)?;
            }
        }
        patch_size(self.bw, 8)?;
        let mut max_frames = 0;
        let mut max_size = 0;
        for stri in self.stream_info.iter() {
            max_frames = max_frames.max(stri.nframes);
            max_size = max_size.max(stri.max_size);
            if stri.pal_change {
                self.bw.seek(SeekFrom::Start(stri.strh_pos))?;
                self.bw.write_u32le(0x00010000)?;
            }
            self.bw.seek(SeekFrom::Start(stri.strh_pos + 0x18))?;
            self.bw.write_u32le(if stri.is_video { stri.nframes } else { 0 })?;
            self.bw.write_u32le(stri.max_size)?;
        }
        self.bw.seek(SeekFrom::Start(0x30))?;
        self.bw.write_u32le(max_frames)?;
        self.bw.seek(SeekFrom::Current(8))?;
        self.bw.write_u32le(max_size)?;
        Ok(())
    }
}

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

pub struct AVIMuxerCreator {}

impl MuxerCreator for AVIMuxerCreator {
    fn new_muxer<'a>(&self, bw: &'a mut ByteWriter<'a>) -> Box<dyn MuxCore<'a> + 'a> {
        Box::new(AVIMuxer::new(bw))
    }
    fn get_name(&self) -> &'static str { "avi" }
    fn get_capabilities(&self) -> MuxerCapabilities { MuxerCapabilities::Universal }
}

#[cfg(test)]
mod test {
    use nihav_core::codecs::*;
    use nihav_core::demuxers::*;
    use nihav_core::muxers::*;
    use nihav_codec_support::test::enc_video::*;
    use crate::*;

    #[test]
    fn test_avi_muxer() {
        let mut dmx_reg = RegisteredDemuxers::new();
        generic_register_all_demuxers(&mut dmx_reg);
        //test sample: https://samples.mplayerhq.hu/V-codecs/RT21/320x240/laser05.avi
        let dec_config = DecoderTestParams {
                demuxer:        "avi",
                in_name:        "assets/Indeo/laser05.avi",
                limit:          None,
                stream_type:    StreamType::None,
                dmx_reg, dec_reg: RegisteredDecoders::new(),
            };
        let mut mux_reg = RegisteredMuxers::new();
        generic_register_all_muxers(&mut mux_reg);
        /*let enc_config = EncoderTestParams {
                muxer:      "avi",
                enc_name:   "",
                out_name:   "muxed.avi",
                mux_reg, enc_reg: RegisteredEncoders::new(),
            };
        test_remuxing(&dec_config, &enc_config);*/
        test_remuxing_md5(&dec_config, "avi", &mux_reg,
                          [0xa0fb0e47, 0x412e24dd, 0x6b89711c, 0x276fb799]);
    }
}
Commit	Line	Data
5cd2f537 KS	1	use nihav_core::muxers::*;
	2	use nihav_registry::register::*;
	3
	4	#[derive(Clone,Copy)]
	5	struct IdxEntry {
	6	stream: u32,
	7	stype: StreamType,
	8	key: bool,
	9	pos: u32,
	10	len: u32,
	11	}
	12
	13	#[derive(Clone,Copy)]
	14	struct AVIStream {
	15	strh_pos: u64,
	16	nframes: u32,
	17	is_video: bool,
c8db9313	18	max_size: u32,
04fafc56	19	pal_change: bool,
5cd2f537 KS	20	}
	21
	22	struct AVIMuxer<'a> {
	23	bw: &'a mut ByteWriter<'a>,
	24	index: Vec<IdxEntry>,
	25	video_str: Option<usize>,
	26	video_id: u32,
	27	data_pos: u64,
	28	stream_info: Vec<AVIStream>,
a1613eee	29	pal_pos: Vec<u32>,
5cd2f537 KS	30	}
	31
	32	impl<'a> AVIMuxer<'a> {
	33	fn new(bw: &'a mut ByteWriter<'a>) -> Self {
	34	Self {
	35	bw,
	36	index: Vec::new(),
	37	video_str: None,
	38	video_id: 0,
	39	data_pos: 0,
	40	stream_info: Vec::with_capacity(2),
a1613eee	41	pal_pos: Vec::with_capacity(2),
5cd2f537 KS	42	}
	43	}
	44	}
	45
	46	fn patch_size(bw: &mut ByteWriter, pos: u64) -> MuxerResult<()> {
	47	let size = bw.tell() - pos;
	48	bw.seek(SeekFrom::Current(-((size + 4) as i64)))?;
	49	bw.write_u32le(size as u32)?;
	50	bw.seek(SeekFrom::End(0))?;
	51	Ok(())
	52	}
	53
	54	fn write_chunk_hdr(bw: &mut ByteWriter, stype: StreamType, str_no: u32) -> MuxerResult<()> {
	55	bw.write_byte(b'0' + ((str_no / 10) as u8))?;
	56	bw.write_byte(b'0' + ((str_no % 10) as u8))?;
	57	match stype {
	58	StreamType::Video => { bw.write_buf(b"dc")?; },
	59	StreamType::Audio => { bw.write_buf(b"wb")?; },
	60	StreamType::Subtitles => { bw.write_buf(b"tx")?; },
	61	_ => return Err(MuxerError::UnsupportedFormat),
	62	};
	63	Ok(())
	64	}
	65
	66	impl<'a> MuxCore<'a> for AVIMuxer<'a> {
61cab15b KS	67	#[allow(clippy::unreadable_literal)]
61cab15b KS	68	#[allow(clippy::cast_lossless)]
5cd2f537 KS	69	fn create(&mut self, strmgr: &StreamManager) -> MuxerResult<()> {
	70	if strmgr.get_num_streams() == 0 {
	71	return Err(MuxerError::InvalidArgument);
	72	}
	73	if strmgr.get_num_streams() > 99 {
	74	return Err(MuxerError::UnsupportedFormat);
	75	}
405cec9e KS	76	for (str_no, strm) in strmgr.iter().enumerate() {
405cec9e KS	77	if strm.get_media_type() == StreamType::Video {
5cd2f537	78	self.video_str = Some(str_no);
405cec9e	79	self.video_id = strm.id;
5cd2f537 KS	80	break;
	81	}
	82	}
	83	let (vinfo, tb_num, tb_den) = if let Some(str_id) = self.video_str {
	84	let vstr = strmgr.get_stream(str_id).unwrap();
	85	(vstr.get_info(), vstr.tb_num, vstr.tb_den)
	86	} else {
	87	(NACodecInfo::new_dummy(), 0, 1)
	88	};
	89	let hdrl_pos = self.bw.tell() + 20;
	90	self.bw.write_buf(b"RIFF\0\0\0\0AVI LIST\0\0\0\0hdrlavih")?;
	91	self.bw.write_u32le(56)?; // avih size
	92	let ms_per_frame = NATimeInfo::ts_to_time(1, 1000000, tb_num, tb_den);
	93	self.bw.write_u32le(ms_per_frame as u32)?;
	94	self.bw.write_u32le(0)?; // max transfer rate
	95	self.bw.write_u32le(0)?; // padding granularity
	96	self.bw.write_u32le(0)?; // flags
	97	self.bw.write_u32le(0)?; // total frames
	98	self.bw.write_u32le(0)?; // initial frames
	99	self.bw.write_u32le(strmgr.get_num_streams() as u32)?;
	100	self.bw.write_u32le(0)?; // suggested buffer size
	101	if let NACodecTypeInfo::Video(ref vinfo) = vinfo.get_properties() {
	102	self.bw.write_u32le(vinfo.width as u32)?;
	103	self.bw.write_u32le(vinfo.height as u32)?;
	104	} else {
	105	self.bw.write_u32le(0)?;
	106	self.bw.write_u32le(0)?;
	107	}
	108	self.bw.write_u32le(0)?; // reserved
	109	self.bw.write_u32le(0)?; // reserved
	110	self.bw.write_u32le(0)?; // reserved
	111	self.bw.write_u32le(0)?; // reserved
	112
37952415	113	self.pal_pos.clear();
a1613eee	114	self.pal_pos.resize(strmgr.get_num_streams(), 0);
405cec9e	115	for (strno, strm) in strmgr.iter().enumerate() {
5cd2f537 KS	116	let strl_pos = self.bw.tell() + 8;
	117	self.bw.write_buf(b"LIST\0\0\0\0strlstrh")?;
	118	self.bw.write_u32le(56)?; // strh size
	119
405cec9e	120	match strm.get_media_type() {
5cd2f537 KS	121	StreamType::Video => {
5cd2f537 KS	122	self.bw.write_buf(b"vids")?;
405cec9e	123	let fcc = find_avi_fourcc(strm.get_info().get_name());
5cd2f537 KS	124	if fcc.is_none() {
	125	return Err(MuxerError::UnsupportedFormat);
	126	}
	127	self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
405cec9e	128	let vinfo = strm.get_info().get_properties().get_video_info().unwrap();
5cd2f537 KS	129	if vinfo.width >= (1 << 16) \|\| vinfo.height >= (1 << 16) {
	130	return Err(MuxerError::UnsupportedFormat);
	131	}
	132	},
	133	StreamType::Audio => {
	134	self.bw.write_buf(b"auds")?;
	135	self.bw.write_u32le(0)?;
	136	},
	137	StreamType::Subtitles => {
	138	self.bw.write_buf(b"txts")?;
	139	self.bw.write_u32le(0)?;
	140	},
	141	_ => return Err(MuxerError::UnsupportedFormat),
	142	};
	143	self.stream_info.push(AVIStream {
	144	strh_pos: self.bw.tell(),
405cec9e	145	is_video: strm.get_media_type() == StreamType::Video,
5cd2f537 KS	146	nframes: 0,
5cd2f537 KS	147	max_size: 0,
04fafc56	148	pal_change: false,
5cd2f537 KS	149	});
	150
	151	self.bw.write_u32le(0)?; // flags
	152	self.bw.write_u16le(0)?; // priority
	153	self.bw.write_u16le(0)?; // language
	154	self.bw.write_u32le(0)?; // initial frames
405cec9e KS	155	self.bw.write_u32le(strm.tb_num)?;
405cec9e KS	156	self.bw.write_u32le(strm.tb_den)?;
5cd2f537 KS	157	self.bw.write_u32le(0)?; // start
	158	self.bw.write_u32le(0)?; // length
	159	self.bw.write_u32le(0)?; // suggested buffer size
	160	self.bw.write_u32le(0)?; // quality
	161	self.bw.write_u32le(0)?; // sample_size
	162	self.bw.write_u16le(0)?; // x
	163	self.bw.write_u16le(0)?; // y
	164	self.bw.write_u16le(0)?; // w
	165	self.bw.write_u16le(0)?; // h
	166
	167	self.bw.write_buf(b"strf")?;
	168	self.bw.write_u32le(0)?;
	169	let strf_pos = self.bw.tell();
405cec9e	170	match strm.get_media_type() {
5cd2f537	171	StreamType::Video => {
405cec9e	172	let vinfo = strm.get_info().get_properties().get_video_info().unwrap();
5cd2f537 KS	173	let hdr_pos = self.bw.tell();
	174	self.bw.write_u32le(0)?;
	175	self.bw.write_u32le(vinfo.width as u32)?;
e5b5248d	176	self.bw.write_u32le(vinfo.height as u32)?;
c905da8d KS	177	if !vinfo.format.palette {
	178	self.bw.write_u16le(vinfo.format.components as u16)?;
	179	self.bw.write_u16le(vinfo.format.get_total_depth() as u16)?;
	180	} else {
	181	self.bw.write_u16le(1)?;
	182	self.bw.write_u16le(8)?;
	183	}
405cec9e	184	let fcc = find_avi_fourcc(strm.get_info().get_name());
5cd2f537 KS	185	if fcc.is_none() {
	186	return Err(MuxerError::UnsupportedFormat);
	187	}
	188	self.bw.write_buf(&fcc.unwrap_or([0; 4]))?;
	189	self.bw.write_u32le(0)?; // image size
	190	self.bw.write_u32le(0)?; // x dpi
	191	self.bw.write_u32le(0)?; // y dpi
	192	if vinfo.format.palette {
a1613eee	193	self.bw.write_u32le(256)?; // total colors
5cd2f537	194	self.bw.write_u32le(0)?; // important colors
a1613eee KS	195	self.pal_pos[strno] = self.bw.tell() as u32;
	196	for _ in 0..256 {
	197	self.bw.write_u32le(0)?;
	198	}
5cd2f537 KS	199	} else {
	200	self.bw.write_u32le(0)?; // total colors
	201	self.bw.write_u32le(0)?; // important colors
	202	}
405cec9e	203	if let Some(ref edata) = strm.get_info().get_extradata() {
5cd2f537 KS	204	self.bw.write_buf(edata.as_slice())?;
	205	}
	206	let bisize = self.bw.tell() - hdr_pos;
	207	self.bw.seek(SeekFrom::Current(-(bisize as i64)))?;
	208	self.bw.write_u32le(bisize as u32)?;
	209	self.bw.seek(SeekFrom::End(0))?;
	210	},
	211	StreamType::Audio => {
405cec9e KS	212	let ainfo = strm.get_info().get_properties().get_audio_info().unwrap();
405cec9e KS	213	let twocc = find_wav_twocc(strm.get_info().get_name());
5cd2f537 KS	214	if twocc.is_none() {
	215	return Err(MuxerError::UnsupportedFormat);
	216	}
586ea187	217	self.bw.write_u16le(twocc.unwrap_or(0))?;
5cd2f537 KS	218	self.bw.write_u16le(ainfo.channels as u16)?;
	219	self.bw.write_u32le(ainfo.sample_rate)?;
	220	self.bw.write_u32le(0)?; // avg bytes per second
	221	self.bw.write_u16le(ainfo.block_len as u16)?;
	222	self.bw.write_u16le(ainfo.format.bits as u16)?;
405cec9e	223	if let Some(ref edata) = strm.get_info().get_extradata() {
5cd2f537 KS	224	self.bw.write_buf(edata.as_slice())?;
	225	}
	226	},
	227	StreamType::Subtitles => {
405cec9e	228	if let Some(ref edata) = strm.get_info().get_extradata() {
5cd2f537 KS	229	self.bw.write_buf(edata.as_slice())?;
	230	}
	231	},
	232	_ => unreachable!(),
	233	};
6f263099 KS	234	patch_size(self.bw, strf_pos)?;
6f263099 KS	235	patch_size(self.bw, strl_pos)?;
5cd2f537	236	}
6f263099	237	patch_size(self.bw, hdrl_pos)?;
5cd2f537 KS	238
	239	self.data_pos = self.bw.tell() + 8;
	240	self.bw.write_buf(b"LIST\0\0\0\0movi")?;
	241
	242	Ok(())
	243	}
	244	fn mux_frame(&mut self, _strmgr: &StreamManager, pkt: NAPacket) -> MuxerResult<()> {
	245	if self.data_pos == 0 {
	246	return Err(MuxerError::NotCreated);
	247	}
817e4872 KS	248	let stream = pkt.get_stream();
817e4872 KS	249	let str_num = stream.get_num();
5cd2f537 KS	250	if str_num > 99 \|\| str_num >= self.stream_info.len() {
	251	return Err(MuxerError::UnsupportedFormat);
	252	}
	253
	254	let chunk_len = pkt.get_buffer().len() as u32;
	255
a1613eee KS	256	if self.pal_pos[str_num] != 0 {
	257	for sdata in pkt.side_data.iter() {
	258	if let NASideData::Palette(_, ref pal) = sdata {
	259	let cur_pos = self.bw.tell();
	260	self.bw.seek(SeekFrom::Start(u64::from(self.pal_pos[str_num])))?;
bfe6df94 KS	261	for quad in pal.chunks(4) {
	262	self.bw.write_byte(quad[2])?;
	263	self.bw.write_byte(quad[1])?;
	264	self.bw.write_byte(quad[0])?;
	265	self.bw.write_byte(0)?;
	266	}
a1613eee KS	267	self.bw.seek(SeekFrom::Start(cur_pos))?;
	268	self.pal_pos[str_num] = 0;
	269	break;
	270	}
	271	}
	272	} else {
	273	for sdata in pkt.side_data.iter() {
	274	if let NASideData::Palette(true, ref pal) = sdata {
	275	//todo search for changed region
	276	let start_clr = 0usize;
	277	let end_clr = 256usize;
	278	if start_clr < end_clr {
	279	let chunk_len = ((end_clr - start_clr) as u32) * 4 + 4;
	280	self.bw.write_byte(b'0' + ((str_num / 10) as u8))?;
	281	self.bw.write_byte(b'0' + ((str_num % 10) as u8))?;
	282	self.bw.write_buf(b"pc")?;
	283	self.bw.write_u32le(chunk_len)?;
	284	self.bw.write_byte(start_clr as u8)?;
	285	self.bw.write_byte((end_clr - start_clr) as u8)?;
	286	self.bw.write_u16le(0)?; //flags
	287	self.bw.write_buf(&pal[start_clr * 4..end_clr * 4])?;
04fafc56	288	self.stream_info[str_num].pal_change = true;
a1613eee KS	289	}
	290	}
	291	}
	292	}
	293
5cd2f537 KS	294	self.stream_info[str_num].nframes += 1;
5cd2f537 KS	295	self.stream_info[str_num].max_size = self.stream_info[str_num].max_size.max(chunk_len);
5cd2f537 KS	296	self.index.push(IdxEntry {
5cd2f537 KS	297	stream: str_num as u32,
817e4872	298	stype: stream.get_media_type(),
5cd2f537 KS	299	key: pkt.keyframe,
	300	pos: self.bw.tell() as u32,
	301	len: chunk_len });
817e4872	302	write_chunk_hdr(self.bw, stream.get_media_type(), str_num as u32)?;
5cd2f537 KS	303	self.bw.write_u32le(chunk_len)?;
5cd2f537 KS	304	self.bw.write_buf(pkt.get_buffer().as_slice())?;
77d52ff6 KS	305	if (self.bw.tell() & 1) != 0 {
	306	self.bw.write_byte(0)?;
	307	}
5cd2f537 KS	308	Ok(())
	309	}
	310	fn flush(&mut self) -> MuxerResult<()> {
	311	Ok(())
	312	}
	313	fn end(&mut self) -> MuxerResult<()> {
6f263099	314	patch_size(self.bw, self.data_pos)?;
61cab15b	315	if !self.index.is_empty() {
5cd2f537 KS	316	self.bw.write_buf(b"idx1")?;
	317	self.bw.write_u32le((self.index.len() * 16) as u32)?;
	318	for item in self.index.iter() {
6f263099	319	write_chunk_hdr(self.bw, item.stype, item.stream)?;
5cd2f537 KS	320	self.bw.write_u32le(if item.key { 0x10 } else { 0 })?;
	321	self.bw.write_u32le(item.pos)?;
	322	self.bw.write_u32le(item.len)?;
	323	}
	324	}
6f263099	325	patch_size(self.bw, 8)?;
5cd2f537 KS	326	let mut max_frames = 0;
	327	let mut max_size = 0;
	328	for stri in self.stream_info.iter() {
	329	max_frames = max_frames.max(stri.nframes);
	330	max_size = max_size.max(stri.max_size);
04fafc56 KS	331	if stri.pal_change {
	332	self.bw.seek(SeekFrom::Start(stri.strh_pos))?;
	333	self.bw.write_u32le(0x00010000)?;
	334	}
5cd2f537 KS	335	self.bw.seek(SeekFrom::Start(stri.strh_pos + 0x18))?;
	336	self.bw.write_u32le(if stri.is_video { stri.nframes } else { 0 })?;
	337	self.bw.write_u32le(stri.max_size)?;
	338	}
	339	self.bw.seek(SeekFrom::Start(0x30))?;
	340	self.bw.write_u32le(max_frames)?;
	341	self.bw.seek(SeekFrom::Current(8))?;
	342	self.bw.write_u32le(max_size)?;
	343	Ok(())
	344	}
	345	}
	346
dc80f48e KS	347	impl<'a> NAOptionHandler for AVIMuxer<'a> {
	348	fn get_supported_options(&self) -> &[NAOptionDefinition] { &[] }
	349	fn set_options(&mut self, _options: &[NAOption]) { }
	350	fn query_option_value(&self, _name: &str) -> Option<NAValue> { None }
	351	}
	352
5cd2f537 KS	353	pub struct AVIMuxerCreator {}
	354
	355	impl MuxerCreator for AVIMuxerCreator {
	356	fn new_muxer<'a>(&self, bw: &'a mut ByteWriter<'a>) -> Box<dyn MuxCore<'a> + 'a> {
	357	Box::new(AVIMuxer::new(bw))
	358	}
	359	fn get_name(&self) -> &'static str { "avi" }
f0081142	360	fn get_capabilities(&self) -> MuxerCapabilities { MuxerCapabilities::Universal }
5cd2f537 KS	361	}
	362
	363	#[cfg(test)]
	364	mod test {
68e5a4ca	365	use nihav_core::codecs::*;
5cd2f537	366	use nihav_core::demuxers::*;
68e5a4ca KS	367	use nihav_core::muxers::*;
	368	use nihav_codec_support::test::enc_video::*;
	369	use crate::*;
5cd2f537 KS	370
	371	#[test]
	372	fn test_avi_muxer() {
	373	let mut dmx_reg = RegisteredDemuxers::new();
	374	generic_register_all_demuxers(&mut dmx_reg);
886cde48	375	//test sample: https://samples.mplayerhq.hu/V-codecs/RT21/320x240/laser05.avi
68e5a4ca KS	376	let dec_config = DecoderTestParams {
	377	demuxer: "avi",
	378	in_name: "assets/Indeo/laser05.avi",
	379	limit: None,
	380	stream_type: StreamType::None,
	381	dmx_reg, dec_reg: RegisteredDecoders::new(),
	382	};
	383	let mut mux_reg = RegisteredMuxers::new();
	384	generic_register_all_muxers(&mut mux_reg);
	385	/*let enc_config = EncoderTestParams {
	386	muxer: "avi",
	387	enc_name: "",
	388	out_name: "muxed.avi",
	389	mux_reg, enc_reg: RegisteredEncoders::new(),
	390	};
	391	test_remuxing(&dec_config, &enc_config);*/
	392	test_remuxing_md5(&dec_config, "avi", &mux_reg,
	393	[0xa0fb0e47, 0x412e24dd, 0x6b89711c, 0x276fb799]);
5cd2f537 KS	394	}
5cd2f537 KS	395	}