flexaudio_core/

normalizer.rs

//! Normalize arbitrary device frames (any sample rate/channel count, interleaved f32)
//! in two stages, then convert them back to the requested output format.
//!
//! ```text
//! Input (any sample rate/channel count)
//!   │  Stage 1 (internal normalization; fixed)
//!   │   - Channel mix (→stereo)
//!   │   - Sample-rate conversion (rubato, →48000)
//!   ▼
//! Internal canonical form: f32 / 48000 Hz / stereo / 20ms = 960 frames
//!   │  Stage 2 (output; per tap)
//!   │   - Channel conversion (stereo→mono average / mono→stereo duplicate / unchanged)
//!   │   - Sample-rate conversion (rubato, 48000→output.sample_rate; passthrough if equal)
//!   ▼
//! Output: f32 / output.sample_rate / output.channels / fixed 20ms duration
//!        (48k=960 / 16k=320 / 8k=160 frames)
//! ```
//!
//! With the default output `{48000, 2}`, stage 2 is entirely passthrough (the internal
//! canonical form is emitted as-is). Stage 1 sample-rate conversion is bypassed when
//! `in_sample_rate == 48000`; stage 2 conversion is bypassed when `output.sample_rate == 48000`.
//!
//! Both rubato resamplers use `FixedAsync::Input` (fixed input chunks). Variable-length output
//! is collected in an internal accumulator and split at 20ms boundaries. The resampler and
//! accumulator carry any remainder forward.
//!
//! PTS is assigned by tracking the input-to-output sample-offset ratio from `device_pts`
//! corresponding to the first sample of each output chunk. The stream layer assigns `seq`.

use rubato::audioadapter_buffers::direct::InterleavedSlice;
use rubato::{
    Async, FixedAsync, Indexing, Resampler, SincInterpolationParameters, SincInterpolationType,
    WindowFunction,
};

use crate::types::{
    ChunkFlags, Error, ErrorContext, ErrorGroup, Operation, OutputFormat, Result, CHANNELS,
    SAMPLE_RATE,
};

/// A library-constructed chunk retaining padding provenance.
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct NormalizedChunk {
    /// Interleaved output PCM.
    pub samples: Vec<f32>,
    /// Presentation timestamp of the first frame.
    pub pts_ns: i64,
    /// Integrity flags; padding alone is not a discontinuity.
    pub flags: ChunkFlags,
}

/// Frames per internal canonical-form chunk (20ms @ 48kHz). Stage 1 split boundary.
pub const CHUNK_FRAMES: usize = 960;

/// Channel count in the internal canonical form (stereo).
const INNER_CH: usize = CHANNELS as usize; // 2

/// A processor applied to the internal normalized form (48kHz / stereo /
/// interleaved f32) *before* it is split into the output taps.
///
/// Kept behind a trait so `flexaudio-core` stays independent of any concrete
/// DSP implementation (e.g. noise suppression). The facade injects an
/// implementation; the core only sees this contract. Because the same
/// processed samples feed every output tap, a single processor affects both the
/// primary and secondary outputs.
///
/// [`process`](Self::process) runs on each batch of normalized samples in
/// place. [`flush`](Self::flush) is called once at stop to recover any tail the
/// processor is holding back (e.g. a fixed delay line); it returns the trailing
/// 48kHz / stereo interleaved samples (empty if none).
pub trait InnerProcessor: Send {
    /// Process one batch of 48kHz / stereo interleaved samples in place. The
    /// slice length is always a multiple of two (stereo).
    fn process(&mut self, samples: &mut [f32]);
    /// Flush and return any buffered tail (48kHz / stereo interleaved). Called
    /// once when the stream stops. Returns an empty vector when nothing is held.
    fn flush(&mut self) -> Vec<f32>;
}

/// Stateful two-stage converter that normalizes input device frames to the internal canonical
/// form (48k/stereo/960 frames), then converts them to one or more output taps (primary and
/// optional secondary).
///
/// Stage 1 (internal normalization) runs once and feeds the resulting canonical form to stage 2
/// of both taps. An optional [`InnerProcessor`] is applied once to the canonical form before it
/// is split for stage 2, so both taps receive the same processed audio.
///
/// Accumulate interleaved samples with `push`, then retrieve completed output chunks one at a
/// time with `pop_chunk` (primary) and `pop_secondary` (secondary).
pub struct Normalizer {
    flush_result: Option<Result<()>>,
    in_sample_rate: u32,
    in_channels: usize,
    /// Interleaved input awaiting a complete native frame.
    pending_input: Vec<f32>,
    /// Timestamp of the first sample retained in `pending_input`.
    pending_input_pts: Option<i64>,

    // --- Stage 1 (internal normalization: → 48k/stereo, shared by all output taps) ---
    /// Passthrough when the input is 48000 Hz (no resampler).
    stage1_resampler: Option<ResamplerState>,
    /// Temporary buffer for the internal canonical form (48k/stereo interleaved) generated by
    /// this push. Apply [`InnerProcessor`] before distributing it to each output tap. Reuse capacity.
    inner_scratch: Vec<f32>,
    /// Total internal 48k frames generated by stage 1 so far (for PTS anchor calculation).
    total_inner_frames: u64,

    /// Optional processor (e.g. noise suppression) applied once to the internal canonical form
    /// before the stage 2 split.
    inner_processor: Option<Box<dyn InnerProcessor>>,

    // --- Output taps (each has its own stage 2, output buffer, and PTS state) ---
    /// Primary output tap (its [`OutputFormat`] is `output`).
    primary: OutputTap,
    capture: Option<OutputTap>,
    /// Secondary output tap (when configured). Has stage 2 and PTS state independent of primary.
    secondary: Option<OutputTap>,
}

/// One output tap. Receives the shared internal canonical form (48k/stereo), converts channels
/// and sample rate to its [`OutputFormat`], then emits fixed 20ms chunks. Each tap has its own
/// PTS anchor and output buffer (which can cause tens of milliseconds of PTS offset between taps).
struct OutputTap {
    output: OutputFormat,
    /// Output stage. `None` (full passthrough) when `output == {48000, 2}`, matching the canonical form.
    stage2: Option<OutputStage>,
    /// Output awaiting completion (interleaved with `output.channels`). `pop` drains from here.
    out_buf: Vec<f32>,
    /// Frames in one output chunk (`output.chunk_frames()`).
    out_chunk_frames: usize,
    /// Output channel count.
    out_channels: usize,
    /// Output frame index for the first sample in `out_buf` (the oldest sample not yet popped).
    out_frame_origin: u64,
    /// PTS anchor: associates an output frame index with `device_pts` (ns).
    pts_anchor: Option<PtsAnchor>,
    /// Exact integer projection for capture anchors; output taps retain their legacy projection.
    canonical_clock: bool,
    padding: std::collections::VecDeque<std::ops::Range<u64>>,
}

#[derive(Clone, Copy)]
struct PtsAnchor {
    /// Output frame index, based on the output sample rate.
    out_frame: u64,
    /// `device_pts` (ns) corresponding to that output frame.
    pts_ns: i64,
}

/// Sample-rate converter wrapping rubato `Async` (`FixedAsync::Input`) for one stage.
///
/// Calls `process` for each fixed input chunk of `chunk_in_frames` frames and appends variable-length
/// output to `out_buf` (the caller's accumulator). `channels` varies by stage (always stereo=2 in
/// stage 1; the output channel count in stage 2).
struct ResamplerState {
    inner: Async<f32>,
    channels: usize,
    /// Input frames required per rubato call (fixed by `FixedAsync::Input`).
    chunk_in_frames: usize,
    /// Maximum output frames one `process` call can produce.
    max_out_frames: usize,
    /// Unprocessed input (interleaved, `channels` channels).
    in_accum: Vec<f32>,
    /// Output scratch buffer for rubato (reused to avoid allocations).
    out_scratch: Vec<f32>,
}

/// Stage 2 (output). Receives 960-frame chunks in the internal 48k/stereo canonical form and
/// converts channels and sample rate to produce interleaved output in the requested format.
struct OutputStage {
    out_channels: usize,
    /// Resampler from 48000 to `output.sample_rate`. `None` (sample-rate passthrough) when
    /// `output.sample_rate == 48000`. Applied to samples after channel conversion.
    resampler: Option<ResamplerState>,
    /// Scratch buffer after channel conversion and before sample-rate conversion (48k / `out_channels` interleaved).
    ch_scratch: Vec<f32>,
}

impl Normalizer {
    /// Create a normalizer for the input sample rate, input channel count, and output format.
    ///
    /// Stage 1 normalizes input to 48k/stereo (`in_sample_rate == 48000` bypasses sample-rate
    /// conversion; `in_channels` of 1 duplicates mono to stereo and 2 is unchanged.
    /// Input above two channels is rejected with UnsupportedFormat). Stage 2 converts the canonical form to `output` (passthrough
    /// when `output == {48000, 2}`).
    ///
    /// Expects the caller to have run [`OutputFormat::validate`] on `output` (this method does not
    /// clamp it to a valid range).
    ///
    /// rubato resampler construction can fail, for example with extreme rate ratios. Panicking
    /// would silently stop the non-RT capture thread, so failures return [`Error::Backend`] and
    /// propagate to the caller.
    pub fn new(in_sample_rate: u32, in_channels: u16, output: OutputFormat) -> Result<Self> {
        if in_sample_rate == 0 || in_channels == 0 {
            return Err(Error::InvalidArg(
                "native rate and channels must be positive".into(),
            ));
        }
        if in_channels > 2 {
            return Err(Error::UnsupportedFormat(
                "native input supports only mono or stereo".into(),
            ));
        }
        output.validate()?;
        let in_channels = usize::from(in_channels);

        // Stage 1 resampler (→48000). Runs once for all output taps.
        let stage1_resampler = if in_sample_rate == SAMPLE_RATE {
            None
        } else {
            Some(ResamplerState::new(in_sample_rate, SAMPLE_RATE, INNER_CH)?)
        };

        Ok(Self {
            flush_result: None,
            in_sample_rate,
            in_channels,
            pending_input: Vec::new(),
            pending_input_pts: None,
            stage1_resampler,
            inner_scratch: Vec::with_capacity(CHUNK_FRAMES * INNER_CH * 4),
            total_inner_frames: 0,
            inner_processor: None,
            primary: OutputTap::new(output)?,
            capture: None,
            secondary: None,
        })
    }

    /// Add a secondary output tap (expects [`OutputFormat::validate`] to have been called).
    ///
    /// The internal canonical form (48k/stereo) is generated once and fed to stage 2 of both taps.
    /// The secondary tap has independent stage 2 and PTS state, and emits its own fixed 20ms chunks.
    /// rubato construction failures return [`Error::Backend`].
    pub fn with_secondary(mut self, secondary: OutputFormat) -> Result<Self> {
        self.secondary = Some(OutputTap::new(secondary)?);
        Ok(self)
    }

    /// Enable a pre-output-conversion canonical branch for off-callback consumers.
    /// Its stop tail contains valid frames only, never transport padding.
    pub fn with_capture_tap(mut self) -> Result<Self> {
        let mut capture = OutputTap::new(OutputFormat::default())?;
        capture.canonical_clock = true;
        self.capture = Some(capture);
        Ok(self)
    }

    /// Canonical frames produced but not emitted, discarded when this generation is rebuilt.
    pub fn buffered_capture_frames(&self) -> usize {
        self.capture
            .as_ref()
            .map_or(0, OutputTap::buffered_out_frames)
    }

    /// Retrieve valid canonical stereo frames; allow a short final tail only at stop.
    pub fn pop_capture(&mut self, stopping: bool) -> Option<(Vec<f32>, i64)> {
        self.pop_capture_with_metadata(stopping)
            .map(|chunk| (chunk.samples, chunk.pts_ns))
    }

    /// Retrieve valid canonical frames with no transport padding.
    pub fn pop_capture_with_metadata(&mut self, stopping: bool) -> Option<NormalizedChunk> {
        let tap = self.capture.as_mut()?;
        let frames = tap.buffered_out_frames();
        if frames == 0 || (frames < CHUNK_FRAMES && !stopping) {
            return None;
        }
        let take = frames.min(CHUNK_FRAMES);
        let pts = tap.pts_for_out_frame(tap.out_frame_origin);
        let data = tap.out_buf.drain(..take * INNER_CH).collect();
        tap.out_frame_origin += take as u64;
        Some(NormalizedChunk {
            samples: data,
            pts_ns: pts,
            flags: ChunkFlags::empty(),
        })
    }

    /// Inject a processor to apply once to the internal canonical form before the stage 2 split.
    pub fn with_inner_processor(mut self, processor: Box<dyn InnerProcessor>) -> Self {
        self.inner_processor = Some(processor);
        self
    }

    /// Input sample rate (Hz).
    pub fn in_sample_rate(&self) -> u32 {
        self.in_sample_rate
    }

    /// Primary output format.
    pub fn output(&self) -> OutputFormat {
        self.primary.output
    }

    /// Secondary output format (only when a secondary tap is configured).
    pub fn secondary_output(&self) -> Option<OutputFormat> {
        self.secondary.as_ref().map(|t| t.output)
    }

    /// Whether a secondary tap is enabled.
    pub fn has_secondary(&self) -> bool {
        self.secondary.is_some()
    }

    /// Whether stage 1 sample-rate conversion is bypassed (`in == 48000`).
    pub fn is_passthrough(&self) -> bool {
        self.stage1_resampler.is_none()
    }

    /// Whether stage 2 for the primary output is fully bypassed (`output == {48000, 2}`).
    pub fn is_output_passthrough(&self) -> bool {
        self.primary.stage2.is_none()
    }

    /// Accumulate interleaved input samples.
    ///
    /// Incomplete trailing frames are retained until a later push completes them.
    /// `device_pts_ns` is the device-provided PTS for the first frame in this push.
    ///
    /// Returns [`Error::Backend`] if rubato's `process` fails. This avoids silently stopping the
    /// capture thread with a panic and lets the caller stop the stream explicitly.
    pub fn push(&mut self, interleaved: &[f32], device_pts_ns: i64) -> Result<()> {
        if interleaved.is_empty() {
            return Ok(());
        }
        if let Some(Err(error)) = &self.flush_result {
            return Err(error.clone());
        }
        self.flush_result = None;
        if self.pending_input.is_empty() {
            self.pending_input_pts = Some(device_pts_ns);
        }
        self.pending_input.extend_from_slice(interleaved);
        let in_frames = self.pending_input.len() / self.in_channels;
        if in_frames == 0 {
            return Ok(());
        }
        let device_pts_ns = self.pending_input_pts.unwrap_or(device_pts_ns);

        // Approximate the future output frame position for the start of this push by the sample-rate
        // ratio and update each tap's PTS anchor. This is approximate because the resampler retains
        // a remainder internally. Set anchors independently for primary and secondary.
        self.primary
            .update_pts_anchor(self.total_inner_frames, device_pts_ns)?;
        if let Some(sec) = self.secondary.as_mut() {
            sec.update_pts_anchor(self.total_inner_frames, device_pts_ns)?;
        }

        if let Some(capture) = self.capture.as_mut() {
            // One rational PTS anchor per source generation. Arrival-time jitter must
            // not reanchor old samples or create a VAD epoch on every push.
            if capture.pts_anchor.is_none() {
                capture.update_pts_anchor(self.total_inner_frames, device_pts_ns)?;
            }
        }

        // Stage 1: mix channels → stereo interleaved → normalize to 48k. Collect what this push
        // generates in inner_scratch.
        self.inner_scratch.clear();
        let mut stereo = Vec::with_capacity(in_frames * INNER_CH);
        Self::mix_to_stereo(
            &self.pending_input,
            self.in_channels,
            in_frames,
            &mut stereo,
        );
        self.pending_input.drain(..in_frames * self.in_channels);
        self.pending_input_pts = if self.pending_input.is_empty() {
            None
        } else {
            Some(device_pts_ns.saturating_add(
                (in_frames as i64).saturating_mul(1_000_000_000) / self.in_sample_rate as i64,
            ))
        };
        match &mut self.stage1_resampler {
            None => {
                // Sample-rate passthrough. Use directly as the internal canonical form.
                self.total_inner_frames += in_frames as u64;
                self.inner_scratch.extend_from_slice(&stereo);
            }
            Some(rs) => {
                rs.in_accum.extend_from_slice(&stereo);
                let produced = rs.drain_into(&mut self.inner_scratch)?;
                self.total_inner_frames += produced;
            }
        }

        // Apply a processor (e.g. denoise) once to the internal canonical form before the stage 2 split.
        if let Some(proc) = self.inner_processor.as_mut() {
            proc.process(&mut self.inner_scratch);
        }

        // Distribute the processed canonical form to stage 2 of both taps.
        self.distribute_inner()
    }

    /// Retrieve one completed primary output chunk.
    ///
    /// Returns (`out_chunk_frames` interleaved frames at `output.channels`, PTS in ns for the
    /// first sample). Returns `None` if a full chunk is not buffered.
    pub fn pop_chunk(&mut self) -> Option<(Vec<f32>, i64)> {
        self.pop_chunk_with_metadata()
            .map(|chunk| (chunk.samples, chunk.pts_ns))
    }

    /// Retrieve one completed secondary output chunk (`None` if no secondary tap is configured).
    pub fn pop_secondary(&mut self) -> Option<(Vec<f32>, i64)> {
        self.pop_secondary_with_metadata()
            .map(|chunk| (chunk.samples, chunk.pts_ns))
    }

    /// Retrieve primary PCM with its integrity flags.
    pub fn pop_chunk_with_metadata(&mut self) -> Option<NormalizedChunk> {
        self.primary.pop()
    }

    /// Retrieve secondary PCM with its integrity flags.
    pub fn pop_secondary_with_metadata(&mut self) -> Option<NormalizedChunk> {
        self.secondary.as_mut().and_then(OutputTap::pop)
    }

    /// Drain stage 1, processor tail and all output stages, retaining independent failures.
    /// Repeated flush returns the same result without generating duplicate PCM.
    /// Valid output remains drainable on cleanup-only failure.
    pub fn flush(&mut self) -> Result<()> {
        if let Some(result) = &self.flush_result {
            return result.clone();
        }
        let mut errors = Vec::new();
        if !self.pending_input.is_empty() {
            errors.push(Error::Backend("incomplete native frame at flush".into()));
            self.pending_input.clear();
            self.pending_input_pts = None;
        }
        self.inner_scratch.clear();
        if let Some(stage) = self.stage1_resampler.as_mut() {
            match stage.flush_into(&mut self.inner_scratch) {
                Ok(produced) => self.total_inner_frames += produced,
                Err(error) => errors.push(error),
            }
            if let Some(processor) = self.inner_processor.as_mut() {
                processor.process(&mut self.inner_scratch);
            }
            if let Err(error) = self.distribute_inner() {
                errors.push(error);
            }
        }
        if let Some(processor) = self.inner_processor.as_mut() {
            self.inner_scratch = processor.flush();
            // Keep the canonical capture epoch empty for a primed silent processor.
            let empty_capture = if self.total_inner_frames == 0 {
                self.capture.take()
            } else {
                None
            };
            if let Err(error) = self.distribute_inner() {
                errors.push(error);
            }
            if let Some(capture) = empty_capture {
                self.capture = Some(capture);
            }
        }
        if let Err(error) = self.primary.flush() {
            errors.push(error);
        }
        if let Some(secondary) = self.secondary.as_mut() {
            if let Err(error) = secondary.flush() {
                errors.push(error);
            }
        }
        if let Some(capture) = self.capture.as_mut() {
            // Preserve canonical valid-frame tails without transport padding.
            if let Some(stage) = capture.stage2.as_mut() {
                if let Err(error) = stage.flush_into(&mut capture.out_buf) {
                    errors.push(error);
                }
            }
        }
        let result = combine_failures(errors)
            .map_err(|error| error.with_context(ErrorContext::new(Operation::Flush)));
        self.flush_result = Some(result.clone());
        result
    }

    /// Number of output frames buffered but not yet retrieved from `out_buf` (primary tap).
    pub fn buffered_out_frames(&self) -> usize {
        self.primary.buffered_out_frames()
    }

    // --- Internal helpers ---

    /// Distribute the canonical form in `inner_scratch` to stage 2 of both taps. Temporarily take
    /// the buffer to avoid borrow conflicts, then restore its capacity.
    fn distribute_inner(&mut self) -> Result<()> {
        if self.inner_scratch.is_empty() {
            return Ok(());
        }
        let inner = std::mem::take(&mut self.inner_scratch);
        let r_primary = self.primary.feed_inner(&inner);
        let r_capture = self
            .capture
            .as_mut()
            .map(|capture| capture.feed_inner(&inner))
            .unwrap_or(Ok(()));
        let r_secondary = self
            .secondary
            .as_mut()
            .map(|sec| sec.feed_inner(&inner))
            .unwrap_or(Ok(()));
        // Restore the buffer so its capacity can be reused.
        self.inner_scratch = inner;
        self.inner_scratch.clear();
        combine_failures(
            [r_primary, r_capture, r_secondary]
                .into_iter()
                .filter_map(Result::err)
                .collect(),
        )
    }

    /// Convert validated mono/stereo input to stereo interleaved and push it to `dst`.
    fn mix_to_stereo(src: &[f32], in_ch: usize, in_frames: usize, dst: &mut Vec<f32>) {
        match in_ch {
            1 => {
                // mono → stereo (duplicate L=R)
                for &s in &src[..in_frames] {
                    dst.push(s);
                    dst.push(s);
                }
            }
            2 => {
                // Keep 2 channels unchanged (only the required samples).
                dst.extend_from_slice(&src[..in_frames * 2]);
            }
            _ => unreachable!("native channel count validated by constructor"),
        }
    }
}

#[cfg(test)]
mod tests;

#[cfg(test)]
mod repro_tests;

mod output;
mod resampler;

fn combine_failures(errors: Vec<Error>) -> Result<()> {
    let mut errors = errors.into_iter();
    match (errors.next(), errors.next()) {
        (None, _) => Ok(()),
        (Some(error), None) => Err(error),
        (Some(primary), Some(first)) => Err(Error::Multiple(ErrorGroup::new(
            primary,
            first,
            errors.collect(),
        ))),
    }
}