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, pub frames_per_chunk: u16,
20 pub vbr: bool,
21 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, 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 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}