firewheel-nodes 0.12.1

use super::{MAX_HZ, MIN_HZ};
use firewheel_core::node::NodeError;
use firewheel_core::{
    StreamInfo,
    channel_config::{ChannelConfig, ChannelCount},
    diff::{Diff, Patch},
    dsp::{
        coeff_update::{CoeffUpdateFactor, CoeffUpdateMask},
        filter::{
            single_pole_iir::{
                OnePoleIirHPFCoeff, OnePoleIirHPFCoeffSimd, OnePoleIirHPFSimd, OnePoleIirLPFCoeff,
                OnePoleIirLPFCoeffSimd, OnePoleIirLPFSimd,
            },
            smoothing_filter::DEFAULT_SMOOTH_SECONDS,
        },
    },
    event::ProcEvents,
    node::{
        AudioNode, AudioNodeInfo, AudioNodeProcessor, ConstructProcessorContext, EmptyConfig,
        ProcBuffers, ProcExtra, ProcInfo, ProcStreamCtx, ProcessStatus,
    },
    param::smoother::{SmoothedParam, SmootherConfig},
};

pub type FastBandpassMonoNode = FastBandpassNode<1>;
pub type FastBandpassStereoNode = FastBandpassNode<2>;

/// A simple single-pole IIR bandpass filter node that is computationally
/// efficient
#[derive(Diff, Patch, Debug, Clone, Copy, PartialEq)]
#[cfg_attr(feature = "bevy", derive(bevy_ecs::prelude::Component))]
#[cfg_attr(feature = "bevy_reflect", derive(bevy_reflect::Reflect))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct FastBandpassNode<const CHANNELS: usize = 2> {
    /// The cutoff frequency in hertz in the range `[20.0, 20480.0]`.
    pub cutoff_hz: f32,

    /// The time in seconds of the internal smoothing filter.
    ///
    /// By default this is set to `0.062` (62ms). This value is chosen such that
    /// the stair-stepping effect isn't noticeable for a typical block size of 1024
    /// samples.
    pub smooth_seconds: f32,

    /// An exponent representing the rate at which DSP coefficients are
    /// updated when parameters are being smoothed.
    ///
    /// Smaller values will produce less "stair-stepping" artifacts,
    /// but will also consume more CPU.
    ///
    /// The resulting number of frames (samples in a single channel of audio)
    /// that will elapse between each update is calculated as
    /// `2^coeff_update_factor`.
    ///
    /// By default this is set to `4`.
    pub coeff_update_factor: CoeffUpdateFactor,
}

impl<const CHANNELS: usize> Default for FastBandpassNode<CHANNELS> {
    fn default() -> Self {
        Self {
            cutoff_hz: 1_000.0,
            smooth_seconds: DEFAULT_SMOOTH_SECONDS,
            coeff_update_factor: CoeffUpdateFactor::default(),
        }
    }
}

impl<const CHANNELS: usize> FastBandpassNode<CHANNELS> {
    /// Construct a new `FastBandpassNode` from the given parameters.
    ///
    /// * `cutoff_hz` - The cutoff frequency in hertz in the range `[20.0, 20480.0]`
    pub const fn from_cutoff_hz(cutoff_hz: f32) -> Self {
        Self {
            cutoff_hz,
            smooth_seconds: DEFAULT_SMOOTH_SECONDS,
            coeff_update_factor: CoeffUpdateFactor::DEFAULT,
        }
    }
}

impl<const CHANNELS: usize> AudioNode for FastBandpassNode<CHANNELS> {
    type Configuration = EmptyConfig;

    fn info(&self, _config: &Self::Configuration) -> Result<AudioNodeInfo, NodeError> {
        Ok(AudioNodeInfo::new()
            .debug_name("fast_bandpass")
            .channel_config(ChannelConfig {
                num_inputs: ChannelCount::new(CHANNELS as u32).unwrap(),
                num_outputs: ChannelCount::new(CHANNELS as u32).unwrap(),
            }))
        // TODO: Once the scheduler gets in-place processing support, use
        // in-place processing for this node.
    }

    fn construct_processor(
        &self,
        _config: &Self::Configuration,
        cx: ConstructProcessorContext,
    ) -> Result<impl AudioNodeProcessor, NodeError> {
        let sample_rate_recip = cx.stream_info.sample_rate_recip as f32;

        let cutoff_hz = self.cutoff_hz.clamp(MIN_HZ, MAX_HZ);

        Ok(Processor {
            lpf: OnePoleIirLPFSimd::default(),
            lpf_coeff: OnePoleIirLPFCoeffSimd::<CHANNELS>::splat(OnePoleIirLPFCoeff::new(
                cutoff_hz,
                sample_rate_recip,
            )),
            hpf: OnePoleIirHPFSimd::default(),
            hpf_coeff: OnePoleIirHPFCoeffSimd::<CHANNELS>::splat(OnePoleIirHPFCoeff::new(
                cutoff_hz,
                sample_rate_recip,
            )),
            cutoff_hz: SmoothedParam::new(
                cutoff_hz,
                MAX_HZ - MIN_HZ,
                SmootherConfig {
                    smooth_seconds: self.smooth_seconds,
                    ..Default::default()
                },
                cx.stream_info.sample_rate,
            ),
            coeff_update_mask: self.coeff_update_factor.mask(),
            cutoff_changed: false,
        })
    }
}

#[cold]
#[inline(never)]
fn calc_coeffs<const CHANNELS: usize>(
    cutoff_hz: f32,
    sample_rate_recip: f32,
) -> (
    OnePoleIirLPFCoeffSimd<CHANNELS>,
    OnePoleIirHPFCoeffSimd<CHANNELS>,
) {
    (
        OnePoleIirLPFCoeffSimd::splat(OnePoleIirLPFCoeff::new(cutoff_hz, sample_rate_recip)),
        OnePoleIirHPFCoeffSimd::splat(OnePoleIirHPFCoeff::new(cutoff_hz, sample_rate_recip)),
    )
}

struct Processor<const CHANNELS: usize> {
    lpf: OnePoleIirLPFSimd<CHANNELS>,
    hpf: OnePoleIirHPFSimd<CHANNELS>,
    lpf_coeff: OnePoleIirLPFCoeffSimd<CHANNELS>,
    hpf_coeff: OnePoleIirHPFCoeffSimd<CHANNELS>,

    cutoff_hz: SmoothedParam,
    coeff_update_mask: CoeffUpdateMask,
    cutoff_changed: bool,
}

impl<const CHANNELS: usize> Processor<CHANNELS> {
    fn reset(&mut self) {
        self.cutoff_hz.reset_to_target();
        self.lpf.reset();
        self.hpf.reset();
    }
}

impl<const CHANNELS: usize> AudioNodeProcessor for Processor<CHANNELS> {
    fn events(&mut self, info: &ProcInfo, events: &mut ProcEvents, _extra: &mut ProcExtra) {
        for patch in events.drain_patches::<FastBandpassNode<CHANNELS>>() {
            match patch {
                FastBandpassNodePatch::CutoffHz(cutoff) => {
                    self.cutoff_changed = true;
                    self.cutoff_hz.set_value(cutoff.clamp(MIN_HZ, MAX_HZ));
                }
                FastBandpassNodePatch::SmoothSeconds(seconds) => {
                    self.cutoff_hz.set_smooth_seconds(seconds, info.sample_rate);
                }
                FastBandpassNodePatch::CoeffUpdateFactor(f) => {
                    self.coeff_update_mask = f.mask();
                }
            }
        }
    }

    fn bypassed(&mut self, _bypassed: bool) {
        self.reset();
    }

    fn process(
        &mut self,
        info: &ProcInfo,
        buffers: ProcBuffers,
        _extra: &mut ProcExtra,
    ) -> ProcessStatus {
        if info.in_silence_mask.all_channels_silent(CHANNELS) {
            // Outputs will be silent, so no need to process.

            // Reset the smoothers and filters since they don't need to smooth any
            // output.
            self.reset();

            return ProcessStatus::ClearAllOutputs;
        }

        assert!(buffers.inputs.len() == CHANNELS);
        assert!(buffers.outputs.len() == CHANNELS);
        for ch in buffers.inputs.iter() {
            assert!(ch.len() >= info.frames);
        }
        for ch in buffers.outputs.iter() {
            assert!(ch.len() >= info.frames);
        }

        if self.cutoff_hz.is_smoothing() {
            for i in 0..info.frames {
                let cutoff_hz = self.cutoff_hz.next_smoothed();

                // Only recalculate coefficients every 2^coeff_update_factor frames
                if self.coeff_update_mask.do_update(i) {
                    (self.lpf_coeff, self.hpf_coeff) =
                        calc_coeffs(cutoff_hz, info.sample_rate_recip as f32);
                }

                let s: [f32; CHANNELS] = core::array::from_fn(|ch_i| {
                    // Safety: These bounds have been checked above.
                    unsafe { *buffers.inputs.get_unchecked(ch_i).get_unchecked(i) }
                });

                let out = self.lpf.process(s, &self.lpf_coeff);
                let out = self.hpf.process(out, &self.hpf_coeff);

                for (ch_i, &o) in out.iter().enumerate().take(CHANNELS) {
                    // Safety: These bounds have been checked above.
                    unsafe {
                        *buffers.outputs.get_unchecked_mut(ch_i).get_unchecked_mut(i) = o;
                    }
                }
            }

            if self.cutoff_hz.settle() {
                (self.lpf_coeff, self.hpf_coeff) =
                    calc_coeffs(self.cutoff_hz.target_value(), info.sample_rate_recip as f32);
            }
        } else {
            // The cutoff parameter is not currently smoothing, so we can optimize by
            // only updating the filter coefficients once.
            if self.cutoff_changed {
                self.cutoff_changed = false;
                (self.lpf_coeff, self.hpf_coeff) =
                    calc_coeffs(self.cutoff_hz.target_value(), info.sample_rate_recip as f32);
            }

            for i in 0..info.frames {
                let s: [f32; CHANNELS] = core::array::from_fn(|ch_i| {
                    // Safety: These bounds have been checked above.
                    unsafe { *buffers.inputs.get_unchecked(ch_i).get_unchecked(i) }
                });

                let out = self.lpf.process(s, &self.lpf_coeff);
                let out = self.hpf.process(out, &self.hpf_coeff);

                for (ch_i, &o) in out.iter().enumerate().take(CHANNELS) {
                    // Safety: These bounds have been checked above.
                    unsafe {
                        *buffers.outputs.get_unchecked_mut(ch_i).get_unchecked_mut(i) = o;
                    }
                }
            }
        }

        ProcessStatus::OutputsModified
    }

    fn new_stream(&mut self, stream_info: &StreamInfo, _context: &mut ProcStreamCtx) {
        self.cutoff_hz.update_sample_rate(stream_info.sample_rate);
        (self.lpf_coeff, self.hpf_coeff) = calc_coeffs(
            self.cutoff_hz.target_value(),
            stream_info.sample_rate_recip as f32,
        );
        self.reset();
    }
}