sea_codec/
encoder.rs

1use alloc::{rc::Rc, string::String};
2
3use crate::codec::{
4    common::SeaError,
5    file::{SeaFile, SeaFileHeader},
6};
7
8pub enum SeaEncoderState {
9    Start,
10    WritingFrames,
11    Finished,
12}
13
14#[derive(Debug, Clone, PartialEq)]
15pub struct EncoderSettings {
16    pub scale_factor_bits: u8,
17    pub scale_factor_frames: u8,
18    pub residual_bits: f32, // 1-8
19    pub frames_per_chunk: u16,
20    pub vbr: bool,
21}
22
23impl Default for EncoderSettings {
24    fn default() -> Self {
25        Self {
26            frames_per_chunk: 5120,
27            scale_factor_bits: 4,
28            scale_factor_frames: 20,
29            residual_bits: 3.0,
30            vbr: false,
31        }
32    }
33}
34
35trait InternalWrite {
36    fn write_all(&mut self, buf: &[u8]) -> Result<(), SeaError>;
37}
38
39#[cfg(feature = "std")]
40impl<W: std::io::Write> InternalWrite for W {
41    fn write_all(&mut self, buf: &[u8]) -> Result<(), SeaError> {
42        Ok(self.write_all(buf)?)
43    }
44}
45
46#[cfg(not(feature = "std"))]
47impl InternalWrite for &mut alloc::vec::Vec<u8> {
48    fn write_all(&mut self, buf: &[u8]) -> Result<(), SeaError> {
49        self.extend_from_slice(buf);
50        Ok(())
51    }
52}
53
54pub struct SeaEncoder<'inp> {
55    data: &'inp [i16],
56    file: SeaFile,
57    state: SeaEncoderState,
58    written_frames: u32,
59    passed_total_frames: Option<u32>,
60}
61
62impl<'inp> SeaEncoder<'inp> {
63    pub fn from_slice(
64        channels: u8,
65        sample_rate: u32,
66        total_frames: Option<u32>,
67        settings: EncoderSettings,
68        data: &'inp [i16],
69    ) -> Result<Self, SeaError> {
70        let header = SeaFileHeader {
71            version: 1,
72            channels,
73            chunk_size: 0, // will be set later by the first chunk
74            frames_per_chunk: settings.frames_per_chunk,
75            sample_rate,
76            total_frames: total_frames.unwrap_or(0),
77            metadata: Rc::new(String::new()),
78        };
79
80        let file = SeaFile::new(header, &settings)?;
81
82        let state = SeaEncoderState::Start;
83
84        Ok(SeaEncoder {
85            file,
86            state,
87            data,
88            written_frames: 0,
89            passed_total_frames: total_frames,
90        })
91    }
92
93    fn read_samples(&mut self, max_sample_count: usize) -> Result<&'inp [i16], SeaError> {
94        let max_to_read = self.data.len().min(max_sample_count);
95
96        if max_to_read == 0 {
97            return Ok(&self.data[..0]);
98        }
99
100        if !max_to_read.is_multiple_of(self.file.header.channels as usize) {
101            return Err(SeaError::EndOfFile);
102        }
103
104        let (samples, new_data) = self.data.split_at(max_to_read);
105        self.data = new_data;
106
107        Ok(samples)
108    }
109
110    #[cfg(feature = "std")]
111    pub fn encode_frame(&mut self, writer: impl std::io::Write) -> Result<bool, SeaError> {
112        self.encode_frame_inner(writer)
113    }
114
115    #[cfg(not(feature = "std"))]
116    pub fn encode_frame(&mut self, writer: &mut alloc::vec::Vec<u8>) -> Result<bool, SeaError> {
117        self.encode_frame_inner(writer)
118    }
119
120    fn encode_frame_inner<W: InternalWrite>(&mut self, mut writer: W) -> Result<bool, SeaError> {
121        if matches!(self.state, SeaEncoderState::Finished) {
122            return Err(SeaError::EncoderClosed);
123        }
124
125        if matches!(self.state, SeaEncoderState::Start) {
126            if let Some(total_frames) = self.passed_total_frames {
127                if total_frames == 0 {
128                    writer.write_all(&self.file.header.serialize())?;
129                    self.state = SeaEncoderState::WritingFrames;
130                }
131            }
132        }
133
134        let channels = self.file.header.channels;
135        let frames = if self.file.header.total_frames > 0 {
136            (self.file.header.frames_per_chunk as usize)
137                .min(self.file.header.total_frames as usize - self.written_frames as usize)
138        } else {
139            self.file.header.frames_per_chunk as usize
140        };
141
142        let full_size_samples =
143            self.file.header.frames_per_chunk as usize * self.file.header.channels as usize;
144        let samples_to_read = frames * channels as usize;
145        let samples = self.read_samples(samples_to_read)?;
146        let eof: bool = samples.is_empty() || samples.len() < full_size_samples;
147
148        if !samples.is_empty() {
149            let encoded_chunk = self.file.make_chunk(samples)?;
150
151            if eof {
152                assert!(encoded_chunk.len() <= self.file.header.chunk_size as usize);
153            } else {
154                assert_eq!(encoded_chunk.len(), self.file.header.chunk_size as usize);
155            }
156
157            // we need to write file header after the first chunk is generated
158            if matches!(self.state, SeaEncoderState::Start) {
159                writer.write_all(&self.file.header.serialize())?;
160                self.state = SeaEncoderState::WritingFrames;
161            }
162
163            writer.write_all(&encoded_chunk)?;
164            self.written_frames += frames as u32;
165        }
166
167        if eof {
168            self.state = SeaEncoderState::Finished;
169        }
170
171        Ok(!eof)
172    }
173
174    pub fn finalize(&mut self) -> Result<(), SeaError> {
175        self.state = SeaEncoderState::Finished;
176        Ok(())
177    }
178}