nihav-codec-support/src/test/wavwriter.rs

   1 //! Audio output in WAV format.
   2 use nihav_core::io::byteio::*;
   3 use nihav_core::frame::*;
   4 use std::io::SeekFrom;
   5
   6 /// WAVE output writer.
   7 pub struct WavWriter<'a> {
   8     io: &'a mut ByteWriter<'a>,
   9     data_pos: u64,
  10 }
  11
  12 fn write_byte(wr: &mut ByteWriter, sample: u8) -> ByteIOResult<()> {
  13     wr.write_byte(sample)
  14 }
  15
  16 fn write_s16(wr: &mut ByteWriter, sample: i16) -> ByteIOResult<()> {
  17     wr.write_u16le(sample as u16)
  18 }
  19
  20 fn write_s32(wr: &mut ByteWriter, sample: i32) -> ByteIOResult<()> {
  21     wr.write_u16le((sample >> 16) as u16)
  22 }
  23
  24 fn write_f32(wr: &mut ByteWriter, sample: f32) -> ByteIOResult<()> {
  25     let mut out = (sample * 32768.0) as i32;
  26     if out < -32768 { out = -32768; }
  27     if out >  32767 { out =  32767; }
  28     if out < 0 { out += 65536; }
  29     wr.write_u16le(out as u16)
  30 }
  31
  32 macro_rules! write_data {
  33     ($wr:expr, $buf:expr, $write:ident) => ({
  34         let len = $buf.get_length();
  35         let ainfo = $buf.get_info();
  36         let nch = ainfo.get_channels() as usize;
  37         let mut offs: Vec<usize> = Vec::with_capacity(nch);
  38         for ch in 0..nch { offs.push($buf.get_offset(ch)); }
  39         let is_planar = $buf.get_step() == 1;
  40         let data = $buf.get_data();
  41
  42         if is_planar {
  43             for i in 0..len {
  44                 for ch in 0..nch {
  45                     let sample = data[offs[ch] + i];
  46                     $write($wr, sample)?;
  47                 }
  48             }
  49         } else {
  50             for i in 0..len*nch {
  51                 let sample = data[i];
  52                 $write($wr, sample)?;
  53             }
  54         }
  55     })
  56 }
  57
  58 impl<'a> WavWriter<'a> {
  59     /// Constructs a new `WavWriter` instance.
  60     pub fn new(io: &'a mut ByteWriter<'a>) -> Self {
  61         WavWriter { io, data_pos: 0 }
  62     }
  63     /// Writes audio format information to the file header.
  64     ///
  65     /// This function should be called exactly once before writing actual audio data.
  66     pub fn write_header(&mut self, ainfo: NAAudioInfo) -> ByteIOResult<()> {
  67         let bits = ainfo.get_format().get_bits() as usize;
  68
  69         self.io.write_buf(b"RIFF")?;
  70         self.io.write_u32le(0)?;
  71         self.io.write_buf(b"WAVE")?;
  72
  73         self.io.write_buf(b"fmt ")?;
  74         self.io.write_u32le(16)?;
  75         self.io.write_u16le(0x0001)?; // PCM
  76         self.io.write_u16le(u16::from(ainfo.get_channels()))?;
  77         self.io.write_u32le(ainfo.get_sample_rate())?;
  78
  79         if bits < 16 {
  80             self.io.write_u32le(u32::from(ainfo.get_channels()) * ainfo.get_sample_rate())?;
  81             self.io.write_u16le(u16::from(ainfo.get_channels()))?; // block align
  82             self.io.write_u16le(8)?;
  83         } else {
  84             self.io.write_u32le(2 * u32::from(ainfo.get_channels()) * ainfo.get_sample_rate())?;
  85             self.io.write_u16le(u16::from(2 * ainfo.get_channels()))?; // block align
  86             self.io.write_u16le(16)?;
  87         }
  88
  89         self.io.write_buf(b"data")?;
  90         self.io.write_u32le(0)?;
  91
  92         self.data_pos = self.io.tell();
  93         Ok(())
  94     }
  95     /// Writes audio data.
  96     pub fn write_frame(&mut self, abuf: NABufferType) -> ByteIOResult<()> {
  97         match abuf {
  98             NABufferType::AudioU8(ref buf) => {
  99                 write_data!(self.io, buf, write_byte);
 100             }
 101             NABufferType::AudioI16(ref buf) => {
 102                 write_data!(self.io, buf, write_s16);
 103             }
 104             NABufferType::AudioI32(ref buf) => {
 105                 write_data!(self.io, buf, write_s32);
 106             }
 107             NABufferType::AudioF32(ref buf) => {
 108                 write_data!(self.io, buf, write_f32);
 109             }
 110             NABufferType::AudioPacked(ref buf) => {
 111                 self.io.write_buf(buf.get_data().as_slice())?;
 112             }
 113             _ => {},
 114         };
 115         Ok(())
 116     }
 117 }
 118
 119 impl<'a> Drop for WavWriter<'a> {
 120     #[allow(unused_variables)]
 121     fn drop(&mut self) {
 122         let size = self.io.tell();
 123         if (self.data_pos > 0) && (size >= self.data_pos) {
 124             let res = self.io.seek(SeekFrom::Start(4));
 125             let res = self.io.write_u32le((size - 8) as u32);
 126             let res = self.io.seek(SeekFrom::Start(self.data_pos - 4));
 127             let res = self.io.write_u32le(((size as u64) - self.data_pos) as u32);
 128         }
 129     }
 130 }