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}
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, 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 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}