toklen 0.2.0

//! Adapted from the `fastokens` library (<https://github.com/crusoecloud/fastokens>).
//! Modifications are made under the same license terms.

use std::fmt::Debug;

use daachorse::DoubleArrayAhoCorasick;
use daachorse::DoubleArrayAhoCorasickBuilder;
use serde::Deserialize;

// ── AddedTokenConfig ────────────────────────────────────────────────────────

/// An entry in the `added_tokens` array of `tokenizer.json`.
#[derive(Clone, Debug, Deserialize)]
pub struct AddedTokenConfig {
    /// Token ID.
    pub id: u32,
    /// Literal text content that triggers this token.
    pub content: String,
    // ------------------------------------------------------------------------
    // Attributes unrelated to counting
    //
    // /// Whether the token should only match as a whole word.
    // #[serde(default)]
    // pub single_word: bool,
    // /// Whether to strip whitespace on the left when matching.
    // #[serde(default)]
    // pub lstrip: bool,
    // /// Whether to strip whitespace on the right when matching.
    // #[serde(default)]
    // pub rstrip: bool,
    // /// Whether the content should be matched against normalized text.
    // #[serde(default)]
    // pub normalized: bool,
    // /// Whether this is a "special" token.
    // #[serde(default)]
    // pub special: bool,
    // ------------------------------------------------------------------------
}

// ── Segment ─────────────────────────────────────────────────────────────────

/// A segment of the input after added-token splitting.
#[derive(Debug, PartialEq, Eq)]
pub enum Segment<'a> {
    /// A span that matched an added token.
    Token(u32),
    /// A span that did not match any added token.
    Text(&'a str),
}

// ── AddedTokens ─────────────────────────────────────────────────────────────

/// A compiled set of added tokens that can be matched against input text.
pub struct AddedTokens {
    daac: DoubleArrayAhoCorasick<u32>,
    token_lens: Vec<u16>,
    start_bytes: Vec<u8>,
    max_token_len: usize,
}

impl AddedTokens {
    /// Build from the `added_tokens` array in `tokenizer.json`.
    pub fn from_configs(configs: &[AddedTokenConfig]) -> Result<Option<Self>, String> {
        if configs.is_empty() {
            return Ok(None);
        }

        let max_id = configs.iter().map(|c| c.id).max().unwrap_or(0);
        let mut token_lens = vec![0u16; (max_id + 1) as usize];

        let mut patterns: Vec<(&str, u32)> = Vec::with_capacity(configs.len());
        for c in configs {
            token_lens[c.id as usize] = u16::try_from(c.content.len()).map_err(|_| {
                format!("AddedToken `{}` is too long (exceed u16::MAX).", c.content)
            })?;

            patterns.push((c.content.as_str(), c.id));
        }

        let daac = DoubleArrayAhoCorasickBuilder::new()
            .match_kind(daachorse::MatchKind::LeftmostLongest)
            .build_with_values(patterns)
            .map_err(|e| format!("error building added-tokens DAAC: {e}"))?;

        // Collect distinct first bytes for memchr prefilter.
        let mut start_set = [false; 256];
        let mut max_token_len = 0;
        for c in configs {
            if let Some(&b) = c.content.as_bytes().first() {
                start_set[b as usize] = true;
            }
            max_token_len = max_token_len.max(c.content.len());
        }

        // The quantity is usually small and does not require pre-alloc.
        let start_bytes: Vec<u8> = start_set
            .iter()
            .enumerate()
            .filter(|&(_, v)| *v)
            .map(|(i, _)| i as u8)
            .collect();

        Ok(Some(Self {
            daac,
            token_lens,
            start_bytes,
            max_token_len,
        }))
    }

    /// Split `input` into segments: spans matching added tokens and spans of
    /// regular text.
    pub fn split<'a>(&self, input: &'a str) -> Vec<Segment<'a>> {
        match self.start_bytes.len() {
            1 => self.split_prefilter(
                input,
                memchr::memchr_iter(self.start_bytes[0], input.as_bytes()),
            ),
            2 => self.split_prefilter(
                input,
                memchr::memchr2_iter(self.start_bytes[0], self.start_bytes[1], input.as_bytes()),
            ),
            3 => self.split_prefilter(
                input,
                memchr::memchr3_iter(
                    self.start_bytes[0],
                    self.start_bytes[1],
                    self.start_bytes[2],
                    input.as_bytes(),
                ),
            ),
            _ => self.split_full_scan(input),
        }
    }

    /// Prefiltered split: only check positions identified by memchr.
    fn split_prefilter<'a>(
        &self,
        input: &'a str,
        candidates: impl Iterator<Item = usize>,
    ) -> Vec<Segment<'a>> {
        let mut segments = Vec::new();
        let mut prev_end = 0;

        for pos in candidates {
            if pos < prev_end {
                continue;
            }
            let mut window_end = (pos + self.max_token_len).min(input.len());
            while window_end < input.len() && !input.is_char_boundary(window_end) {
                window_end += 1;
            }
            let window = &input[pos..window_end];
            if let Some(m) = self.daac.leftmost_find_iter(window).next()
                && m.start() == 0
            {
                if pos > prev_end {
                    segments.push(Segment::Text(&input[prev_end..pos]));
                }
                segments.push(Segment::Token(m.value()));
                prev_end = pos + m.end();
            }
        }

        if prev_end < input.len() {
            segments.push(Segment::Text(&input[prev_end..]));
        }
        if segments.is_empty() && !input.is_empty() {
            segments.push(Segment::Text(input));
        }

        segments
    }

    /// Full-scan fallback for >3 distinct start bytes.
    fn split_full_scan<'a>(&self, input: &'a str) -> Vec<Segment<'a>> {
        let mut segments = Vec::new();
        let mut prev_end = 0;

        for m in self.daac.leftmost_find_iter(input) {
            if m.start() > prev_end {
                segments.push(Segment::Text(&input[prev_end..m.start()]));
            }
            segments.push(Segment::Token(m.value()));
            prev_end = m.end();
        }

        if prev_end < input.len() {
            segments.push(Segment::Text(&input[prev_end..]));
        }

        segments
    }
}

impl Debug for AddedTokens {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let count = self.token_lens.iter().filter(|&&len| len > 0).count();
        f.debug_struct("AddedTokens")
            .field("count", &count)
            .finish()
    }
}