Skip to main content

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