nihav-commonfmt/src/muxers/wav.rs

   1 use nihav_core::muxers::*;
   2 use nihav_registry::register::*;
   3
   4 struct WAVMuxer<'a> {
   5     bw:         &'a mut ByteWriter<'a>,
   6     data_pos:   u64,
   7 }
   8
   9 impl<'a> WAVMuxer<'a> {
  10     fn new(bw: &'a mut ByteWriter<'a>) -> Self {
  11         Self {
  12             bw,
  13             data_pos:   0,
  14         }
  15     }
  16 }
  17
  18 fn patch_size(bw: &mut ByteWriter, pos: u64) -> MuxerResult<()> {
  19     let size = bw.tell() - pos;
  20     bw.seek(SeekFrom::Current(-((size + 4) as i64)))?;
  21     bw.write_u32le(size as u32)?;
  22     bw.seek(SeekFrom::End(0))?;
  23     Ok(())
  24 }
  25
  26 impl<'a> MuxCore<'a> for WAVMuxer<'a> {
  27     fn create(&mut self, strmgr: &StreamManager) -> MuxerResult<()> {
  28         if strmgr.get_num_streams() != 1 {
  29             return Err(MuxerError::InvalidArgument);
  30         }
  31
  32         let stream = strmgr.get_stream(0).unwrap();
  33
  34         if stream.get_info().get_properties().get_audio_info().is_none() {
  35             return Err(MuxerError::InvalidArgument);
  36         }
  37         let ainfo = stream.get_info().get_properties().get_audio_info().unwrap();
  38
  39         let edata_len = if let Some(ref buf) = stream.get_info().get_extradata() { buf.len() } else { 0 };
  40         if edata_len >= (1 << 16) {
  41             return Err(MuxerError::UnsupportedFormat);
  42         }
  43
  44         let twocc = find_wav_twocc(stream.get_info().get_name());
  45         if twocc.is_none() {
  46             return Err(MuxerError::UnsupportedFormat);
  47         }
  48         let twocc = if stream.get_info().get_name() == "pcm" {
  49                 if !ainfo.format.float { 0x0001 } else { 0x0003 }
  50             } else {
  51                 twocc.unwrap_or(0)
  52             };
  53         let avg_bytes_per_sec = if stream.get_info().get_name() == "pcm" {
  54                 u32::from(ainfo.channels) * ainfo.sample_rate * u32::from(ainfo.format.bits) >> 3
  55             } else {
  56                 0
  57             };
  58
  59         self.bw.write_buf(b"RIFF\0\0\0\0WAVEfmt ")?;
  60         self.bw.write_u32le(if edata_len == 0 { 16 } else { 18 + edata_len } as u32)?;
  61         self.bw.write_u16le(twocc)?;
  62         self.bw.write_u16le(ainfo.channels as u16)?;
  63         self.bw.write_u32le(ainfo.sample_rate)?;
  64         self.bw.write_u32le(avg_bytes_per_sec)?;
  65         self.bw.write_u16le(ainfo.block_len as u16)?;
  66         self.bw.write_u16le(ainfo.format.bits as u16)?;
  67         if let Some(ref buf) = stream.get_info().get_extradata() {
  68             self.bw.write_u16le(edata_len as u16)?;
  69             self.bw.write_buf(buf.as_slice())?;
  70         }
  71         self.bw.write_buf(b"data\0\0\0\0")?;
  72         self.data_pos = self.bw.tell();
  73
  74         Ok(())
  75     }
  76     fn mux_frame(&mut self, _strmgr: &StreamManager, pkt: NAPacket) -> MuxerResult<()> {
  77         if self.data_pos == 0 {
  78             return Err(MuxerError::NotCreated);
  79         }
  80
  81         let stream = pkt.get_stream();
  82         if stream.get_num() != 0 {
  83             return Err(MuxerError::UnsupportedFormat);
  84         }
  85
  86         self.bw.write_buf(pkt.get_buffer().as_slice())?;
  87         Ok(())
  88     }
  89     fn flush(&mut self) -> MuxerResult<()> {
  90         Ok(())
  91     }
  92     fn end(&mut self) -> MuxerResult<()> {
  93         patch_size(&mut self.bw, self.data_pos)?;
  94         patch_size(&mut self.bw, 8)?;
  95         // todo patch avg_bytes_per_second if calculated
  96         // todo write fact value if calculated
  97         Ok(())
  98     }
  99 }
 100
 101 pub struct WAVMuxerCreator {}
 102
 103 impl MuxerCreator for WAVMuxerCreator {
 104     fn new_muxer<'a>(&self, bw: &'a mut ByteWriter<'a>) -> Box<dyn MuxCore<'a> + 'a> {
 105         Box::new(WAVMuxer::new(bw))
 106     }
 107     fn get_name(&self) -> &'static str { "wav" }
 108 }
 109
 110 #[cfg(test)]
 111 mod test {
 112     use super::*;
 113     use std::fs::File;
 114     use nihav_core::demuxers::*;
 115     use crate::demuxers::*;
 116
 117     #[test]
 118     fn test_wav_muxer() {
 119         let mut dmx_reg = RegisteredDemuxers::new();
 120         generic_register_all_demuxers(&mut dmx_reg);
 121         let mut file = File::open("assets/Indeo/laser05.avi").unwrap();
 122         let mut fr = FileReader::new_read(&mut file);
 123         let mut br = ByteReader::new(&mut fr);
 124         let dmx_f = dmx_reg.find_demuxer("avi").unwrap();
 125         let mut dmx = create_demuxer(dmx_f, &mut br).unwrap();
 126
 127         let mut out_sm = StreamManager::new();
 128         let mut out_streamno = 0;
 129         for stream in dmx.get_streams() {
 130             if stream.get_media_type() == StreamType::Audio {
 131                 let mut stream = NAStream::clone(&stream);
 132                 out_streamno = stream.id;
 133                 stream.id = 0;
 134                 out_sm.add_stream(stream);
 135             }
 136         }
 137
 138         let ofile = File::create("assets/test_out/muxed.wav").unwrap();
 139         let mut fw = FileWriter::new_write(ofile);
 140         let mut bw = ByteWriter::new(&mut fw);
 141         let mut mux = WAVMuxer::new(&mut bw);
 142
 143         mux.create(&out_sm).unwrap();
 144
 145         loop {
 146             let pktres = dmx.get_frame();
 147             if let Err(e) = pktres {
 148                 if e == DemuxerError::EOF { break; }
 149                 panic!("error");
 150             }
 151             let mut pkt = pktres.unwrap();
 152             println!("Got {}", pkt);
 153             let pkt_str = pkt.get_stream();
 154             if pkt_str.id == out_streamno {
 155                 pkt.reassign(out_sm.get_stream(0).unwrap(), pkt.get_time_information());
 156                 mux.mux_frame(&out_sm, pkt).unwrap();
 157             }
 158         }
 159
 160         mux.end().unwrap();
 161 panic!("end");
 162     }
 163 }