]> git.nihav.org Git - nihav-tool.git/blob - src/wavwriter.rs
some pretty-printing
[nihav-tool.git] / 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 }