moshi_db/
wav.rs

1// Copyright (c) Kyutai, all rights reserved.
2// This source code is licensed under the license found in the
3// LICENSE file in the root directory of this source tree.
4
5use std::io::prelude::*;
6
7pub trait Sample {
8    fn to_i16(&self) -> i16;
9}
10
11impl Sample for f32 {
12    fn to_i16(&self) -> i16 {
13        (self.clamp(-1.0, 1.0) * 32767.0) as i16
14    }
15}
16
17impl Sample for f64 {
18    fn to_i16(&self) -> i16 {
19        (self.clamp(-1.0, 1.0) * 32767.0) as i16
20    }
21}
22
23impl Sample for i16 {
24    fn to_i16(&self) -> i16 {
25        *self
26    }
27}
28pub fn write_wav_header<W: Write>(
29    w: &mut W,
30    sample_rate: u32,
31    chunk_size: u32,
32    data_size: u32,
33) -> std::io::Result<()> {
34    let n_channels = 1u16;
35    let bits_per_sample = 16u16;
36    let byte_rate = sample_rate * n_channels as u32 * (bits_per_sample / 8) as u32;
37    let block_align = n_channels * (bits_per_sample / 8);
38
39    w.write_all(b"RIFF")?;
40    w.write_all(&chunk_size.to_le_bytes())?; // unknown chunk size
41    w.write_all(b"WAVE")?;
42
43    w.write_all(b"fmt ")?;
44    w.write_all(&16u32.to_le_bytes())?;
45    w.write_all(&1u16.to_le_bytes())?; // PCM format
46    w.write_all(&n_channels.to_le_bytes())?;
47    w.write_all(&sample_rate.to_le_bytes())?;
48    w.write_all(&byte_rate.to_le_bytes())?;
49    w.write_all(&block_align.to_le_bytes())?;
50    w.write_all(&bits_per_sample.to_le_bytes())?;
51
52    w.write_all(b"data")?;
53    w.write_all(&data_size.to_le_bytes())?; // unknown data size
54
55    Ok(())
56}
57
58pub fn write_pcm_in_wav<W: Write, S: Sample>(w: &mut W, samples: &[S]) -> std::io::Result<usize> {
59    for sample in samples {
60        w.write_all(&sample.to_i16().to_le_bytes())?
61    }
62    Ok(samples.len() * std::mem::size_of::<i16>())
63}
64
65pub fn write_pcm_as_wav<W: Write, S: Sample>(
66    w: &mut W,
67    samples: &[S],
68    sample_rate: u32,
69) -> std::io::Result<()> {
70    let chunk_size = 12u32 + 24u32 + samples.len() as u32 * 2;
71    let data_size = samples.len() as u32 * 2;
72    write_wav_header(w, sample_rate, chunk_size, data_size)?;
73    write_pcm_in_wav(w, samples)?;
74    Ok(())
75}