zk-audio 0.1.0

Audio processing library for voice recording and enhancement
Documentation
use crate::contracts::AudioProcessor;
use crate::core::{AudioError, AudioFrame, AudioProfile, AudioResult, AudioSpec};
use crate::profiles::profile_tuning;
use biquad::{Biquad, Coefficients, DirectForm1, Hertz, ToHertz, Type, Q_BUTTERWORTH_F32};

pub struct VoiceFinisherProcessor {
    profile: AudioProfile,
    warmth_low_pass: Option<DirectForm1<f32>>,
    harsh_low_pass: Option<DirectForm1<f32>>,
    envelope: f32,
    glue_gain: f32,
    harshness_reduction: f32,
}

impl VoiceFinisherProcessor {
    pub fn new(profile: AudioProfile) -> Self {
        Self {
            profile,
            warmth_low_pass: None,
            harsh_low_pass: None,
            envelope: 0.0,
            glue_gain: 1.0,
            harshness_reduction: 0.0,
        }
    }
}

impl AudioProcessor for VoiceFinisherProcessor {
    fn name(&self) -> &'static str {
        "voice_finisher"
    }

    fn prepare(&mut self, spec: AudioSpec) -> AudioResult<()> {
        let tuning = profile_tuning(self.profile);
        let sample_rate = Hertz::<f32>::from_hz(spec.sample_rate as f32)
            .map_err(|_| AudioError::new("Invalid sample rate for voice finisher"))?;
        let warmth = Coefficients::<f32>::from_params(
            Type::LowPass,
            sample_rate,
            900.0f32.hz(),
            Q_BUTTERWORTH_F32,
        )
        .map_err(|_| AudioError::new("Invalid warmth configuration for voice finisher"))?;
        let harsh = Coefficients::<f32>::from_params(
            Type::LowPass,
            sample_rate,
            tuning.finisher_harshness_hz.hz(),
            Q_BUTTERWORTH_F32,
        )
        .map_err(|_| AudioError::new("Invalid harshness configuration for voice finisher"))?;

        self.warmth_low_pass = Some(DirectForm1::<f32>::new(warmth));
        self.harsh_low_pass = Some(DirectForm1::<f32>::new(harsh));
        self.envelope = 0.0;
        self.glue_gain = 1.0;
        self.harshness_reduction = 0.0;
        Ok(())
    }

    fn process(&mut self, frame: &mut AudioFrame) -> AudioResult<()> {
        if frame.samples.is_empty() {
            return Ok(());
        }

        let warmth = self
            .warmth_low_pass
            .as_mut()
            .ok_or_else(|| AudioError::new("Voice finisher warmth filter not prepared"))?;
        let harsh = self
            .harsh_low_pass
            .as_mut()
            .ok_or_else(|| AudioError::new("Voice finisher harshness filter not prepared"))?;

        let tuning = profile_tuning(self.profile);
        let mut harsh_energy = 0.0;
        let mut rms_acc = 0.0;
        let mut peak = 0.0f32;
        for sample in &frame.samples {
            let harsh_low = harsh.run(*sample);
            let harsh_band = *sample - harsh_low;
            harsh_energy += harsh_band * harsh_band;
            rms_acc += sample * sample;
            peak = peak.max(sample.abs());
        }
        let rms = (rms_acc / frame.samples.len() as f32).sqrt();
        let harshness = (harsh_energy / frame.samples.len() as f32).sqrt() / (rms + 1e-4);

        let over_threshold = (rms - tuning.finisher_glue_threshold).max(0.0);
        let target_glue = if over_threshold <= 0.0 {
            1.0
        } else {
            let compressed =
                1.0 / (1.0 + over_threshold * (tuning.finisher_glue_ratio - 1.0) * 4.0);
            compressed.clamp(0.80, 1.0)
        };
        let glue_smoothing = if target_glue < self.glue_gain {
            tuning.finisher_glue_attack
        } else {
            tuning.finisher_glue_release
        };
        let start_glue = self.glue_gain;
        self.glue_gain = start_glue * (1.0 - glue_smoothing) + target_glue * glue_smoothing;

        let harsh_target =
            ((harshness - 0.34) / 0.78).clamp(0.0, 1.0) * tuning.finisher_harshness_strength;
        let harsh_smoothing = if harsh_target > self.harshness_reduction {
            0.24
        } else {
            0.10
        };
        let start_harsh = self.harshness_reduction;
        self.harshness_reduction =
            start_harsh * (1.0 - harsh_smoothing) + harsh_target * harsh_smoothing;

        let len = frame.samples.len().max(1) as f32;
        let warmth_gain = 10.0f32.powf((tuning.finisher_warmth_db * 0.5) / 20.0);
        let presence_gain = 10.0f32.powf((tuning.finisher_presence_db * 0.5) / 20.0);
        let saturation = tuning.finisher_saturation;
        for (index, sample) in frame.samples.iter_mut().enumerate() {
            let t = index as f32 / len;
            let glue = start_glue + (self.glue_gain - start_glue) * t;
            let harsh_reduction = start_harsh + (self.harshness_reduction - start_harsh) * t;
            let low = warmth.run(*sample);
            let high = *sample - low;
            let contoured = low * warmth_gain + high * presence_gain * (1.0 - harsh_reduction);
            let driven = contoured * glue;
            let saturated = if saturation > 0.0 {
                (driven * (1.0 + saturation * 1.8)).tanh() / (1.0 + saturation * 0.6)
            } else {
                driven
            };
            self.envelope = self.envelope * 0.90 + saturated.abs() * 0.10;
            let safety = if self.envelope > tuning.finisher_output_ceiling {
                tuning.finisher_output_ceiling / self.envelope.max(1e-4)
            } else {
                1.0
            };
            *sample = (saturated * safety).clamp(
                -tuning.finisher_output_ceiling,
                tuning.finisher_output_ceiling,
            );
        }

        if peak > tuning.finisher_output_ceiling {
            let trim = tuning.finisher_output_ceiling / peak.max(1e-4);
            for sample in &mut frame.samples {
                *sample *= trim;
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::profiles::profile_tuning;

    fn test_spec() -> AudioSpec {
        AudioSpec {
            sample_rate: 44_100,
            channels: 1,
        }
    }

    #[test]
    fn voice_finisher_limits_peaks_and_keeps_signal_bounded() {
        let spec = test_spec();
        let mut processor = VoiceFinisherProcessor::new(AudioProfile::VoiceClean);
        processor.prepare(spec).unwrap();

        let mut frame = AudioFrame {
            samples: (0..512)
                .map(|i| {
                    let t = i as f32 / spec.sample_rate as f32;
                    (std::f32::consts::TAU * 230.0 * t).sin() * 0.55
                        + (std::f32::consts::TAU * 3_600.0 * t).sin() * 0.32
                })
                .collect(),
            spec,
        };
        processor.process(&mut frame).unwrap();

        let peak = frame
            .samples
            .iter()
            .map(|sample| sample.abs())
            .fold(0.0f32, f32::max);
        assert!(peak <= profile_tuning(AudioProfile::VoiceClean).finisher_output_ceiling + 1e-4);
    }
}