token-sampling 0.1.0

use core::cmp::Ordering;

use crate::candidates::{candidate_order, Candidates, Error};
use crate::rng::{unit, Rng};

/// Randomness and token history supplied to one sampling stage.
pub struct SamplerCtx<'a> {
    /// Random source for stages that draw.
    pub rng: &'a mut dyn Rng,
    /// Token history in chronological order.
    pub history: &'a [u32],
}

impl core::fmt::Debug for SamplerCtx<'_> {
    fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        formatter
            .debug_struct("SamplerCtx")
            .field("rng", &"dyn Rng")
            .field("history", &self.history)
            .finish()
    }
}

/// One object-safe stage in a sampling chain.
pub trait Sampler: Send {
    /// Applies the stage to the current candidates.
    fn apply(&mut self, candidates: &mut Candidates, ctx: &mut SamplerCtx<'_>)
        -> Result<(), Error>;

    /// Restores per-sequence state.
    fn reset(&mut self) {}
}

/// Divides logits by a fixed temperature, with greedy masking at zero.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Temperature {
    /// Temperature value.
    pub t: f32,
}

impl Default for Temperature {
    fn default() -> Self {
        Self { t: 0.8 }
    }
}

impl Sampler for Temperature {
    fn apply(&mut self, candidates: &mut Candidates, _: &mut SamplerCtx<'_>) -> Result<(), Error> {
        apply_temperature(candidates, self.t);
        Ok(())
    }
}

/// Maps distribution entropy to a temperature range.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct DynTemperature {
    /// Center temperature.
    pub t: f32,
    /// Half-width of the temperature range.
    pub delta: f32,
    /// Exponent applied to normalized entropy.
    pub exponent: f32,
}

impl Default for DynTemperature {
    fn default() -> Self {
        Self {
            t: 0.8,
            delta: 0.0,
            exponent: 1.0,
        }
    }
}

impl Sampler for DynTemperature {
    fn apply(&mut self, candidates: &mut Candidates, _: &mut SamplerCtx<'_>) -> Result<(), Error> {
        if !self.delta.is_finite() || self.delta <= 0.0 {
            apply_temperature(candidates, self.t);
            return Ok(());
        }
        if candidates.buf.len() <= 1 {
            return Ok(());
        }

        let min_t = (self.t as f64 - self.delta as f64).max(0.0);
        let max_t = self.t as f64 + self.delta as f64;
        candidates.softmax(true)?;
        let entropy = candidates
            .buf
            .iter()
            .filter(|candidate| candidate.p > 0.0)
            .map(|candidate| {
                let p = candidate.p as f64;
                -p * p.ln()
            })
            .sum::<f64>();
        let max_entropy = (candidates.buf.len() as f64).ln();
        let exponent = if self.exponent.is_finite() {
            self.exponent as f64
        } else {
            1.0
        };
        let normalized = (entropy / max_entropy).powf(exponent);
        let dynamic = min_t + (max_t - min_t) * normalized;
        let dynamic = if dynamic.is_nan() {
            self.t
        } else {
            finite_f32(dynamic)
        };
        apply_temperature(candidates, dynamic);
        candidates.softmax(false)
    }
}

/// Keeps the `k` candidates with the largest logits.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TopK {
    /// Maximum survivor count. Non-positive values disable the stage.
    pub k: i32,
}

impl Default for TopK {
    fn default() -> Self {
        Self { k: 40 }
    }
}

impl Sampler for TopK {
    fn apply(&mut self, candidates: &mut Candidates, _: &mut SamplerCtx<'_>) -> Result<(), Error> {
        apply_top_k(candidates, self.k);
        Ok(())
    }
}

/// Keeps the smallest descending prefix that reaches a probability mass.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct TopP {
    /// Target cumulative probability.
    pub p: f32,
    /// Minimum survivor count.
    pub min_keep: usize,
}

impl Default for TopP {
    fn default() -> Self {
        Self {
            p: 0.95,
            min_keep: 0,
        }
    }
}

impl Sampler for TopP {
    fn apply(&mut self, candidates: &mut Candidates, _: &mut SamplerCtx<'_>) -> Result<(), Error> {
        if self.p.is_nan() || self.p >= 1.0 {
            return Ok(());
        }
        candidates.softmax(false)?;
        if candidates.sorted {
            truncate_top_p(candidates, self.p, self.min_keep);
            return Ok(());
        }

        let len = candidates.buf.len();
        let minimum = self.min_keep.max(1).min(len);
        let mut prefix = minimum.max(32).min(len);
        loop {
            if prefix < len {
                candidates
                    .buf
                    .select_nth_unstable_by(prefix, candidate_order);
                candidates.buf[..prefix].sort_unstable_by(candidate_order);
            } else {
                candidates.buf.sort_unstable_by(candidate_order);
            }

            let mut cumulative = 0.0_f64;
            for (index, candidate) in candidates.buf[..prefix].iter().enumerate() {
                cumulative += candidate.p as f64;
                if cumulative >= self.p as f64 && index + 1 >= minimum {
                    candidates.buf.truncate(index + 1);
                    candidates.sorted = true;
                    return Ok(());
                }
            }
            if prefix == len {
                candidates.sorted = true;
                return Ok(());
            }

            let estimated = if cumulative > 0.0 && self.p > 0.0 {
                ((prefix as f64 * self.p as f64 / cumulative).ceil() as usize).min(len)
            } else {
                prefix.saturating_mul(2).min(len)
            };
            prefix = estimated.max(prefix + 1).max(minimum).min(len);
        }
    }
}

/// Keeps logits above a fraction of the maximum token probability.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct MinP {
    /// Fraction of the maximum token probability.
    pub p: f32,
    /// Minimum survivor count.
    pub min_keep: usize,
}

impl Default for MinP {
    fn default() -> Self {
        Self {
            p: 0.05,
            min_keep: 0,
        }
    }
}

impl Sampler for MinP {
    fn apply(&mut self, candidates: &mut Candidates, _: &mut SamplerCtx<'_>) -> Result<(), Error> {
        if self.p.is_nan() || self.p <= 0.0 || candidates.buf.is_empty() {
            return Ok(());
        }

        if !candidates.sorted {
            let max_logit = candidates
                .buf
                .iter()
                .map(|candidate| candidate.logit)
                .fold(f32::NEG_INFINITY, f32::max);
            let min_logit = max_logit + self.p.ln();
            let matching = candidates
                .buf
                .iter()
                .filter(|candidate| candidate.logit >= min_logit)
                .count();
            if matching > 0 && matching >= self.min_keep {
                candidates
                    .buf
                    .retain(|candidate| candidate.logit >= min_logit);
                return Ok(());
            }
            candidates.sort_desc();
        }

        let min_logit = candidates.buf[0].logit + self.p.ln();
        let mut keep = 1;
        while keep < candidates.buf.len()
            && (candidates.buf[keep].logit >= min_logit || keep < self.min_keep)
        {
            keep += 1;
        }
        candidates.buf.truncate(keep);
        Ok(())
    }
}

/// Keeps candidates closest to the distribution's information content.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Typical {
    /// Target cumulative probability.
    pub p: f32,
    /// Minimum survivor count.
    pub min_keep: usize,
}

impl Default for Typical {
    fn default() -> Self {
        Self {
            p: 1.0,
            min_keep: 0,
        }
    }
}

impl Sampler for Typical {
    fn apply(&mut self, candidates: &mut Candidates, _: &mut SamplerCtx<'_>) -> Result<(), Error> {
        if self.p.is_nan() || self.p >= 1.0 {
            return Ok(());
        }
        candidates.softmax(true)?;
        let entropy = candidates
            .buf
            .iter()
            .filter(|candidate| candidate.p > 0.0)
            .map(|candidate| {
                let p = candidate.p as f64;
                -p * p.ln()
            })
            .sum::<f64>();
        candidates.buf.sort_unstable_by(|left, right| {
            typical_score(left.p, entropy)
                .total_cmp(&typical_score(right.p, entropy))
                .then_with(|| left.id.cmp(&right.id))
        });
        let mut cumulative = 0.0_f64;
        let mut keep = candidates.buf.len();
        for (index, candidate) in candidates.buf.iter().enumerate() {
            cumulative += candidate.p as f64;
            if cumulative > self.p as f64 && (self.min_keep == 0 || index >= self.min_keep - 1) {
                keep = index + 1;
                break;
            }
        }
        candidates.buf.truncate(keep);
        candidates.sorted = false;
        Ok(())
    }
}

/// Removes the low-probability tail using normalized second differences.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct TailFree {
    /// Cumulative normalized second-difference threshold.
    pub z: f32,
    /// Minimum index before truncation is allowed.
    pub min_keep: usize,
}

impl Default for TailFree {
    fn default() -> Self {
        Self {
            z: 1.0,
            min_keep: 0,
        }
    }
}

impl Sampler for TailFree {
    fn apply(&mut self, candidates: &mut Candidates, _: &mut SamplerCtx<'_>) -> Result<(), Error> {
        if self.z.is_nan() || self.z >= 1.0 || candidates.buf.len() <= 2 {
            return Ok(());
        }
        candidates.softmax(true)?;
        let second_difference = |index: usize| {
            let first = candidates.buf[index].p - candidates.buf[index + 1].p;
            let next = candidates.buf[index + 1].p - candidates.buf[index + 2].p;
            (first - next).abs() as f64
        };
        let count = candidates.buf.len() - 2;
        let sum = (0..count).map(second_difference).sum::<f64>();
        let mut cumulative = 0.0;
        let mut keep = candidates.buf.len();
        for index in 0..count {
            let value = if sum > 1e-6 {
                second_difference(index) / sum
            } else {
                1.0 / count as f64
            };
            cumulative += value;
            if cumulative > self.z as f64 && index >= self.min_keep {
                keep = index.max(1);
                break;
            }
        }
        candidates.buf.truncate(keep);
        Ok(())
    }
}

/// Excludes leading candidates whose probabilities exceed a threshold.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Xtc {
    /// Chance that the exclusion stage runs.
    pub probability: f32,
    /// Candidate probability threshold.
    pub threshold: f32,
    /// Minimum survivor count.
    pub min_keep: usize,
}

impl Default for Xtc {
    fn default() -> Self {
        Self {
            probability: 0.0,
            threshold: 0.1,
            min_keep: 0,
        }
    }
}

impl Sampler for Xtc {
    fn apply(
        &mut self,
        candidates: &mut Candidates,
        ctx: &mut SamplerCtx<'_>,
    ) -> Result<(), Error> {
        if self.probability.is_nan()
            || self.threshold.is_nan()
            || self.probability <= 0.0
            || self.threshold > 0.5
            || candidates.buf.len() < 2
        {
            return Ok(());
        }
        if unit(ctx.rng) > self.probability as f64 {
            return Ok(());
        }
        candidates.softmax(true)?;
        let mut pos_last = 0;
        for (index, candidate) in candidates.buf.iter().enumerate() {
            if candidate.p > 0.0 && candidate.p >= self.threshold {
                pos_last = index;
            } else {
                break;
            }
        }
        if pos_last > 0 && candidates.buf.len() - pos_last >= self.min_keep {
            candidates.buf.drain(..pos_last);
        }
        Ok(())
    }
}

/// Applies repetition, frequency, and presence penalties over recent history.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Penalties {
    /// Number of recent tokens to count. Non-positive values disable history.
    pub last_n: i32,
    /// Sign-aware repetition multiplier.
    pub repeat: f32,
    /// Penalty multiplied by token frequency.
    pub freq: f32,
    /// Penalty applied once to each present token.
    pub presence: f32,
}

impl Default for Penalties {
    fn default() -> Self {
        Self {
            last_n: 64,
            repeat: 1.0,
            freq: 0.0,
            presence: 0.0,
        }
    }
}

impl Sampler for Penalties {
    fn apply(
        &mut self,
        candidates: &mut Candidates,
        ctx: &mut SamplerCtx<'_>,
    ) -> Result<(), Error> {
        if self.last_n <= 0 || (self.repeat == 1.0 && self.freq == 0.0 && self.presence == 0.0) {
            return Ok(());
        }
        let window_len = ctx.history.len().min(self.last_n as usize);
        if candidates.counts.capacity() < window_len {
            candidates
                .counts
                .reserve_exact(window_len - candidates.counts.capacity());
        }
        candidates.counts.clear();
        candidates.counts.extend(
            ctx.history[ctx.history.len() - window_len..]
                .iter()
                .map(|&token| (token, 1)),
        );
        candidates.counts.sort_unstable_by_key(|entry| entry.0);
        coalesce_counts(&mut candidates.counts);

        let repeat = if self.repeat.is_finite() && self.repeat > 0.0 {
            self.repeat
        } else {
            1.0
        };
        let frequency = if self.freq.is_finite() {
            self.freq
        } else {
            0.0
        };
        let presence = if self.presence.is_finite() {
            self.presence
        } else {
            0.0
        };
        for candidate in &mut candidates.buf {
            let index = match candidates
                .counts
                .binary_search_by_key(&candidate.id, |entry| entry.0)
            {
                Ok(index) => index,
                Err(_) => continue,
            };
            let count = candidates.counts[index].1 as f32;
            if candidate.logit == f32::NEG_INFINITY {
                continue;
            }
            let repeated = if candidate.logit <= 0.0 {
                candidate.logit as f64 * repeat as f64
            } else {
                candidate.logit as f64 / repeat as f64
            };
            candidate.logit =
                finite_f32(repeated - count as f64 * frequency as f64 - presence as f64);
        }
        candidates.sorted = false;
        Ok(())
    }
}

/// Applies sequence-extension penalties using recent token repetitions.
#[derive(Clone, Debug, PartialEq)]
pub struct Dry {
    /// Penalty multiplier.
    pub multiplier: f32,
    /// Exponential penalty base.
    pub base: f32,
    /// Repetition length allowed before penalties begin.
    pub allowed_length: i32,
    /// Number of recent history tokens to inspect.
    pub last_n: i32,
    /// Token sequences that reset repetition matching.
    pub breakers: Vec<Vec<u32>>,
}

impl Default for Dry {
    fn default() -> Self {
        Self {
            multiplier: 0.0,
            base: 1.75,
            allowed_length: 2,
            last_n: 64,
            breakers: Vec::new(),
        }
    }
}

impl Sampler for Dry {
    fn apply(
        &mut self,
        candidates: &mut Candidates,
        ctx: &mut SamplerCtx<'_>,
    ) -> Result<(), Error> {
        if !self.multiplier.is_finite()
            || !self.base.is_finite()
            || self.multiplier == 0.0
            || self.base < 1.0
            || self.last_n <= 0
        {
            return Ok(());
        }
        let allowed = self.allowed_length.max(0) as usize;
        let n = ctx.history.len().min(self.last_n as usize);
        if n <= allowed {
            return Ok(());
        }
        let history = &ctx.history[ctx.history.len() - n..];
        if candidates.breaker_indices.capacity() < self.breakers.len() {
            candidates
                .breaker_indices
                .reserve_exact(self.breakers.len() - candidates.breaker_indices.capacity());
        }
        candidates.breaker_indices.clear();
        candidates.breaker_indices.extend(
            self.breakers
                .iter()
                .enumerate()
                .filter_map(|(index, sequence)| (!sequence.is_empty()).then_some(index)),
        );
        candidates
            .breaker_indices
            .sort_unstable_by(|&left, &right| {
                breaker_order(&self.breakers[left], &self.breakers[right])
            });

        let recent = |offset: usize| history[n - 1 - offset];
        let mut repeat_limit = n;
        for index in 0..n {
            let token = recent(index);
            let range = breaker_range(&self.breakers, &candidates.breaker_indices, token);
            let mut longest = None;
            for &breaker_index in &candidates.breaker_indices[range] {
                let sequence = &self.breakers[breaker_index];
                let tail = &sequence[1..];
                if tail.len() <= index
                    && longest.map_or(true, |length| tail.len() > length)
                    && tail
                        .iter()
                        .enumerate()
                        .all(|(offset, &expected)| expected == recent(index - offset - 1))
                {
                    longest = Some(tail.len());
                }
            }
            if let Some(length) = longest {
                repeat_limit = index - length;
                break;
            }
        }
        if repeat_limit < allowed {
            return Ok(());
        }

        if candidates.repeat_count.capacity() < n {
            candidates
                .repeat_count
                .reserve_exact(n - candidates.repeat_count.capacity());
        }
        if candidates.max_repeats.capacity() < n {
            candidates
                .max_repeats
                .reserve_exact(n - candidates.max_repeats.capacity());
        }
        candidates.repeat_count.clear();
        candidates.repeat_count.resize(n, 0);
        let last_index = n - 1;
        let (mut left, mut right) = (0, 0);
        for k in 1..n {
            if k > right {
                let mut matched = 0;
                while matched + k < n && recent(matched) == recent(matched + k) {
                    matched += 1;
                }
                candidates.repeat_count[last_index - k] = matched.min(repeat_limit);
                if matched > 0 {
                    left = k;
                    right = k + matched - 1;
                }
            } else {
                let previous = k - left;
                let available = right - k + 1;
                if candidates.repeat_count[last_index - previous] < available {
                    candidates.repeat_count[last_index - k] =
                        candidates.repeat_count[last_index - previous].min(repeat_limit);
                } else {
                    let mut index = right + 1;
                    while index < n && recent(index) == recent(index - k) {
                        index += 1;
                    }
                    candidates.repeat_count[last_index - k] = (index - k).min(repeat_limit);
                    left = k;
                    right = index - 1;
                }
            }
        }

        candidates.max_repeats.clear();
        for index in 0..n - 1 {
            let length = candidates.repeat_count[index];
            if length < allowed {
                continue;
            }
            let token = recent(n - 2 - index);
            candidates.max_repeats.push((token, length));
        }
        candidates.max_repeats.sort_unstable_by_key(|entry| entry.0);
        coalesce_max_repeats(&mut candidates.max_repeats);

        let max_exponent = if self.base > 1.000_001 {
            (88.722_84_f32 / self.base.ln()) as usize
        } else {
            0
        };
        for candidate in &mut candidates.buf {
            let position = match candidates
                .max_repeats
                .binary_search_by_key(&candidate.id, |entry| entry.0)
            {
                Ok(position) => position,
                Err(_) => continue,
            };
            let range = breaker_range(&self.breakers, &candidates.breaker_indices, candidate.id);
            if candidates.breaker_indices[range]
                .iter()
                .any(|&index| self.breakers[index].len() == 1)
            {
                continue;
            }
            let mut exponent = candidates.max_repeats[position].1.saturating_sub(allowed);
            if max_exponent > 0 {
                exponent = exponent.min(max_exponent);
            }
            if candidate.logit != f32::NEG_INFINITY {
                let penalty = self.multiplier as f64 * (self.base as f64).powi(exponent as i32);
                candidate.logit = finite_f32(candidate.logit as f64 - penalty);
            }
        }
        candidates.sorted = false;
        Ok(())
    }
}

/// Stateful Mirostat version 1 sampling.
#[derive(Clone, Debug, PartialEq)]
pub struct MirostatV1 {
    /// Vocabulary size used by the Mirostat estimate.
    pub n_vocab: u32,
    /// Target surprise.
    pub tau: f32,
    /// Learning rate.
    pub eta: f32,
    /// Number of leading probabilities used for the estimate.
    pub m: usize,
    mu: f32,
    last_k: usize,
}

impl MirostatV1 {
    /// Creates a sampler with default Mirostat parameters and a required vocabulary size.
    pub fn new(n_vocab: u32) -> Self {
        Self {
            n_vocab: n_vocab.max(1),
            tau: 5.0,
            eta: 0.1,
            m: 100,
            mu: 10.0,
            last_k: 0,
        }
    }

    /// Returns the current adaptive target.
    pub fn mu(&self) -> f32 {
        self.mu
    }

    #[cfg(test)]
    pub(crate) fn last_k(&self) -> usize {
        self.last_k
    }
}

impl Sampler for MirostatV1 {
    fn apply(
        &mut self,
        candidates: &mut Candidates,
        ctx: &mut SamplerCtx<'_>,
    ) -> Result<(), Error> {
        candidates.softmax(true)?;
        let pair_count = self
            .m
            .saturating_sub(1)
            .min(candidates.buf.len().saturating_sub(1));
        let (mut sum_tb, mut sum_tt) = (0.0_f64, 0.0_f64);
        for index in 0..pair_count {
            let t = (((index + 2) as f64) / ((index + 1) as f64)).ln();
            let b = ((candidates.buf[index].p as f64) / (candidates.buf[index + 1].p as f64)).ln();
            sum_tb += t * b;
            sum_tt += t * t;
        }
        let slope = sum_tb / sum_tt;
        let epsilon = slope - 1.0;
        let numerator = epsilon * 2.0_f64.powf(self.mu as f64);
        let denominator = 1.0 - (self.n_vocab as f64).powf(-epsilon);
        let estimate = (numerator / denominator).powf(1.0 / slope);
        let k = if estimate.is_finite() && estimate >= 1.0 {
            estimate.min(i32::MAX as f64) as i32
        } else {
            1
        };
        self.last_k = k as usize;
        apply_top_k(candidates, k);
        candidates.softmax(true)?;
        apply_dist(candidates, ctx.rng)?;
        let selected = candidates.selected.ok_or(Error::AllMasked)?;
        let surprise = -(candidates.buf[selected].p as f64).log2();
        self.mu -= self.eta * (surprise as f32 - self.tau);
        Ok(())
    }

    fn reset(&mut self) {
        self.mu = 2.0 * self.tau;
        self.last_k = 0;
    }
}

/// Stateful Mirostat version 2 sampling.
#[derive(Clone, Debug, PartialEq)]
pub struct MirostatV2 {
    /// Target surprise.
    pub tau: f32,
    /// Learning rate.
    pub eta: f32,
    mu: f32,
}

impl MirostatV2 {
    /// Creates a sampler from a target surprise and learning rate.
    pub fn new(tau: f32, eta: f32) -> Self {
        Self {
            tau,
            eta,
            mu: 2.0 * tau,
        }
    }

    /// Returns the current adaptive target.
    pub fn mu(&self) -> f32 {
        self.mu
    }
}

impl Default for MirostatV2 {
    fn default() -> Self {
        Self::new(5.0, 0.1)
    }
}

impl Sampler for MirostatV2 {
    fn apply(
        &mut self,
        candidates: &mut Candidates,
        ctx: &mut SamplerCtx<'_>,
    ) -> Result<(), Error> {
        candidates.softmax(true)?;
        let mut keep = candidates.buf.len();
        for (index, candidate) in candidates.buf.iter().enumerate() {
            if -(candidate.p as f64).log2() > self.mu as f64 {
                keep = index.max(1);
                break;
            }
        }
        candidates.buf.truncate(keep);
        candidates.softmax(true)?;
        apply_dist(candidates, ctx.rng)?;
        let selected = candidates.selected.ok_or(Error::AllMasked)?;
        let surprise = -(candidates.buf[selected].p as f64).log2();
        self.mu -= self.eta * (surprise as f32 - self.tau);
        Ok(())
    }

    fn reset(&mut self) {
        self.mu = 2.0 * self.tau;
    }
}

/// Selects the first maximum logit in candidate order.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Greedy;

impl Sampler for Greedy {
    fn apply(&mut self, candidates: &mut Candidates, _: &mut SamplerCtx<'_>) -> Result<(), Error> {
        if candidates.buf.is_empty() {
            return Err(Error::AllMasked);
        }
        let mut selected: Option<usize> = None;
        for (index, candidate) in candidates.buf.iter().enumerate() {
            if !candidate.logit.is_finite() {
                continue;
            }
            let is_larger = match selected {
                Some(current) => {
                    candidate.logit.total_cmp(&candidates.buf[current].logit) == Ordering::Greater
                }
                None => true,
            };
            if is_larger {
                selected = Some(index);
            }
        }
        candidates.selected = selected;
        selected.ok_or(Error::AllMasked).map(|_| ())
    }
}

/// Runs softmax and selects one candidate by cumulative probability.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Dist;

impl Sampler for Dist {
    fn apply(
        &mut self,
        candidates: &mut Candidates,
        ctx: &mut SamplerCtx<'_>,
    ) -> Result<(), Error> {
        apply_dist(candidates, ctx.rng)
    }
}

fn apply_temperature(candidates: &mut Candidates, temperature: f32) {
    if candidates.buf.is_empty() {
        return;
    }
    if temperature.partial_cmp(&0.0) != Some(Ordering::Greater) {
        let maximum = candidates
            .buf
            .iter()
            .enumerate()
            .filter(|(_, candidate)| candidate.logit.is_finite())
            .max_by(|(left_index, left), (right_index, right)| {
                left.logit
                    .total_cmp(&right.logit)
                    .then_with(|| right_index.cmp(left_index))
            })
            .map(|(index, _)| index);
        for (index, candidate) in candidates.buf.iter_mut().enumerate() {
            if Some(index) != maximum {
                candidate.logit = f32::NEG_INFINITY;
            }
        }
        return;
    }

    let would_overflow = temperature.is_finite()
        && candidates.buf.iter().any(|candidate| {
            candidate.logit.is_finite()
                && (candidate.logit as f64 / temperature as f64).abs() > f32::MAX as f64
        });
    let offset = if would_overflow {
        candidates
            .buf
            .iter()
            .filter(|candidate| candidate.logit.is_finite())
            .map(|candidate| candidate.logit)
            .max_by(f32::total_cmp)
            .unwrap_or(0.0)
    } else {
        0.0
    };
    for candidate in &mut candidates.buf {
        if candidate.logit == f32::NEG_INFINITY {
            continue;
        }
        candidate.logit = finite_f32((candidate.logit as f64 - offset as f64) / temperature as f64);
    }
}

fn apply_top_k(candidates: &mut Candidates, k: i32) {
    if k <= 0 || candidates.buf.is_empty() {
        return;
    }
    let keep = (k as usize).min(candidates.buf.len());
    if !candidates.sorted {
        if keep < candidates.buf.len() {
            candidates.buf.select_nth_unstable_by(keep, candidate_order);
            candidates.buf[..keep].sort_unstable_by(candidate_order);
        } else {
            candidates.buf.sort_unstable_by(candidate_order);
        }
        candidates.sorted = true;
    }
    candidates.buf.truncate(keep);
}

fn apply_dist(candidates: &mut Candidates, rng: &mut dyn Rng) -> Result<(), Error> {
    if candidates.buf.is_empty() {
        return Err(Error::AllMasked);
    }
    if !candidates.sanitize_logits() {
        return Err(Error::AllMasked);
    }
    let draw = unit(rng);
    if candidates.buf.len() == 1 {
        candidates.buf[0].p = 1.0;
        candidates.selected = Some(0);
        return Ok(());
    }

    let max_logit = candidates
        .buf
        .iter()
        .map(|candidate| candidate.logit)
        .fold(f32::NEG_INFINITY, f32::max);
    let total = candidates
        .buf
        .iter()
        .map(|candidate| ((candidate.logit - max_logit) as f64).exp())
        .sum::<f64>();
    if !total.is_finite() || total <= 0.0 {
        return Err(Error::AllMasked);
    }
    let target = draw * total;
    let mut running = 0.0;
    let mut selected = None;
    for (index, candidate) in candidates.buf.iter_mut().enumerate() {
        let weight = ((candidate.logit - max_logit) as f64).exp();
        candidate.p = (weight / total) as f32;
        running += weight;
        if selected.is_none() && weight > 0.0 && running >= target {
            selected = Some(index);
        }
    }
    candidates.selected = selected.or_else(|| {
        candidates
            .buf
            .iter()
            .enumerate()
            .rev()
            .find(|(_, candidate)| candidate.p > 0.0)
            .map(|(index, _)| index)
    });
    candidates.selected.ok_or(Error::AllMasked).map(|_| ())
}

fn typical_score(probability: f32, entropy: f64) -> f64 {
    if probability > 0.0 {
        (-(probability as f64).ln() - entropy).abs()
    } else {
        f64::INFINITY
    }
}

fn truncate_top_p(candidates: &mut Candidates, probability: f32, min_keep: usize) {
    let minimum = min_keep.max(1).min(candidates.buf.len());
    let mut cumulative = 0.0_f64;
    let mut keep = candidates.buf.len();
    for (index, candidate) in candidates.buf.iter().enumerate() {
        cumulative += candidate.p as f64;
        if cumulative >= probability as f64 && index + 1 >= minimum {
            keep = index + 1;
            break;
        }
    }
    candidates.buf.truncate(keep);
}

fn coalesce_counts(counts: &mut Vec<(u32, u32)>) {
    let mut write = 0;
    for read in 0..counts.len() {
        if write > 0 && counts[write - 1].0 == counts[read].0 {
            counts[write - 1].1 += counts[read].1;
        } else {
            counts[write] = counts[read];
            write += 1;
        }
    }
    counts.truncate(write);
}

fn coalesce_max_repeats(repeats: &mut Vec<(u32, usize)>) {
    let mut write = 0;
    for read in 0..repeats.len() {
        if write > 0 && repeats[write - 1].0 == repeats[read].0 {
            repeats[write - 1].1 = repeats[write - 1].1.max(repeats[read].1);
        } else {
            repeats[write] = repeats[read];
            write += 1;
        }
    }
    repeats.truncate(write);
}

fn finite_f32(value: f64) -> f32 {
    if value.is_nan() {
        f32::NEG_INFINITY
    } else if value > f32::MAX as f64 {
        f32::MAX
    } else if value < -(f32::MAX as f64) {
        -f32::MAX
    } else {
        value as f32
    }
}

fn breaker_order(left: &[u32], right: &[u32]) -> Ordering {
    left.first()
        .cmp(&right.first())
        .then_with(|| right.len().cmp(&left.len()))
        .then_with(|| left.cmp(right))
}

fn breaker_range(breakers: &[Vec<u32>], indices: &[usize], token: u32) -> core::ops::Range<usize> {
    let start = indices.partition_point(|&index| breakers[index].first().copied().unwrap() < token);
    let end = indices.partition_point(|&index| breakers[index].first().copied().unwrap() <= token);
    start..end
}