| 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() as u32)?; |
| 81 | |
| 82 | if bits < 16 { |
| 83 | self.io.write_u32le((ainfo.get_channels() as u32) * (ainfo.get_sample_rate() as u32))?; |
| 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() as u32))?; |
| 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 | if !self.be || self.bits == 8 { |
| 116 | self.io.write_buf(buf.get_data().as_slice())?; |
| 117 | } else { |
| 118 | let data = buf.get_data(); |
| 119 | match self.bits { |
| 120 | 16 => { |
| 121 | for samp in data.chunks(2) { |
| 122 | self.io.write_byte(samp[1])?; |
| 123 | self.io.write_byte(samp[0])?; |
| 124 | } |
| 125 | }, |
| 126 | 24 => { |
| 127 | for samp in data.chunks(3) { |
| 128 | self.io.write_byte(samp[2])?; |
| 129 | self.io.write_byte(samp[1])?; |
| 130 | self.io.write_byte(samp[0])?; |
| 131 | } |
| 132 | }, |
| 133 | 32 => { |
| 134 | for samp in data.chunks(4) { |
| 135 | self.io.write_byte(samp[3])?; |
| 136 | self.io.write_byte(samp[2])?; |
| 137 | self.io.write_byte(samp[1])?; |
| 138 | self.io.write_byte(samp[0])?; |
| 139 | } |
| 140 | }, |
| 141 | _ => unimplemented!(), |
| 142 | }; |
| 143 | } |
| 144 | } |
| 145 | _ => {}, |
| 146 | }; |
| 147 | Ok(()) |
| 148 | } |
| 149 | } |
| 150 | |
| 151 | impl<'a> Drop for WavWriter<'a> { |
| 152 | #[allow(unused_variables)] |
| 153 | fn drop(&mut self) { |
| 154 | let size = self.io.tell(); |
| 155 | if (self.data_pos > 0) && (size >= self.data_pos) { |
| 156 | let res = self.io.seek(SeekFrom::Start(4)); |
| 157 | let res = self.io.write_u32le((size - 8) as u32); |
| 158 | let res = self.io.seek(SeekFrom::Start(self.data_pos - 4)); |
| 159 | let res = self.io.write_u32le(((size as u64) - self.data_pos) as u32); |
| 160 | let res = self.io.flush(); |
| 161 | } |
| 162 | } |
| 163 | } |