src/wavwriter.rs

   1 extern crate nihav_core;
   2
   3 use nihav_core::io::byteio::*;
   4 use nihav_core::frame::*;
   5 use std::fs::File;
   6 use std::io::SeekFrom;
   7
   8 pub struct WavWriter<'a> {
   9     io: Box<ByteWriter<'a>>,
  10     data_pos: u64,
  11     be: bool,
  12     bits: u8,
  13 }
  14
  15 fn write_byte(wr: &mut ByteWriter, sample: u8) -> ByteIOResult<()> {
  16     wr.write_byte(sample)
  17 }
  18
  19 fn write_s16(wr: &mut ByteWriter, sample: i16) -> ByteIOResult<()> {
  20     wr.write_u16le(sample as u16)
  21 }
  22
  23 fn write_s32(wr: &mut ByteWriter, sample: i32) -> ByteIOResult<()> {
  24     wr.write_u16le((sample >> 16) as u16)
  25 }
  26
  27 fn write_f32(wr: &mut ByteWriter, sample: f32) -> ByteIOResult<()> {
  28     let mut out = (sample * 32768.0) as i32;
  29     if out < -32768 { out = -32768; }
  30     if out >  32767 { out =  32767; }
  31     if out < 0 { out += 65536; }
  32     wr.write_u16le(out as u16)
  33 }
  34
  35 macro_rules! write_data {
  36     ($wr:expr, $buf:expr, $write:ident) => ({
  37         let len = $buf.get_length();
  38         let ainfo = $buf.get_info();
  39         let nch = ainfo.get_channels() as usize;
  40         let mut offs: Vec<usize> = Vec::with_capacity(nch);
  41         for ch in 0..nch { offs.push($buf.get_offset(ch)); }
  42         let is_planar = $buf.get_step() == 1;
  43         let data = $buf.get_data();
  44
  45         if is_planar {
  46             for i in 0..len {
  47                 for ch in 0..nch {
  48                     let sample = data[offs[ch] + i];
  49                     $write($wr, sample)?;
  50                 }
  51             }
  52         } else {
  53             for i in 0..len*nch {
  54                 let sample = data[i];
  55                 $write($wr, sample)?;
  56             }
  57         }
  58     })
  59 }
  60
  61 impl<'a> WavWriter<'a> {
  62     pub fn new(name: &str) -> Self {
  63         let file = File::create(name).unwrap();
  64         let file = std::io::BufWriter::new(file);
  65         let fw   = Box::new(FileWriter::new_write(file));
  66         let io   = ByteWriter::new(Box::leak(fw));
  67         WavWriter { io: Box::new(io), data_pos: 0, be: false, bits: 0 }
  68     }
  69     pub fn write_header(&mut self, ainfo: NAAudioInfo) -> ByteIOResult<()> {
  70         let bits = ainfo.get_format().get_bits() as usize;
  71
  72         self.io.write_buf(b"RIFF")?;
  73         self.io.write_u32le(0)?;
  74         self.io.write_buf(b"WAVE")?;
  75
  76         self.io.write_buf(b"fmt ")?;
  77         self.io.write_u32le(16)?;
  78         self.io.write_u16le(0x0001)?; // PCM
  79         self.io.write_u16le(ainfo.get_channels() as u16)?;
  80         self.io.write_u32le(ainfo.get_sample_rate())?;
  81
  82         if bits <= 8 {
  83             self.io.write_u32le((ainfo.get_channels() as u32) * ainfo.get_sample_rate())?;
  84             self.io.write_u16le(ainfo.get_channels() as u16)?; // block align
  85             self.io.write_u16le(8)?;
  86         } else {
  87             self.io.write_u32le(2 * (ainfo.get_channels() as u32) * ainfo.get_sample_rate())?;
  88             self.io.write_u16le((2 * ainfo.get_channels()) as u16)?; // block align
  89             self.io.write_u16le(16)?;
  90         }
  91
  92         self.io.write_buf(b"data")?;
  93         self.io.write_u32le(0)?;
  94
  95         self.bits = bits as u8;
  96         self.be = ainfo.get_format().is_be();
  97         self.data_pos = self.io.tell();
  98         Ok(())
  99     }
 100     pub fn write_frame(&mut self, abuf: NABufferType) -> ByteIOResult<()> {
 101         match abuf {
 102             NABufferType::AudioU8(ref buf) => {
 103                 write_data!(&mut self.io, buf, write_byte);
 104             }
 105             NABufferType::AudioI16(ref buf) => {
 106                 write_data!(&mut self.io, buf, write_s16);
 107             }
 108             NABufferType::AudioI32(ref buf) => {
 109                 write_data!(&mut self.io, buf, write_s32);
 110             }
 111             NABufferType::AudioF32(ref buf) => {
 112                 write_data!(&mut self.io, buf, write_f32);
 113             }
 114             NABufferType::AudioPacked(ref buf) => {
 115                 let data = buf.get_data();
 116                 match self.bits {
 117                     _ if !self.be && (self.bits & 7) == 0 => {
 118                         self.io.write_buf(data.as_slice())?;
 119                     },
 120                     8 => {
 121                         self.io.write_buf(data.as_slice())?;
 122                     },
 123                     12 if !self.be => {
 124                         let mut src = data.chunks_exact(3);
 125                         for chunk in src.by_ref() {
 126                             self.io.write_byte(chunk[0] << 4)?;
 127                             self.io.write_byte((chunk[1] << 4) | (chunk[0] >> 4))?;
 128                             self.io.write_byte(chunk[1] & 0xF0)?;
 129                             self.io.write_byte(chunk[2])?;
 130                         }
 131                         let tail = src.remainder();
 132                         if tail.len() == 2 {
 133                             self.io.write_byte(tail[0] << 4)?;
 134                             self.io.write_byte(tail[1] << 4)?;
 135                         }
 136                     }
 137                     16 => {
 138                         for samp in data.chunks(2) {
 139                             self.io.write_byte(samp[1])?;
 140                             self.io.write_byte(samp[0])?;
 141                         }
 142                     },
 143                     24 => {
 144                         for samp in data.chunks(3) {
 145                             self.io.write_byte(samp[2])?;
 146                             self.io.write_byte(samp[1])?;
 147                             self.io.write_byte(samp[0])?;
 148                         }
 149                     },
 150                     32 => {
 151                         for samp in data.chunks(4) {
 152                             self.io.write_byte(samp[3])?;
 153                             self.io.write_byte(samp[2])?;
 154                             self.io.write_byte(samp[1])?;
 155                             self.io.write_byte(samp[0])?;
 156                         }
 157                     },
 158                     _ => unimplemented!(),
 159                 };
 160             }
 161             _ => {},
 162         };
 163         Ok(())
 164     }
 165 }
 166
 167 impl<'a> Drop for WavWriter<'a> {
 168     #[allow(unused_variables)]
 169     fn drop(&mut self) {
 170         let size = self.io.tell();
 171         if (self.data_pos > 0) && (size >= self.data_pos) {
 172             let res = self.io.seek(SeekFrom::Start(4));
 173             let res = self.io.write_u32le((size - 8) as u32);
 174             let res = self.io.seek(SeekFrom::Start(self.data_pos - 4));
 175             let res = self.io.write_u32le((size - self.data_pos) as u32);
 176             let res = self.io.flush();
 177         }
 178     }
 179 }