nihav-commonfmt/src/demuxers/wav.rs

   1 use nihav_core::demuxers::*;
   2 use nihav_registry::register;
   3 use nihav_core::demuxers::DemuxerError::*;
   4
   5 macro_rules! mktag {
   6     ($a:expr, $b:expr, $c:expr, $d:expr) => {
   7         (($a as u32) << 24) | (($b as u32) << 16) | (($c as u32) << 8) | ($d as u32)
   8     };
   9     ($arr:expr) => {
  10         (($arr[0] as u32) << 24) | (($arr[1] as u32) << 16) | (($arr[2] as u32) << 8) | ($arr[3] as u32)
  11     };
  12 }
  13
  14 struct WAVDemuxer<'a> {
  15     src:            &'a mut ByteReader<'a>,
  16     data_pos:       u64,
  17     data_end:       u64,
  18     srate:          u32,
  19     block_size:     usize,
  20     is_pcm:         bool,
  21     avg_bytes:      u32,
  22 }
  23
  24 impl<'a> DemuxCore<'a> for WAVDemuxer<'a> {
  25     fn open(&mut self, strmgr: &mut StreamManager, seek_index: &mut SeekIndex) -> DemuxerResult<()> {
  26         let riff                        = self.src.read_u32be()?;
  27         let riff_size                   = self.src.read_u32le()? as usize;
  28         let riff_end = self.src.tell() + if riff_size > 0 { (riff_size as u64) } else { u64::from(std::u32::MAX) };
  29         let wave                        = self.src.read_u32be()?;
  30         validate!(riff == mktag!(b"RIFF"));
  31         validate!(wave == mktag!(b"WAVE"));
  32
  33         seek_index.mode = SeekIndexMode::Automatic;
  34
  35         let mut fmt_parsed = false;
  36         let mut _duration = 0;
  37         while self.src.tell() < riff_end {
  38             let ctype                   = self.src.read_tag()?;
  39             let csize                   = self.src.read_u32le()? as usize;
  40             match &ctype {
  41                 b"fmt " => {
  42                     validate!(!fmt_parsed);
  43                     self.parse_fmt(strmgr, csize)?;
  44                     fmt_parsed = true;
  45                 },
  46                 b"fact" => {
  47                     validate!(csize == 4);
  48                     _duration           = self.src.read_u32le()? as usize;
  49                 },
  50                 b"data" => {
  51                     validate!(fmt_parsed);
  52                     self.data_pos = self.src.tell();
  53                     self.data_end = self.data_pos + (csize as u64);
  54                     return Ok(());
  55                 },
  56                 _ => {
  57                                           self.src.read_skip(csize)?;
  58                 },
  59             };
  60         }
  61         Err(DemuxerError::InvalidData)
  62     }
  63
  64     fn get_frame(&mut self, strmgr: &mut StreamManager) -> DemuxerResult<NAPacket> {
  65         if self.src.tell() >= self.data_end {
  66             return Err(DemuxerError::EOF);
  67         }
  68         let str = strmgr.get_stream(0);
  69         if str.is_none() { return Err(InvalidData); }
  70         let stream = str.unwrap();
  71         let ts = NATimeInfo::new(None, None, None, 1, self.srate);
  72         if self.is_pcm {
  73             let mut bsize = self.block_size;
  74             while bsize < 256 {
  75                 bsize <<= 1;
  76             }
  77             let mut buf = vec![0; bsize];
  78             let size                    = self.src.read_buf_some(buf.as_mut_slice())?;
  79             buf.truncate(size);
  80             Ok(NAPacket::new(stream, ts, true, buf))
  81         } else {
  82             self.src.read_packet(stream, ts, true, self.block_size)
  83         }
  84     }
  85
  86     fn seek(&mut self, time: u64, _seek_index: &SeekIndex) -> DemuxerResult<()> {
  87         if self.block_size != 0 && self.avg_bytes != 0 {
  88             let seek_dst = u64::from(self.avg_bytes) * time / 1000;
  89             let seek_off = seek_dst / (self.block_size as u64) * (self.block_size as u64);
  90             self.src.seek(SeekFrom::Start(self.data_pos + seek_off))?;
  91             Ok(())
  92         } else {
  93             Err(DemuxerError::NotImplemented)
  94         }
  95     }
  96 }
  97
  98 impl<'a> WAVDemuxer<'a> {
  99     fn new(io: &'a mut ByteReader<'a>) -> Self {
 100         WAVDemuxer {
 101             src:        io,
 102             data_pos:   0,
 103             data_end:   0,
 104             srate:      0,
 105             block_size: 0,
 106             is_pcm:     false,
 107             avg_bytes:  0,
 108         }
 109     }
 110     fn parse_fmt(&mut self, strmgr: &mut StreamManager, csize: usize) -> DemuxerResult<()> {
 111         validate!(csize >= 14);
 112         let format_tag                  = self.src.read_u16le()?;
 113         let channels                    = self.src.read_u16le()?;
 114         validate!(channels < 256);
 115         let samples_per_sec             = self.src.read_u32le()?;
 116         let avg_bytes_per_sec           = self.src.read_u32le()?;
 117         let block_align                 = self.src.read_u16le()? as usize;
 118         if block_align == 0 {
 119             return Err(DemuxerError::NotImplemented);
 120         }
 121         let bits_per_sample             = if csize >= 16 { self.src.read_u16le()? } else { 8 };
 122         validate!(channels < 256);
 123
 124         let edata;
 125         if csize > 16 {
 126             validate!(csize >= 18);
 127             let cb_size                 = self.src.read_u16le()? as usize;
 128             let mut buf = vec![0; cb_size];
 129                                           self.src.read_buf(buf.as_mut_slice())?;
 130             edata = Some(buf);
 131         } else {
 132             edata = None;
 133         }
 134
 135         let cname = register::find_codec_from_wav_twocc(format_tag).unwrap_or("unknown");
 136         let soniton = if cname == "pcm" {
 137                 if format_tag != 0x0003 {
 138                     if bits_per_sample == 8 {
 139                         NASoniton::new(8, 0)
 140                     } else {
 141                         NASoniton::new(bits_per_sample as u8, SONITON_FLAG_SIGNED)
 142                     }
 143                 } else {
 144                     NASoniton::new(bits_per_sample as u8, SONITON_FLAG_FLOAT)
 145                 }
 146             } else {
 147                 NASoniton::new(bits_per_sample as u8, SONITON_FLAG_SIGNED)
 148             };
 149         let ahdr = NAAudioInfo::new(samples_per_sec, channels as u8, soniton, block_align);
 150         let ainfo = NACodecInfo::new(cname, NACodecTypeInfo::Audio(ahdr), edata);
 151         let res = strmgr.add_stream(NAStream::new(StreamType::Audio, 0, ainfo, 1, samples_per_sec));
 152         if res.is_none() { return Err(MemoryError); }
 153
 154         self.srate = samples_per_sec;
 155         self.block_size = block_align;
 156         self.avg_bytes = avg_bytes_per_sec;
 157         self.is_pcm = cname == "pcm";
 158
 159         Ok(())
 160     }
 161 }
 162
 163 pub struct WAVDemuxerCreator { }
 164
 165 impl DemuxerCreator for WAVDemuxerCreator {
 166     fn new_demuxer<'a>(&self, br: &'a mut ByteReader<'a>) -> Box<dyn DemuxCore<'a> + 'a> {
 167         Box::new(WAVDemuxer::new(br))
 168     }
 169     fn get_name(&self) -> &'static str { "wav" }
 170 }
 171
 172 #[cfg(test)]
 173 mod test {
 174     use super::*;
 175     use std::fs::File;
 176
 177     #[test]
 178     fn test_wav_demux() {
 179         let mut file = File::open("assets/MS/scatter.wav").unwrap();
 180         let mut fr = FileReader::new_read(&mut file);
 181         let mut br = ByteReader::new(&mut fr);
 182         let mut dmx = WAVDemuxer::new(&mut br);
 183         let mut sm = StreamManager::new();
 184         let mut si = SeekIndex::new();
 185         dmx.open(&mut sm, &mut si).unwrap();
 186
 187         loop {
 188             let pktres = dmx.get_frame(&mut sm);
 189             if let Err(e) = pktres {
 190                 if e == DemuxerError::EOF { break; }
 191                 panic!("error");
 192             }
 193             let pkt = pktres.unwrap();
 194             println!("Got {}", pkt);
 195         }
 196     }
 197 }