1use fundsp::prelude::*;
2
3use crate::abnormal::AbnormalSample;
4use crate::chart::generate_svg;
5use crate::config::{Processing, SnapshotConfig};
6use crate::input::InputSource;
7use crate::wav::generate_wav;
8
9pub fn snapshot_audio_unit<N>(unit: N) -> Vec<u8>
21where
22 N: AudioUnit,
23{
24 snapshot_audio_unit_with_input_and_options(unit, InputSource::None, SnapshotConfig::default())
25}
26
27pub fn snapshot_audio_unit_with_options<N>(unit: N, options: SnapshotConfig) -> Vec<u8>
40where
41 N: AudioUnit,
42{
43 snapshot_audio_unit_with_input_and_options(unit, InputSource::None, options)
44}
45
46pub fn snapshot_audio_unit_with_input<N>(unit: N, input_source: InputSource) -> Vec<u8>
59where
60 N: AudioUnit,
61{
62 snapshot_audio_unit_with_input_and_options(
63 unit,
64 input_source,
65 SnapshotConfig {
66 ..SnapshotConfig::default()
67 },
68 )
69}
70
71pub fn snapshot_audio_unit_with_input_and_options<N>(
85 mut unit: N,
86 mut input_source: InputSource,
87 config: SnapshotConfig,
88) -> Vec<u8>
89where
90 N: AudioUnit,
91{
92 let num_inputs = N::inputs(&unit);
93 let num_outputs = N::outputs(&unit);
94
95 unit.set_sample_rate(config.sample_rate);
96 unit.reset();
97 unit.allocate();
98
99 let input_data = input_source.make_data(num_inputs, config.num_samples);
100
101 let mut output_data: Vec<Vec<f32>> = vec![vec![]; num_outputs];
102
103 let warmup_samples = config
104 .warm_up
105 .warm_up_samples(config.sample_rate, num_inputs);
106
107 let num_warmup_samples = warmup_samples
108 .iter()
109 .map(|ch| ch.len())
110 .next()
111 .unwrap_or_default();
112
113 let mut abnormalities: Vec<Vec<(usize, AbnormalSample)>> = vec![vec![]; num_outputs];
114
115 let mut checked_sample = |mut sample: f32, ch: usize, i: usize| {
116 if sample.is_nan() || sample.is_infinite() {
117 let abnormality = AbnormalSample::from(sample);
118
119 if config.allow_abnormal_samples {
120 abnormalities[ch].push((i, abnormality));
121 sample = 0.0;
122 } else {
123 panic!("Output channel #[{ch}] at sample [{i}] produced [{abnormality}] sample");
124 }
125 }
126 sample
127 };
128
129 (0..num_warmup_samples).for_each(|i| {
130 let mut input_frame = vec![0.0; num_inputs];
131 for ch in 0..num_inputs {
132 input_frame[ch] = warmup_samples[ch][i];
133 }
134 let mut output_frame = vec![0.0; num_outputs];
135 unit.tick(&input_frame, &mut output_frame);
136 });
138
139 match config.processing_mode {
140 Processing::Tick => {
141 (0..config.num_samples).for_each(|i| {
142 let mut input_frame = vec![0.0; num_inputs];
143 for ch in 0..num_inputs {
144 input_frame[ch] = input_data[ch][i];
145 }
146 let mut output_frame = vec![0.0; num_outputs];
147 unit.tick(&input_frame, &mut output_frame);
148 for ch in 0..num_outputs {
149 let sample = checked_sample(output_frame[ch], ch, i);
150 output_data[ch].push(sample);
151 }
152 });
153 }
154 Processing::Batch(batch_size) => {
155 assert!(
156 batch_size <= MAX_BUFFER_SIZE as u8,
157 "Batch size must be less than or equal to [{MAX_BUFFER_SIZE}]"
158 );
159
160 let samples_index = (0..config.num_samples).collect::<Vec<_>>();
161 for chunk in samples_index.chunks(batch_size as usize) {
162 let mut input_buff = BufferVec::new(num_inputs);
163 for (frame_index, input_index) in chunk.iter().enumerate() {
164 for (ch, input) in input_data.iter().enumerate() {
165 let value: f32 = input[*input_index];
166 input_buff.set_f32(ch, frame_index, value);
167 }
168 }
169 let input_ref = input_buff.buffer_ref();
170 let mut output_buf = BufferVec::new(num_outputs);
171 let mut output_ref = output_buf.buffer_mut();
172
173 unit.process(chunk.len(), &input_ref, &mut output_ref);
174
175 for (ch, data) in output_data.iter_mut().enumerate() {
176 data.extend(
177 output_buf
178 .channel_f32(ch)
179 .iter()
180 .enumerate()
181 .map(|(i, &value)| checked_sample(value, ch, i + chunk[0])),
182 );
183 }
184 }
185 }
186 }
187
188 match config.output_mode {
189 crate::config::SnapshotOutputMode::SvgChart(svg_chart_config) => {
190 let start_sample = config.warm_up.num_samples(config.sample_rate);
191
192 generate_svg(
193 &input_data,
194 &output_data,
195 &abnormalities,
196 &svg_chart_config,
197 config.sample_rate,
198 config.num_samples,
199 start_sample,
200 )
201 .as_bytes()
202 .to_vec()
203 }
204 crate::config::SnapshotOutputMode::Wav(wav_output) => generate_wav(
205 &output_data,
206 &wav_output,
207 config.sample_rate,
208 config.num_samples,
209 ),
210 }
211}