coremlit 0.1.2

//! The CLAP [`AudioEncoder`]: a Rust log-mel front-end (the private `mel`
//! submodule) around the fp16 CoreML HTSAT graph, with L2 normalization applied
//! in Rust.

mod mel;

use std::path::Path;

use crate::{
  ComputeUnits, DataType, Model, MultiArray,
  model::contract::{Checked, Dim, FeatureContract, LoadContract, StateContract},
};

use crate::embeddings::clap::{
  embedding::{EMBEDDING_DIM, Embedding, check_finite_output},
  error::{AudioTooLong, Error, OutputShape, Result, WinditError, contract_violation},
  window::{WindowEmbedding, WindowPlan},
};

pub use self::mel::{N_MELS, SAMPLE_RATE_HZ, T_FRAMES, TARGET_SAMPLES};

/// Declared feature names on `clap_audio.mlmodelc` (pinned by
/// `tests/clap/model_io.rs`).
mod names {
  pub const INPUT_FEATURES: &str = "input_features";
  pub const AUDIO_EMBEDS: &str = "audio_embeds";
}

/// Default [`AudioEncoderOptions::compute`]. [`ComputeUnits::All`] lets CoreML
/// schedule across the available hardware.
///
/// As converted (T1), the HTSAT audio graph does **not** compile for the ANE and
/// CoreML falls back to GPU/CPU (fp16-clean there); the text graph does compile
/// for the ANE. `All` is the honest default — it never *asserts* ANE residency,
/// which `tests/clap/placement.rs` characterizes rather than claims.
///
/// The issue #30 perf pass (the `clap_encode` bench) measured every placement and
/// **kept** `All` here — unlike the text tower, no non-`All` unit clears the
/// meaningful-speedup bar. Warm-median latency (Apple M1 Max, macOS 26.5 25F71,
/// fp16, 25 runs × 3 sweeps): `CpuAndGpu` ~44.2 ms, `All` ~48.6 ms, `CpuOnly`
/// ~47.8 ms, `CpuAndNeuralEngine` ~120.6 ms. `CpuAndGpu` is only ~9 % faster than
/// `All` (below the ~15 % threshold a default change would need), `All`/`CpuOnly`
/// are within noise, and naming the ANE is ~2.5× *slower* (the HTSAT `ANECCompile`
/// fallback thrash). So `All` stays the default.
///
/// One honest caveat the bench also surfaced: because `All` names the ANE, every
/// *load* re-attempts the HTSAT `ANECCompile` (which fails and falls back), a
/// ~1.5–2 s cost `CpuAndGpu` skips entirely (it loads in ~0.1 s). Since the crate
/// mandates construct-once / reuse / prewarm, that load is paid once and does not
/// touch per-request latency — but a load-latency-sensitive caller can pin
/// [`AudioEncoderOptions::with_compute`]`(ComputeUnits::CpuAndGpu)` for the faster
/// load (and ~9 % warm win) at identical parity.
pub const DEFAULT_AUDIO_COMPUTE: ComputeUnits = ComputeUnits::All;

#[cfg(feature = "serde")]
fn default_audio_compute() -> ComputeUnits {
  DEFAULT_AUDIO_COMPUTE
}

/// Construction options for [`AudioEncoder`] (rust-options-pattern): a single
/// `compute` knob with one source of truth shared by `const new`/`Default`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct AudioEncoderOptions {
  #[cfg_attr(feature = "serde", serde(default = "default_audio_compute"))]
  compute: ComputeUnits,
}

impl Default for AudioEncoderOptions {
  fn default() -> Self {
    Self::new()
  }
}

impl AudioEncoderOptions {
  /// Options matching the crate default: [`DEFAULT_AUDIO_COMPUTE`].
  pub const fn new() -> Self {
    Self {
      compute: DEFAULT_AUDIO_COMPUTE,
    }
  }

  /// Which hardware CoreML may schedule the audio graph on.
  #[inline]
  pub const fn compute(&self) -> ComputeUnits {
    self.compute
  }

  /// Builder form of [`Self::set_compute`].
  #[must_use]
  #[inline]
  pub const fn with_compute(mut self, compute: ComputeUnits) -> Self {
    self.set_compute(compute);
    self
  }

  /// Sets [`Self::compute`] in place.
  #[inline]
  pub const fn set_compute(&mut self, compute: ComputeUnits) -> &mut Self {
    self.compute = compute;
    self
  }
}

/// CLAP audio encoder: 48 kHz mono `&[f32]` in, a unit-norm 512-d
/// [`Embedding`] out.
///
/// The front-end is a Rust log-mel port (the private `mel` submodule); the fp16
/// CoreML HTSAT graph maps the `[1, 1, 1001, 64]` spectrogram to a
/// pre-normalization 512-d projection, which this encoder L2-normalizes.
///
/// `&self` inference (no mutable scratch): the FFT plan and filterbank are built
/// once at load and per-call buffers are local, so fan-out means one
/// [`AudioEncoder`] per worker over a `Send` [`crate::Model`]
/// (`crate::Model` is deliberately `!Sync`).
#[derive(Debug)]
pub struct AudioEncoder {
  /// A [`Checked`], never a bare [`Model`]: [`audio_contract`] is the only
  /// contract this door states and [`Checked::new`] is the only way one is
  /// built, so removing the check from [`Self::from_file_with`] does not
  /// compile.
  model: Checked,
  mel: mel::MelExtractor,
}

impl AudioEncoder {
  /// Loads `clap_audio.mlmodelc` with [`AudioEncoderOptions::new`]
  /// ([`ComputeUnits::All`]).
  ///
  /// # Errors
  /// As [`Self::from_file_with`].
  pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
    Self::from_file_with(path, AudioEncoderOptions::new())
  }

  /// Loads the model with custom options.
  ///
  /// The model is checked against this door's load contract
  /// (`audio_contract`) and held as a crate-internal `Checked` wrapper whose
  /// only constructor runs that check:
  ///
  /// ```text
  /// input   input_features  f32  [1, 1, 1001, 64]  every axis Exactly
  /// output  audio_embeds    f32  [1, 512]          every axis Exactly
  /// state   none
  /// ```
  ///
  /// The contract is COMPLETE over the members of [`crate::ModelDescription`]
  /// that can make a conformant prediction fail, not just over the features
  /// this door sends: a graph carrying `input_features` plus another REQUIRED
  /// input clears every per-feature clause and then fails on every prediction,
  /// and a STATE buffer is not an input at all — it lives in its own
  /// dictionary, so a stateful ML Program declaring exactly this door's two
  /// features plus a state clears the input set too, and then meets
  /// [`Self::embed_window`], which predicts through the stateless API. The
  /// ground truth stays pinned by `tests/clap/model_io.rs`, which now also
  /// loads both staged tiers THROUGH this constructor.
  ///
  /// # Errors
  /// [`Error::Load`] if CoreML rejects the model;
  /// [`Error::ContractMismatch`] if a named feature's type or geometry
  /// mismatches; [`Error::UnsatisfiableInput`] if it requires an input this
  /// door never sends; [`Error::UnsatisfiableState`] if it declares a state
  /// buffer.
  pub fn from_file_with(path: impl AsRef<Path>, options: AudioEncoderOptions) -> Result<Self> {
    let model = Model::load(path, options.compute())?;
    let model = Checked::new(model, &audio_contract()).map_err(contract_violation)?;

    Ok(Self {
      model,
      mel: mel::MelExtractor::new(),
    })
  }

  /// Embeds one audio window into a unit-norm [`Embedding`].
  ///
  /// `samples` is 48 kHz ([`SAMPLE_RATE_HZ`]) mono and must be
  /// `1..=`[`TARGET_SAMPLES`] long. A shorter clip is `repeatpad`ed up to the
  /// fixed 480 000-sample window (exactly as HF's `ClapFeatureExtractor` does);
  /// a **longer** clip is rejected with [`Error::AudioTooLong`] rather than
  /// silently head-truncated. This is the per-window primitive — feed a longer
  /// clip to [`Self::embed_windows`] (the long-audio pipeline), which hops it
  /// into 480 000-sample windows first. (HF is configured for `rand_trunc`, so
  /// head-truncating here would be neither deterministic nor HF-faithful;
  /// clapkit will not truncate behind your back.)
  ///
  /// # Errors
  /// [`Error::EmptyAudio`] if `samples` is empty.
  /// [`Error::AudioTooLong`] if `samples.len()` exceeds [`TARGET_SAMPLES`] (use
  /// [`Self::embed_windows`] for long audio).
  /// [`Error::NonFiniteInput`] if any sample is NaN/infinite (it would
  /// otherwise propagate through the mel into a garbage embedding).
  /// [`Error::Tensor`] / [`Error::Prediction`] on a tensor or CoreML failure.
  /// [`Error::OutputShape`] if the predicted `audio_embeds` shape diverges from
  /// `[1, `[`EMBEDDING_DIM`]`]`. [`Error::NonFiniteOutput`] if the model output
  /// has a NaN/infinite component — model corruption, classified apart from a
  /// caller's own non-finite embedding data ([`Error::NonFiniteEmbedding`]).
  /// [`Error::EmbeddingZero`] if the (finite) projection has zero magnitude.
  pub fn embed_window(&self, samples: &[f32]) -> Result<Embedding> {
    if samples.is_empty() {
      return Err(Error::EmptyAudio);
    }
    check_window_len(samples.len())?;
    if let Some(index) = first_non_finite(samples) {
      return Err(Error::NonFiniteInput(index));
    }

    let mut features = vec![0.0f32; N_MELS * T_FRAMES];
    self.mel.extract_into(samples, &mut features)?;

    // Time-major mel [1001, 64] maps directly onto the row-major
    // `input_features [1, 1, 1001, 64]` contract (T1).
    let input = MultiArray::from_slice(&[1, 1, T_FRAMES, N_MELS], &features)?;
    let mut outputs = self
      .model
      .predict_with(&[(names::INPUT_FEATURES, &input)])?;
    let embeds = outputs
      .take(names::AUDIO_EMBEDS)
      .ok_or_else(|| crate::PredictionError::MissingOutput(names::AUDIO_EMBEDS.to_string()))?;
    if embeds.shape() != [1, EMBEDDING_DIM] {
      return Err(Error::OutputShape(OutputShape::new(
        embeds.shape().to_vec(),
        vec![1, EMBEDDING_DIM],
      )));
    }

    let mut row = [0.0f32; EMBEDDING_DIM];
    embeds.copy_into::<f32>(&mut row)?;
    // Classify a NaN/∞ the CoreML runtime produced as model-output corruption
    // (`NonFiniteOutput`) before it reaches `from_slice_normalizing`, which would
    // otherwise mislabel it as caller-supplied embedding data
    // (`NonFiniteEmbedding`).
    check_finite_output(&row)?;
    Embedding::from_slice_normalizing(&row)
  }

  /// Embeds a long clip as overlapped windows per `plan`, one
  /// [`WindowEmbedding`] (embedding + its span + coverage) per
  /// [`Span`](crate::embeddings::clap::window::Span).
  ///
  /// This is the long-audio pipeline entry: it slices `samples` at the plan's
  /// offsets and runs [`Self::embed_window`] on each (which `repeatpad`s a short
  /// tail to the fixed window). The per-window embeddings are RETURNED, not
  /// hidden inside aggregation, so a caller can aggregate them with an
  /// [`AggregatePolicy`](crate::embeddings::clap::aggregate::AggregatePolicy), score each window
  /// ([`crate::embeddings::clap::score::score_windows`]), or both.
  ///
  /// Runs sequentially: [`crate::Model`] is `!Sync`, so windows share one
  /// encoder on one thread (fan out with one [`AudioEncoder`] per worker for
  /// parallelism).
  ///
  /// # Errors
  /// [`Error::EmptyAudio`] if `samples` is empty;
  /// [`Error::Windowing`]`(`[`WinditError::TooManyWindows`]`)`
  /// if `plan` would plan more than its
  /// [`max_windows`](crate::embeddings::clap::WindowPlan::max_windows) over this clip (the
  /// resource rail — refused before any window is embedded), or
  /// [`Error::Windowing`]`(`[`WinditError::AllocFailed`]`)`
  /// if an admitted (at-cap) plan's output vector cannot be reserved — done
  /// fallibly, so an allocator refusal is typed rather than a process abort;
  /// otherwise any error [`Self::embed_window`] raises for a window.
  pub fn embed_windows(&self, samples: &[f32], plan: &WindowPlan) -> Result<Vec<WindowEmbedding>> {
    if samples.is_empty() {
      return Err(Error::EmptyAudio);
    }
    let spans = plan.spans(samples.len())?;
    // The at-cap plan can name up to `plan.max_windows` spans, each an inline
    // ~2 KiB `WindowEmbedding` — hundreds of MiB for a small clip. Reserve it
    // FALLIBLY (mirroring `spans()` itself), so an allocator refusal surfaces as
    // the typed `AllocFailed` rather than aborting the process the way the prior
    // infallible `with_capacity` would. The reservation is exact, so the in-loop
    // `push`es stay within capacity and never reallocate.
    let mut out = Vec::new();
    out.try_reserve_exact(spans.len()).map_err(|_| {
      Error::Windowing(WinditError::AllocFailed {
        elements: spans.len(),
      })
    })?;
    for span in spans {
      let embedding = self.embed_window(&samples[span.start()..span.end()])?;
      out.push(WindowEmbedding::new(embedding, span));
    }
    Ok(out)
  }

  /// Runs one throwaway [`Self::embed_window`] to fully specialize the prediction
  /// path, so the first user-facing request is warm.
  ///
  /// Construction ([`Self::from_file`] &c.) already pays the model *load* /
  /// device specialization; what it does **not** pay is the first prediction's
  /// own graph specialization. The `clap_encode` bench measures that first
  /// inference at several times the warm latency (e.g. ~200 ms first vs ~48 ms
  /// warm on `All`), so calling `prewarm` once — after construction, before
  /// serving — moves that one-time cost off the first real clip. Then **reuse**
  /// this same encoder for every request (it is `&self`, so it stays resident and
  /// there is nothing to reconstruct).
  ///
  /// This is the whole prewarm delta over the construct-once-and-reuse pattern:
  /// the load is the constructor's job, and this is deliberately *only* the dummy
  /// inference the reuse pattern otherwise leaves for the first live request. The
  /// warm-up runs a fixed synthetic 1 s tone (`repeatpad`ed to the fixed window),
  /// so it neither reads caller audio nor allocates a full-window buffer up front.
  ///
  /// # Errors
  /// As [`Self::embed_window`]; a failure here surfaces a broken model at prewarm
  /// time rather than on the first request.
  pub fn prewarm(&self) -> Result<()> {
    // 1 s of a fixed 440 Hz tone: a valid, non-silent window (avoids the
    // zero-magnitude path) that `embed_window` `repeatpad`s to the full window.
    let signal: Vec<f32> = (0..SAMPLE_RATE_HZ)
      .map(|i| 0.5 * (std::f32::consts::TAU * 440.0 * (i as f32 / SAMPLE_RATE_HZ as f32)).sin())
      .collect();
    self.embed_window(&signal)?;
    Ok(())
  }
}

/// The load contract this door states: `input_features` `[1, 1, 1001, 64]` f32
/// in, `audio_embeds` `[1, 512]` f32 out, no state.
///
/// Data rather than a sequence of checks, and the ONLY thing
/// [`AudioEncoder::from_file_with`] does beyond calling [`Model::load`]. The
/// four hand-written comparisons this replaced — a presence test and a
/// shape-and-dtype test per feature — were each a check the constructor could
/// forget to make, and deleting any of them failed no runnable test. A
/// [`Checked`] field turns that mutation into a compile error.
///
/// Every axis is [`Dim::Exactly`], and that buys more than the numbers.
/// [`crate::FeatureInfo::shape`] reports the DEFAULT shape of a flexible input,
/// so a `RangeDims` graph converted at `[1, 1, 1001, 64]` declares this
/// contract's exact numbers. An all-`Exactly` contract therefore requires the
/// whole feature to be [`crate::ShapeConstraint::Fixed`], which is the only
/// thing that separates the two. Nothing here is read back off the artifact:
/// the fp16 and int8 tiers are contract-identical, so every number is this
/// door's own.
fn audio_contract() -> LoadContract {
  LoadContract::new(
    vec![FeatureContract::new(
      names::INPUT_FEATURES,
      DataType::F32,
      vec![
        Dim::Exactly(1),
        Dim::Exactly(1),
        Dim::Exactly(T_FRAMES),
        Dim::Exactly(N_MELS),
      ],
    )],
    vec![FeatureContract::new(
      names::AUDIO_EMBEDS,
      DataType::F32,
      vec![Dim::Exactly(1), Dim::Exactly(EMBEDDING_DIM)],
    )],
    StateContract::None,
  )
}

/// Flat index of the first non-finite (NaN/±∞) sample, if any.
fn first_non_finite(samples: &[f32]) -> Option<usize> {
  samples.iter().position(|v| !v.is_finite())
}

/// Reject a per-window sample count over [`TARGET_SAMPLES`]. The mel front-end
/// maps exactly one 480 000-sample window per inference, so an over-length clip
/// is a caller error ([`Error::AudioTooLong`]) — it must be hopped into windows
/// by [`AudioEncoder::embed_windows`], never silently head-truncated. Empty input
/// is rejected separately by the caller ([`Error::EmptyAudio`]).
fn check_window_len(len: usize) -> Result<()> {
  if len > TARGET_SAMPLES {
    return Err(Error::AudioTooLong(AudioTooLong::new(len, TARGET_SAMPLES)));
  }
  Ok(())
}

#[cfg(test)]
mod tests;