llama_cpp_4/

common_sampler.rs

//! llama.cpp's assembled sampler chain, and the reasoning-budget sampler.
//!
//! [`LlamaSampler`](crate::sampling::LlamaSampler) wraps the individual
//! `llama_sampler_*` primitives, leaving the caller to build a chain. This
//! wraps `common_sampler` — the chain upstream assembles for its own tools —
//! which is a different trade: less control, but it gets right several things
//! that are easy to miss when hand-rolling.
//!
//! - **Ordering.** Penalties and DRY run before truncation samplers, which run
//!   before temperature. A chain in the wrong order silently samples from the
//!   wrong distribution.
//! - **Grammar prefill.** Output-format and tool-call grammars are advanced
//!   past the generation prompt; user grammars are not. See
//!   [`ChatParams::grammar_sampler`](crate::chat::ChatParams::grammar_sampler)
//!   for what goes wrong otherwise.
//! - **Model-declared suppress tokens** are merged into the logit bias, so
//!   `tokenizer.ggml.suppress_tokens` is honoured without the caller knowing it
//!   exists.
//! - **Reasoning budget**, created whenever a lazy grammar is active so a
//!   thinking block can be force-closed.
//! - **Speculative acceptance** via [`CommonSampler::sample_and_accept_n`].

use std::ffi::{c_char, CString};
use std::ptr::NonNull;

use llama_cpp_sys_4 as sys;

use crate::context::LlamaContext;
use crate::model::LlamaModel;
use crate::token::LlamaToken;

/// Errors from the common-sampler layer.
///
/// An alias for [`ShimError`](crate::shim::ShimError) — every shim-backed
/// module shares one error type, since they share one status enum and one error
/// buffer.
pub type CommonSamplerError = crate::shim::ShimError;

use crate::shim::{check_status, last_error, read_i32s, read_string, read_tokens, Result};

/// Which of llama.cpp's samplers to run, and in what order.
///
/// Values match `common_sampler_type`; the default chain is penalties → DRY →
/// top-n-sigma → top-k → typical-p → top-p → min-p → XTC → temperature.
// Discriminants are an ABI contract with `common_sampler_type`, and they are
// not contiguous: 5 is a retired TFS-Z slot upstream left as a comment. Taken
// verbatim from `common/common.h`; the round-trip test pins them.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(i32)]
pub enum CommonSamplerType {
    /// DRY repetition penalty.
    Dry = 1,
    /// Top-k truncation.
    TopK = 2,
    /// Top-p (nucleus) truncation.
    TopP = 3,
    /// Min-p truncation.
    MinP = 4,
    /// Typical-p truncation.
    TypicalP = 6,
    /// Temperature scaling.
    Temperature = 7,
    /// Exclude Top Choices.
    Xtc = 8,
    /// Infill sampler, for fill-in-the-middle.
    Infill = 9,
    /// Repetition / frequency / presence penalties.
    Penalties = 10,
    /// Top-n-sigma truncation.
    TopNSigma = 11,
    /// Adaptive-p, which targets a probability rather than a cutoff.
    AdaptiveP = 12,
}

impl CommonSamplerType {
    /// llama.cpp's own name for this sampler, e.g. `"top_k"`.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Failed`] if llama.cpp cannot name it.
    pub fn name(self) -> Result<String> {
        // SAFETY: the closure hands the shim a buffer pointer together with its capacity,
        // and the shim reports how much it needed; nothing is read back unless it
        // reported success.
        read_string(|buf, len, expected| unsafe {
            sys::common_shim_sampler_type_to_str(self as i32, buf, len, expected)
        })
    }

    /// Parse llama.cpp's sampler names into an ordering.
    ///
    /// Unrecognised names are dropped by llama.cpp rather than rejected, so a
    /// shorter result than `names` means something was not understood.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Nul`] if a name contains an interior NUL.
    pub fn from_names(names: &[&str]) -> Result<Vec<i32>> {
        let c_names: Vec<CString> = names
            .iter()
            .map(|n| CString::new(*n))
            .collect::<std::result::Result<_, _>>()?;
        let ptrs: Vec<*const c_char> = c_names.iter().map(|c| c.as_ptr()).collect();

        // SAFETY: the closure hands the shim a buffer pointer together with its capacity,
        // and the shim reports how much it needed; nothing is read back unless it
        // reported success.
        read_i32s(|out, cap, len| unsafe {
            sys::common_shim_sampler_types_from_names(ptrs.as_ptr(), ptrs.len(), out, cap, len)
        })
    }
}

/// How a grammar was obtained, which decides whether the generation prompt is
/// prefilled into it.
///
/// Getting this wrong is silent: prefilling a user grammar consumes tokens it
/// never expected, and *not* prefilling a tool-call grammar forces the model to
/// re-emit the generation prompt.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum GrammarSource {
    /// No grammar.
    #[default]
    None,
    /// Supplied verbatim by the caller. Never prefilled.
    User,
    /// Generated from a JSON schema. Prefilled.
    OutputFormat,
    /// Generated by a chat template for tool calling. Prefilled.
    ToolCalls,
}

impl GrammarSource {
    // The C enum is unsigned; the field it feeds is `int32_t`.
    #[allow(clippy::cast_possible_wrap)]
    fn as_raw(self) -> i32 {
        let raw = match self {
            Self::None => sys::COMMON_SHIM_GRAMMAR_NONE,
            Self::User => sys::COMMON_SHIM_GRAMMAR_USER,
            Self::OutputFormat => sys::COMMON_SHIM_GRAMMAR_OUTPUT_FORMAT,
            Self::ToolCalls => sys::COMMON_SHIM_GRAMMAR_TOOL_CALLS,
        };
        raw as i32
    }
}

/// Parameters for [`CommonSampler::new`], seeded from llama.cpp's defaults.
///
/// The scalar knobs are a plain public struct ([`CommonSamplerScalars`]);
/// everything backed by a container is set through a method, which keeps this
/// wrapper stable when upstream adds fields.
pub struct CommonSamplerParams {
    raw: NonNull<sys::common_shim_sampler_params>,
}

impl std::fmt::Debug for CommonSamplerParams {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CommonSamplerParams")
            .field("scalars", &self.scalars())
            .finish_non_exhaustive()
    }
}

// SAFETY: the handle owns a `common_params_sampling` with no shared state; the
// shim's only global is a thread-local error buffer.
unsafe impl Send for CommonSamplerParams {}

impl Drop for CommonSamplerParams {
    fn drop(&mut self) {
        // SAFETY: this runs once from `Drop` on a non-null handle, and nothing reads it
        // afterwards.
        unsafe { sys::common_shim_sampler_params_free(self.raw.as_ptr()) }
    }
}

/// The scalar half of llama.cpp's sampling parameters.
pub type CommonSamplerScalars = sys::common_shim_sampler_scalars;

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

impl CommonSamplerParams {
    /// Allocate with llama.cpp's defaults (temp 0.8, top-k 40, top-p 0.95, …).
    ///
    /// # Panics
    ///
    /// Panics if the allocation fails.
    #[must_use]
    pub fn new() -> Self {
        // SAFETY: no borrowed pointer is involved beyond arguments alive for the call;
        // the returned handle is checked for null before being wrapped.
        let raw = unsafe { sys::common_shim_sampler_params_init() };
        Self {
            raw: NonNull::new(raw).expect("common_shim_sampler_params_init returned null"),
        }
    }

    /// Read the scalar knobs.
    #[must_use]
    pub fn scalars(&self) -> CommonSamplerScalars {
        // SAFETY: the struct is a plain-old-data `repr(C)` type with no pointer or
        // reference fields, so an all-zero value is a valid instance; the shim
        // fills it in immediately below.
        let mut out: CommonSamplerScalars = unsafe { std::mem::zeroed() };
        // SAFETY: the `CString` arguments are alive for the duration of the call, and any
        // out-param is a stack local the shim writes through.
        unsafe { sys::common_shim_sampler_params_get_scalars(self.raw.as_ptr(), &raw mut out) };
        out
    }

    /// Replace the scalar knobs. Read [`Self::scalars`] first and modify it, so
    /// fields you do not care about keep llama.cpp's defaults.
    pub fn set_scalars(&mut self, scalars: &CommonSamplerScalars) {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { sys::common_shim_sampler_params_set_scalars(self.raw.as_ptr(), scalars) }
    }

    /// Constrain sampling with a GBNF grammar.
    ///
    /// `source` decides whether the generation prompt is prefilled — see
    /// [`GrammarSource`].
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Nul`] for an interior NUL in `grammar`.
    pub fn set_grammar(&mut self, grammar: &str, source: GrammarSource, lazy: bool) -> Result<()> {
        let c_grammar = CString::new(grammar)?;
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let status = unsafe {
            sys::common_shim_sampler_params_set_grammar(
                self.raw.as_ptr(),
                c_grammar.as_ptr(),
                source.as_raw(),
                lazy,
            )
        };
        check_status(status)
    }

    /// Add a trigger that activates a lazy grammar.
    ///
    /// `kind` matches [`GrammarTrigger::kind`](crate::chat::GrammarTrigger):
    /// `"token"`, `"word"`, `"pattern"` or `"pattern_full"`.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::InvalidArg`] for an unknown `kind`, or
    /// [`CommonSamplerError::Nul`] for an interior NUL.
    pub fn add_grammar_trigger(
        &mut self,
        kind: &str,
        value: &str,
        token: LlamaToken,
    ) -> Result<()> {
        let raw_kind = match kind {
            "token" => 0,
            "word" => 1,
            "pattern" => 2,
            "pattern_full" => 3,
            _ => return Err(CommonSamplerError::InvalidArg),
        };
        let c_value = CString::new(value)?;
        // SAFETY: the `CString` arguments are alive for the duration of the call, and any
        // out-param is a stack local the shim writes through.
        let status = unsafe {
            sys::common_shim_sampler_params_add_grammar_trigger(
                self.raw.as_ptr(),
                raw_kind,
                c_value.as_ptr(),
                token.0,
            )
        };
        check_status(status)
    }

    /// Set the generation prompt the grammar should be advanced past.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Nul`] for an interior NUL.
    pub fn set_generation_prompt(&mut self, prompt: &str) -> Result<()> {
        let c_prompt = CString::new(prompt)?;
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let status = unsafe {
            sys::common_shim_sampler_params_set_generation_prompt(
                self.raw.as_ptr(),
                c_prompt.as_ptr(),
            )
        };
        check_status(status)
    }

    /// Bias a token's logit. Use `f32::NEG_INFINITY` to ban it outright.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::InvalidArg`] if the params handle is bad.
    pub fn add_logit_bias(&mut self, token: LlamaToken, bias: f32) -> Result<()> {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let status = unsafe {
            sys::common_shim_sampler_params_add_logit_bias(self.raw.as_ptr(), token.0, bias)
        };
        check_status(status)
    }

    /// Replace the sampler ordering.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::InvalidArg`] if the params handle is bad.
    pub fn set_samplers(&mut self, samplers: &[CommonSamplerType]) -> Result<()> {
        let raw: Vec<i32> = samplers.iter().map(|s| *s as i32).collect();
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let status = unsafe {
            sys::common_shim_sampler_params_set_samplers(self.raw.as_ptr(), raw.as_ptr(), raw.len())
        };
        check_status(status)
    }

    /// Replace the DRY sequence breakers (default: newline, `:`, `"`, `*`).
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Nul`] for an interior NUL.
    pub fn set_dry_breakers(&mut self, breakers: &[&str]) -> Result<()> {
        let c_breakers: Vec<CString> = breakers
            .iter()
            .map(|b| CString::new(*b))
            .collect::<std::result::Result<_, _>>()?;
        let ptrs: Vec<*const c_char> = c_breakers.iter().map(|c| c.as_ptr()).collect();
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let status = unsafe {
            sys::common_shim_sampler_params_set_dry_breakers(
                self.raw.as_ptr(),
                ptrs.as_ptr(),
                ptrs.len(),
            )
        };
        check_status(status)
    }

    /// Cap how many tokens the model may spend inside a reasoning block.
    ///
    /// `start` is the tokenized opening tag (`<think>`), `ends` the closing
    /// tags — the first doubles as the forcing sequence. `forced` is emitted
    /// when the budget runs out, typically a truncation message followed by the
    /// closing tag. Set the budget itself through
    /// [`CommonSamplerScalars::reasoning_budget_tokens`].
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Nul`] for an interior NUL in `message`.
    pub fn set_reasoning_budget(
        &mut self,
        start: &[LlamaToken],
        ends: &[Vec<LlamaToken>],
        forced: &[LlamaToken],
        message: &str,
    ) -> Result<()> {
        let start_raw: Vec<i32> = start.iter().map(|t| t.0).collect();
        let (ends_flat, end_lens) = flatten(ends);
        let forced_raw: Vec<i32> = forced.iter().map(|t| t.0).collect();
        let c_message = CString::new(message)?;

        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let status = unsafe {
            sys::common_shim_sampler_params_set_reasoning_budget(
                self.raw.as_ptr(),
                start_raw.as_ptr(),
                start_raw.len(),
                ends_flat.as_ptr(),
                end_lens.as_ptr(),
                end_lens.len(),
                forced_raw.as_ptr(),
                forced_raw.len(),
                c_message.as_ptr(),
            )
        };
        check_status(status)
    }
}

/// Flatten a slice of token sequences into `(data, lengths)`, which is how the
/// shim takes ragged arrays without arrays-of-pointers.
fn flatten(seqs: &[Vec<LlamaToken>]) -> (Vec<i32>, Vec<usize>) {
    let mut data = Vec::new();
    let mut lens = Vec::with_capacity(seqs.len());
    for seq in seqs {
        lens.push(seq.len());
        data.extend(seq.iter().map(|t| t.0));
    }
    (data, lens)
}

/// llama.cpp's assembled sampler chain.
pub struct CommonSampler {
    raw: NonNull<sys::common_shim_sampler>,
}

impl std::fmt::Debug for CommonSampler {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CommonSampler")
            .field("seed", &self.seed())
            .finish_non_exhaustive()
    }
}

// SAFETY: the handle owns a `common_sampler` that is not shared; every method
// takes `&mut self` where it mutates.
unsafe impl Send for CommonSampler {}

impl Drop for CommonSampler {
    fn drop(&mut self) {
        // SAFETY: this runs once from `Drop` on a non-null handle, and nothing reads it
        // afterwards.
        unsafe { sys::common_shim_sampler_free(self.raw.as_ptr()) }
    }
}

impl CommonSampler {
    /// Assemble the chain for `model`.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Failed`] if llama.cpp cannot build it —
    /// most often a grammar that does not parse.
    pub fn new(model: &LlamaModel, params: &mut CommonSamplerParams) -> Result<Self> {
        let raw =
            // SAFETY: the arguments are either scalars or pointers that outlive the call, and
            // the library owns whatever it returns.
            unsafe { sys::common_shim_sampler_init(model.model.as_ptr(), params.raw.as_ptr()) };
        NonNull::new(raw)
            .map(|raw| Self { raw })
            .ok_or_else(|| CommonSamplerError::Failed(last_error()))
    }

    /// Sample a token from the logits at `idx`.
    ///
    /// `grammar_first` applies the grammar before the other samplers rather
    /// than after. Upstream's default is `false`: sample first and only fall
    /// back to grammar-constrained resampling if the pick is rejected, which is
    /// much cheaper than filtering the whole vocabulary every step.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Failed`] if llama.cpp throws.
    pub fn sample(
        &mut self,
        ctx: &mut LlamaContext<'_>,
        idx: i32,
        grammar_first: bool,
    ) -> Result<LlamaToken> {
        let mut status = sys::LLAMA_SHIM_OK;
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let token = unsafe {
            sys::common_shim_sampler_sample(
                self.raw.as_ptr(),
                ctx.context.as_ptr(),
                idx,
                grammar_first,
                &raw mut status,
            )
        };
        check_status(status)?;
        Ok(LlamaToken(token))
    }

    /// Feed a token back into the sampler's history.
    ///
    /// `is_generated` marks a token the model produced, as opposed to one from
    /// the prompt; penalties and the grammar treat the two differently.
    pub fn accept(&mut self, token: LlamaToken, is_generated: bool) {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { sys::common_shim_sampler_accept(self.raw.as_ptr(), token.0, is_generated) }
    }

    /// Validate a speculative draft and return the accepted prefix.
    ///
    /// This is the acceptance half of speculative decoding: given `draft`
    /// tokens proposed by a drafter and a target-model forward pass covering
    /// them, it returns every draft token the target agrees with, plus one
    /// freshly sampled token at the first divergence. The result is therefore
    /// never empty and never longer than `draft.len() + 1`.
    ///
    /// Accepted tokens are also fed to [`Self::accept`] internally, so the
    /// caller must not do so again.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Failed`] if llama.cpp throws.
    pub fn sample_and_accept_n(
        &mut self,
        ctx: &mut LlamaContext<'_>,
        draft: &[LlamaToken],
        grammar_first: bool,
    ) -> Result<Vec<LlamaToken>> {
        let raw_draft: Vec<i32> = draft.iter().map(|t| t.0).collect();
        // SAFETY: the closure hands the shim a buffer pointer together with its capacity,
        // and the shim reports how much it needed; nothing is read back unless it
        // reported success.
        read_tokens(|out, cap, len| unsafe {
            sys::common_shim_sampler_sample_and_accept_n(
                self.raw.as_ptr(),
                ctx.context.as_ptr(),
                raw_draft.as_ptr(),
                raw_draft.len(),
                grammar_first,
                out,
                cap,
                len,
            )
        })
    }

    /// Clear all history and grammar state, keeping the configuration.
    pub fn reset(&mut self) {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { sys::common_shim_sampler_reset(self.raw.as_ptr()) }
    }

    /// Deep-copy this sampler, state included.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Failed`] if llama.cpp returns null.
    pub fn try_clone(&self) -> Result<Self> {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let raw = unsafe { sys::common_shim_sampler_clone(self.raw.as_ptr()) };
        NonNull::new(raw)
            .map(|raw| Self { raw })
            .ok_or_else(|| CommonSamplerError::Failed(last_error()))
    }

    /// The seed actually in use, after any `LLAMA_DEFAULT_SEED` was resolved.
    #[must_use]
    pub fn seed(&self) -> u32 {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { sys::common_shim_sampler_get_seed(self.raw.as_ptr()) }
    }

    /// The most recently sampled token.
    #[must_use]
    pub fn last(&self) -> LlamaToken {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        LlamaToken(unsafe { sys::common_shim_sampler_last(self.raw.as_ptr()) })
    }

    /// End the current reasoning block now, as if the budget had run out.
    ///
    /// Returns `false` when there is no budget sampler or it is not counting.
    pub fn force_end_reasoning(&mut self) -> bool {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { sys::common_shim_sampler_reasoning_budget_force(self.raw.as_ptr()) }
    }

    /// Describe the assembled chain, e.g. `"penalties -> top_k -> temp"`.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Failed`] if llama.cpp throws.
    pub fn describe(&self) -> Result<String> {
        // SAFETY: the closure hands the shim a buffer pointer together with its capacity,
        // and the shim reports how much it needed; nothing is read back unless it
        // reported success.
        read_string(|buf, len, expected| unsafe {
            sys::common_shim_sampler_print(self.raw.as_ptr(), buf, len, expected)
        })
    }

    /// Detokenize the last `n` sampled tokens.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Failed`] if llama.cpp throws.
    pub fn prev_str(&mut self, ctx: &mut LlamaContext<'_>, n: i32) -> Result<String> {
        // SAFETY: the closure hands the shim a buffer pointer together with its capacity,
        // and the shim reports how much it needed; nothing is read back unless it
        // reported success.
        read_string(|buf, len, expected| unsafe {
            sys::common_shim_sampler_prev_str(
                self.raw.as_ptr(),
                ctx.context.as_ptr(),
                n,
                buf,
                len,
                expected,
            )
        })
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Reasoning budget
// ─────────────────────────────────────────────────────────────────────────────

/// Where a [`ReasoningBudget`] is in its state machine.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReasoningBudgetState {
    /// Passing tokens through, watching for a start tag.
    Idle,
    /// Inside a reasoning block, counting down.
    Counting,
    /// Budget spent; emitting the forced closing sequence.
    Forcing,
    /// Budget spent, but finishing a partial UTF-8 sequence first.
    WaitingUtf8,
    /// Finished; passing everything through.
    Done,
}

impl ReasoningBudgetState {
    // Unsigned in C, signed on the wire — bridge in both directions.
    #[allow(clippy::cast_possible_wrap)]
    fn from_raw(raw: i32) -> Option<Self> {
        if raw == sys::COMMON_SHIM_RBUDGET_IDLE as i32 {
            Some(Self::Idle)
        } else if raw == sys::COMMON_SHIM_RBUDGET_COUNTING as i32 {
            Some(Self::Counting)
        } else if raw == sys::COMMON_SHIM_RBUDGET_FORCING as i32 {
            Some(Self::Forcing)
        } else if raw == sys::COMMON_SHIM_RBUDGET_WAITING_UTF8 as i32 {
            Some(Self::WaitingUtf8)
        } else if raw == sys::COMMON_SHIM_RBUDGET_DONE as i32 {
            Some(Self::Done)
        } else {
            None
        }
    }

    #[allow(clippy::cast_possible_wrap)]
    fn as_raw(self) -> i32 {
        let raw = match self {
            Self::Idle => sys::COMMON_SHIM_RBUDGET_IDLE,
            Self::Counting => sys::COMMON_SHIM_RBUDGET_COUNTING,
            Self::Forcing => sys::COMMON_SHIM_RBUDGET_FORCING,
            Self::WaitingUtf8 => sys::COMMON_SHIM_RBUDGET_WAITING_UTF8,
            Self::Done => sys::COMMON_SHIM_RBUDGET_DONE,
        };
        raw as i32
    }
}

/// A sampler that caps how long a model may think.
///
/// Reasoning models emit an unbounded `<think>…</think>` block before their
/// answer, and nothing in the model stops them running to the context limit.
/// This watches for the start tag, counts down, and — when the budget is spent
/// — forces the closing sequence token by token, masking everything else.
///
/// It is the piece llama.cpp pairs with a *lazy* grammar: the grammar leaves
/// reasoning unconstrained by design, so something else has to bound it.
///
/// Build one with [`ReasoningBudget::new`] and add it to a
/// [`LlamaSampler::chain`](crate::sampling::LlamaSampler::chain), or let
/// [`CommonSampler`] assemble it for you via
/// [`CommonSamplerParams::set_reasoning_budget`].
#[derive(Debug)]
pub struct ReasoningBudget {
    sampler: crate::sampling::LlamaSampler,
}

impl ReasoningBudget {
    /// Build a budget sampler.
    ///
    /// * `starts` — tokenized opening tags; any one arms the countdown.
    /// * `ends` — tokenized closing tags; any one disarms it naturally.
    /// * `forced` — emitted when the budget runs out, usually a short message
    ///   followed by a closing tag.
    /// * `budget` — tokens allowed inside the block.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Failed`] if llama.cpp returns null.
    pub fn new(
        model: &LlamaModel,
        starts: &[Vec<LlamaToken>],
        ends: &[Vec<LlamaToken>],
        forced: &[LlamaToken],
        budget: i32,
    ) -> Result<Self> {
        Self::with_initial_state(
            model,
            starts,
            ends,
            forced,
            budget,
            ReasoningBudgetState::Idle,
        )
    }

    /// [`Self::new`] with an explicit starting state — use
    /// [`ReasoningBudgetState::Counting`] when the prompt already opened a
    /// reasoning block, which is what a template with `thinking_forced_open`
    /// produces.
    ///
    /// # Errors
    ///
    /// Returns [`CommonSamplerError::Failed`] if llama.cpp returns null.
    pub fn with_initial_state(
        model: &LlamaModel,
        starts: &[Vec<LlamaToken>],
        ends: &[Vec<LlamaToken>],
        forced: &[LlamaToken],
        budget: i32,
        initial_state: ReasoningBudgetState,
    ) -> Result<Self> {
        let (starts_flat, start_lens) = flatten(starts);
        let (ends_flat, end_lens) = flatten(ends);
        let forced_raw: Vec<i32> = forced.iter().map(|t| t.0).collect();

        // SAFETY: the arguments are either scalars or pointers that outlive the call, and
        // the library owns whatever it returns.
        let raw = unsafe {
            sys::common_shim_reasoning_budget_init(
                model.get_vocab().vocab.as_ref(),
                starts_flat.as_ptr(),
                start_lens.as_ptr(),
                start_lens.len(),
                ends_flat.as_ptr(),
                end_lens.as_ptr(),
                end_lens.len(),
                forced_raw.as_ptr(),
                forced_raw.len(),
                budget,
                initial_state.as_raw(),
            )
        };
        let ptr = NonNull::new(raw).ok_or_else(|| CommonSamplerError::Failed(last_error()))?;
        // SAFETY: upstream returns a sampler the caller owns and frees with
        // `llama_sampler_free`, which is what `LlamaSampler` does.
        Ok(Self {
            // SAFETY: `from_raw_ptr` is `unsafe` because it takes ownership; `ptr` is the
            // non-null sampler the shim just created and nothing else holds it.
            sampler: unsafe { crate::sampling::LlamaSampler::from_raw_ptr(ptr) },
        })
    }

    /// Where the state machine currently is.
    #[must_use]
    pub fn state(&self) -> Option<ReasoningBudgetState> {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        let raw = unsafe {
            sys::common_shim_reasoning_budget_get_state(self.sampler.as_ptr().cast_const())
        };
        ReasoningBudgetState::from_raw(raw)
    }

    /// Cut the reasoning block short now. Returns `false` if it was not
    /// counting.
    pub fn force_end(&mut self) -> bool {
        // SAFETY: the handle is non-null and kept alive by this borrow, and the remaining
        // arguments are scalars the library validates itself.
        unsafe { sys::common_shim_reasoning_budget_force(self.sampler.as_ptr()) }
    }

    /// Take the underlying sampler, to add to a chain.
    #[must_use]
    pub fn into_sampler(self) -> crate::sampling::LlamaSampler {
        self.sampler
    }
}

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

    #[test]
    fn params_start_from_llama_cpp_defaults() {
        let params = CommonSamplerParams::new();
        let s = params.scalars();
        // Pinned against upstream's documented defaults; a bump that changes
        // them silently changes output for every caller.
        assert!((s.temp - 0.80).abs() < 1e-6, "temp = {}", s.temp);
        assert_eq!(s.top_k, 40);
        assert!((s.top_p - 0.95).abs() < 1e-6, "top_p = {}", s.top_p);
        assert!((s.min_p - 0.05).abs() < 1e-6, "min_p = {}", s.min_p);
        assert_eq!(s.mirostat, 0);
        assert_eq!(s.reasoning_budget_tokens, -1, "budget disabled by default");
    }

    #[test]
    fn scalars_round_trip() {
        let mut params = CommonSamplerParams::new();
        let mut s = params.scalars();
        s.temp = 0.25;
        s.top_k = 7;
        s.reasoning_budget_tokens = 128;
        params.set_scalars(&s);

        let back = params.scalars();
        assert!((back.temp - 0.25).abs() < 1e-6);
        assert_eq!(back.top_k, 7);
        assert_eq!(back.reasoning_budget_tokens, 128);
    }

    /// Setting one field must not disturb the others — the getter/setter pair
    /// copies the whole struct, so a missed field would silently reset.
    #[test]
    fn setting_scalars_preserves_untouched_fields() {
        let mut params = CommonSamplerParams::new();
        let before = params.scalars();
        let mut s = before;
        s.top_k = 3;
        params.set_scalars(&s);
        let after = params.scalars();

        assert_eq!(after.top_k, 3);
        assert!((after.top_p - before.top_p).abs() < 1e-6);
        assert!((after.dry_base - before.dry_base).abs() < 1e-6);
        assert_eq!(after.penalty_last_n, before.penalty_last_n);
        assert_eq!(after.seed, before.seed);
    }

    #[test]
    fn sampler_type_names_match_upstream() {
        assert_eq!(CommonSamplerType::TopK.name().unwrap(), "top_k");
        assert_eq!(CommonSamplerType::TopP.name().unwrap(), "top_p");
        assert_eq!(
            CommonSamplerType::Temperature.name().unwrap(),
            "temperature"
        );
    }

    /// The discriminants are an ABI contract with `common_sampler_type`; if
    /// they drift, `set_samplers` silently reorders the chain.
    #[test]
    fn sampler_type_discriminants_round_trip_through_names() {
        for ty in [
            CommonSamplerType::Dry,
            CommonSamplerType::TopK,
            CommonSamplerType::TopP,
            CommonSamplerType::MinP,
            CommonSamplerType::TypicalP,
            CommonSamplerType::Temperature,
            CommonSamplerType::Xtc,
            CommonSamplerType::Infill,
            CommonSamplerType::Penalties,
            CommonSamplerType::TopNSigma,
            CommonSamplerType::AdaptiveP,
        ] {
            let name = ty.name().unwrap();
            let parsed = CommonSamplerType::from_names(&[&name]).unwrap();
            assert_eq!(parsed, vec![ty as i32], "{name} did not round-trip");
        }
    }

    #[test]
    fn unknown_sampler_names_are_dropped() {
        let parsed = CommonSamplerType::from_names(&["top_k", "not_a_sampler"]).unwrap();
        assert_eq!(parsed, vec![CommonSamplerType::TopK as i32]);
    }

    #[test]
    fn grammar_trigger_rejects_unknown_kind() {
        let mut params = CommonSamplerParams::new();
        assert!(matches!(
            params.add_grammar_trigger("nonsense", "x", LlamaToken(-1)),
            Err(CommonSamplerError::InvalidArg)
        ));
    }

    #[test]
    fn grammar_trigger_accepts_every_known_kind() {
        let mut params = CommonSamplerParams::new();
        for kind in ["token", "word", "pattern", "pattern_full"] {
            params
                .add_grammar_trigger(kind, "<tool_call>", LlamaToken(1))
                .unwrap_or_else(|e| panic!("{kind} rejected: {e}"));
        }
    }

    #[test]
    fn interior_nul_is_rejected_not_truncated() {
        let mut params = CommonSamplerParams::new();
        assert!(matches!(
            params.set_grammar("root ::= \0 \"a\"", GrammarSource::User, false),
            Err(CommonSamplerError::Nul(_))
        ));
        assert!(matches!(
            params.set_generation_prompt("a\0b"),
            Err(CommonSamplerError::Nul(_))
        ));
    }

    #[test]
    fn flatten_produces_matching_data_and_lengths() {
        let seqs = vec![
            vec![LlamaToken(1), LlamaToken(2)],
            vec![],
            vec![LlamaToken(3)],
        ];
        let (data, lens) = flatten(&seqs);
        assert_eq!(data, vec![1, 2, 3]);
        assert_eq!(lens, vec![2, 0, 1]);
    }
}