daggrs 0.1.0

mod iter;

pub use iter::TrieFindIter;

use alloc::{
    collections::{BTreeMap, VecDeque},
    vec::Vec,
};

use crate::daac::DoubleArrayAhoCorasick;
use crate::types::{Match, MatchKind, Output};

/// Sentinel value indicating a "dead" state where searching should stop.
/// Used in leftmost matching modes to prevent overlapping matches.
pub const DEAD_STATE: u32 = u32::MAX;

/// A state in the trie-based Aho-Corasick automaton.
#[derive(Debug, Clone)]
pub struct TrieState {
    /// Transitions to child states, keyed by input byte.
    pub edges: BTreeMap<u8, u32>,
    /// Failure link: the longest proper suffix that is also a prefix in the trie.
    pub fail: u32,
    /// Output index if this state completes a pattern, `None` otherwise.
    pub outpos: Option<u32>,
}

/// A trie-based Aho-Corasick automaton for multi-pattern matching.
///
/// The `Trie` provides a simple tree-based implementation. For better performance
/// on large pattern sets, compile to [`DoubleArrayAhoCorasick`] using [`compile()`](Self::compile).
#[derive(Debug, Clone)]
pub struct Trie {
    pub states: Vec<TrieState>,
    pub outputs: Vec<Output>,
    pub match_kind: MatchKind,
    /// State ID of the continuation anchor for WordPiece (e.g., "##" state).
    pub anchor: Option<u32>,
}

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

impl Trie {
    /// Creates a new empty trie with only the root state.
    pub fn new() -> Self {
        let root = TrieState {
            edges: BTreeMap::new(),
            fail: 0,
            outpos: None,
        };
        Trie {
            states: vec![root],
            outputs: Vec::new(),
            match_kind: MatchKind::default(),
            anchor: None,
        }
    }

    /// Adds a pattern to the trie with the given identifier.
    ///
    /// Patterns must be added before calling `build()`.
    pub fn add(&mut self, pattern: &[u8], pattern_id: u32) {
        let outpos = self.outputs.len() as u32;
        self.outputs.push(Output {
            pattern_id,
            length: pattern.len() as u32,
            parent: u32::MAX,
        });

        let mut state = 0;
        for &byte in pattern {
            state = match self.states[state].edges.get(&byte) {
                Some(&next) => next as usize,
                None => {
                    let next_id = self.states.len() as u32;
                    self.states.push(TrieState {
                        edges: BTreeMap::new(),
                        fail: 0,
                        outpos: None,
                    });
                    self.states[state].edges.insert(byte, next_id);
                    next_id as usize
                }
            };
        }

        self.states[state].outpos = Some(outpos);
    }

    /// Builds the automaton by computing failure links and output chains.
    ///
    /// Must be called after adding all patterns and before searching.
    /// For WordPiece mode, use `build_wordpiece()` instead.
    pub fn build(&mut self, match_kind: MatchKind) {
        self.match_kind = match_kind;

        if match_kind == MatchKind::WordPiece {
            self.build_wordpiece(b"##");
            return;
        }

        if match_kind == MatchKind::LeftmostFirst {
            self.prune_outputs_for_leftmost_first();
        }

        let leftmost = matches!(
            match_kind,
            MatchKind::LeftmostFirst | MatchKind::LeftmostLongest
        );

        let mut queue: VecDeque<usize> = VecDeque::new();

        // Initialize root's children
        let root_children: Vec<u32> = self.states[0].edges.values().copied().collect();
        for child in root_children {
            queue.push_back(child as usize);
            self.states[child as usize].fail =
                if leftmost && self.states[child as usize].outpos.is_some() {
                    DEAD_STATE
                } else {
                    0
                };
        }

        // BFS to compute failure links
        while let Some(state) = queue.pop_front() {
            let parent_fail = self.states[state].fail;
            let edges: Vec<(u8, u32)> = self.states[state]
                .edges
                .iter()
                .map(|(&k, &v)| (k, v))
                .collect();

            for (byte, child) in edges {
                queue.push_back(child as usize);

                // Propagate DEAD state to children
                if leftmost && parent_fail == DEAD_STATE {
                    self.states[child as usize].fail = DEAD_STATE;
                    continue;
                }

                // Follow failure links to find the longest proper suffix
                let mut fail_state = parent_fail as usize;
                while fail_state != 0 && !self.states[fail_state].edges.contains_key(&byte) {
                    let f = self.states[fail_state].fail;
                    if leftmost && f == DEAD_STATE {
                        break;
                    }
                    fail_state = f as usize;
                }

                let computed_fail = self.states[fail_state]
                    .edges
                    .get(&byte)
                    .copied()
                    .unwrap_or(0);

                self.states[child as usize].fail =
                    if leftmost && self.states[child as usize].outpos.is_some() {
                        DEAD_STATE
                    } else {
                        computed_fail
                    };

                // Build output chains for overlapping mode
                if !leftmost {
                    let fail_outpos = if computed_fail == 0 {
                        u32::MAX
                    } else {
                        self.states[computed_fail as usize]
                            .outpos
                            .unwrap_or(u32::MAX)
                    };

                    if let Some(outpos) = self.states[child as usize].outpos {
                        self.outputs[outpos as usize].parent = fail_outpos;
                    } else if fail_outpos != u32::MAX {
                        self.states[child as usize].outpos = Some(fail_outpos);
                    }
                }
            }
        }
    }

    /// Compiles this trie into a double-array Aho-Corasick automaton.
    ///
    /// The double-array representation provides O(1) state transitions and is
    /// more memory-efficient for large pattern sets.
    ///
    /// Note: This consumes the trie. Call after [`build()`](Self::build).
    pub fn compile(self) -> DoubleArrayAhoCorasick {
        DoubleArrayAhoCorasick::from_trie(self)
    }

    /// Removes outputs from states where an ancestor has an earlier-added pattern.
    /// This implements "first-added pattern wins" semantics.
    fn prune_outputs_for_leftmost_first(&mut self) {
        let mut stack: Vec<(usize, Option<u32>)> = vec![(0, None)];

        while let Some((state, min_ancestor_outpos)) = stack.pop() {
            let current_outpos = self.states[state].outpos;

            // Prune if an ancestor was added earlier (smaller outpos)
            if let (Some(ancestor), Some(current)) = (min_ancestor_outpos, current_outpos) {
                if ancestor < current {
                    self.states[state].outpos = None;
                }
            }

            // Track minimum outpos seen on path from root
            let new_min = match (min_ancestor_outpos, self.states[state].outpos) {
                (None, None) => None,
                (None, Some(c)) => Some(c),
                (Some(a), None) => Some(a),
                (Some(a), Some(c)) => Some(a.min(c)),
            };

            for &child in self.states[state].edges.values() {
                stack.push((child as usize, new_min));
            }
        }
    }

    /// Builds the automaton for WordPiece tokenization.
    ///
    /// This sets up failure links so that completed tokens transition to the
    /// continuation anchor (e.g., "##" state), enabling single-pass tokenization.
    ///
    /// # Arguments
    /// * `prefix` - The continuation prefix bytes (e.g., b"##")
    pub fn build_wordpiece(&mut self, prefix: &[u8]) {
        self.match_kind = MatchKind::WordPiece;

        // Step 1: Find the anchor state by traversing the prefix
        let mut anchor = 0u32;
        for &byte in prefix {
            match self.states[anchor as usize].edges.get(&byte) {
                Some(&next) => anchor = next,
                None => {
                    // Prefix not in trie - no continuation patterns exist
                    self.anchor = None;
                    // Fall back to standard leftmost-longest build
                    self.build_standard_failure_links();
                    return;
                }
            }
        }
        self.anchor = Some(anchor);

        // Step 2: Build standard failure links using BFS
        self.build_standard_failure_links();

        // Step 3: Override failure links for states with outputs
        // All completed tokens should fail to the anchor
        for state_id in 0..self.states.len() {
            if self.states[state_id].outpos.is_some() {
                self.states[state_id].fail = anchor;
            }
        }

        // Step 4: Anchor's failure should be root (for clean fallback)
        self.states[anchor as usize].fail = 0;
    }

    /// Builds standard failure links using BFS (used by build_wordpiece).
    fn build_standard_failure_links(&mut self) {
        let mut queue: VecDeque<usize> = VecDeque::new();

        // Initialize root's children with fail = 0
        let root_children: Vec<u32> = self.states[0].edges.values().copied().collect();
        for child in root_children {
            queue.push_back(child as usize);
            self.states[child as usize].fail = 0;
        }

        // BFS to compute failure links
        while let Some(state) = queue.pop_front() {
            let parent_fail = self.states[state].fail;
            let edges: Vec<(u8, u32)> = self.states[state]
                .edges
                .iter()
                .map(|(&k, &v)| (k, v))
                .collect();

            for (byte, child) in edges {
                queue.push_back(child as usize);

                // Follow failure links to find longest proper suffix
                let mut fail_state = parent_fail as usize;
                while fail_state != 0 && !self.states[fail_state].edges.contains_key(&byte) {
                    fail_state = self.states[fail_state].fail as usize;
                }

                let computed_fail = self.states[fail_state]
                    .edges
                    .get(&byte)
                    .copied()
                    .unwrap_or(0);

                self.states[child as usize].fail = computed_fail;
            }
        }
    }

    /// Returns an iterator over all matches in the given text.
    pub fn find_iter<'a>(&'a self, text: &'a [u8]) -> TrieFindIter<'a> {
        TrieFindIter::new(self, text)
    }

    /// Returns all matches in the given text as a vector.
    pub fn find(&self, text: &[u8]) -> Vec<Match> {
        self.find_iter(text).collect()
    }

    /// Returns the number of states in the trie.
    pub fn num_states(&self) -> usize {
        self.states.len()
    }
}

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

    #[test]
    fn test_new_creates_root_state() {
        let trie = Trie::new();
        assert_eq!(trie.states.len(), 1);
        assert_eq!(trie.states[0].edges.len(), 0);
        assert_eq!(trie.states[0].fail, 0);
    }

    #[test]
    fn test_add_single_pattern() {
        let mut trie = Trie::new();
        trie.add(b"he", 0);
        assert_eq!(trie.states.len(), 3);

        // Check the output values
        assert!(trie.states[2].outpos.is_some());
        let outpos = trie.states[2].outpos.unwrap();
        let output = &trie.outputs[outpos as usize];
        assert_eq!(output.pattern_id, 0);
        assert_eq!(output.length, 2);
    }

    #[test]
    fn test_build_fails() {
        let mut trie = Trie::new();
        trie.add(b"he", 0);
        trie.add(b"she", 1);
        trie.build(MatchKind::Overlapping);

        // Find state for "she" (traverse s -> h -> e)
        let s = trie.states[0].edges.get(&b's').copied().unwrap_or(0);
        let sh = trie.states[s as usize]
            .edges
            .get(&b'h')
            .copied()
            .unwrap_or(0);
        let she = trie.states[sh as usize]
            .edges
            .get(&b'e')
            .copied()
            .unwrap_or(0);

        // Find state for "he" (traverse h -> e)
        let h = trie.states[0].edges.get(&b'h').copied().unwrap_or(0);
        let he = trie.states[h as usize]
            .edges
            .get(&b'e')
            .copied()
            .unwrap_or(0);

        // Check she should fail to "he"
        assert_eq!(trie.states[she as usize].fail, he);
    }

    #[test]
    fn test_find_single_pattern() {
        let mut trie = Trie::new();
        trie.add(b"he", 0);
        trie.build(MatchKind::Overlapping);

        let matches = trie.find(b"she");
        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0].pattern_id, 0);
        assert_eq!(matches[0].start, 1);
        assert_eq!(matches[0].end, 3);
    }

    #[test]
    fn test_find_no_match() {
        let mut trie = Trie::new();
        trie.add(b"xyz", 0);
        trie.build(MatchKind::Overlapping);

        let matches = trie.find(b"abc");
        assert_eq!(matches.len(), 0);
    }

    #[test]
    fn test_find_multiple_patterns() {
        let mut trie = Trie::new();
        trie.add(b"he", 0);
        trie.add(b"she", 1);
        trie.add(b"hers", 2);
        trie.build(MatchKind::Overlapping);

        // "ushers" = u s h e r s
        //            0 1 2 3 4 5
        let matches = trie.find(b"ushers");

        // Should find: "she" at 1..4, "he" at 2..4, "hers" at 2..6
        assert_eq!(matches.len(), 3);

        // Collect matches for easier checking
        let match_tuples: Vec<(u32, usize, usize)> = matches
            .iter()
            .map(|m| (m.pattern_id, m.start, m.end))
            .collect();

        assert!(match_tuples.contains(&(1, 1, 4))); // "she"
        assert!(match_tuples.contains(&(0, 2, 4))); // "he"
        assert!(match_tuples.contains(&(2, 2, 6))); // "hers"
    }

    #[test]
    fn test_find_at_start() {
        let mut trie = Trie::new();
        trie.add(b"hello", 0);
        trie.build(MatchKind::Overlapping);

        let matches = trie.find(b"hello world");
        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0].start, 0);
        assert_eq!(matches[0].end, 5);
    }

    #[test]
    fn test_find_at_end() {
        let mut trie = Trie::new();
        trie.add(b"end", 0);
        trie.build(MatchKind::Overlapping);

        let matches = trie.find(b"the end");
        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0].start, 4);
        assert_eq!(matches[0].end, 7);
    }

    #[test]
    fn test_find_overlapping() {
        let mut trie = Trie::new();
        trie.add(b"a", 0);
        trie.add(b"aa", 1);
        trie.add(b"aaa", 2);
        trie.build(MatchKind::Overlapping);

        // "aaaa" should find many overlapping matches
        let matches = trie.find(b"aaaa");

        // "a" at positions 0, 1, 2, 3 (4 matches)
        // "aa" at positions 0, 1, 2 (3 matches)
        // "aaa" at positions 0, 1 (2 matches)
        // Total: 9 matches
        assert_eq!(matches.len(), 9);
    }

    #[test]
    fn test_find_empty_text() {
        let mut trie = Trie::new();
        trie.add(b"he", 0);
        trie.build(MatchKind::Overlapping);

        let matches = trie.find(b"");
        assert_eq!(matches.len(), 0);
    }

    // ============ LeftmostLongest tests ============

    #[test]
    fn test_leftmost_longest_build_dead_state() {
        let mut trie = Trie::new();
        trie.add(b"he", 0);
        trie.add(b"she", 1);
        trie.build(MatchKind::LeftmostLongest);

        // Find state for "he"
        let h = trie.states[0].edges.get(&b'h').copied().unwrap();
        let he = trie.states[h as usize].edges.get(&b'e').copied().unwrap();

        // "he" has output, so fail should be DEAD
        assert_eq!(trie.states[he as usize].fail, DEAD_STATE);
    }

    #[test]
    fn test_leftmost_longest_propagates_dead() {
        let mut trie = Trie::new();
        trie.add(b"he", 0);
        trie.add(b"hers", 1);
        trie.build(MatchKind::LeftmostLongest);

        // "he" -> "her" -> "hers"
        let h = trie.states[0].edges.get(&b'h').copied().unwrap();
        let he = trie.states[h as usize].edges.get(&b'e').copied().unwrap();
        let her = trie.states[he as usize].edges.get(&b'r').copied().unwrap();
        let hers = trie.states[her as usize].edges.get(&b's').copied().unwrap();

        // "he" has output -> DEAD
        assert_eq!(trie.states[he as usize].fail, DEAD_STATE);
        // descendants of DEAD also get DEAD
        assert_eq!(trie.states[her as usize].fail, DEAD_STATE);
        assert_eq!(trie.states[hers as usize].fail, DEAD_STATE);
    }

    #[test]
    fn test_leftmost_longest_find_non_overlapping() {
        let mut trie = Trie::new();
        trie.add(b"a", 0);
        trie.add(b"aa", 1);
        trie.add(b"aaa", 2);
        trie.build(MatchKind::LeftmostLongest);

        // LeftmostLongest prefers the longest match at each position
        let matches = trie.find(b"aaaa");

        // Should find "aaa" at 0..3, then "a" at 3..4
        assert_eq!(matches.len(), 2);
        assert_eq!(matches[0].pattern_id, 2); // "aaa"
        assert_eq!(matches[0].start, 0);
        assert_eq!(matches[0].end, 3);
        assert_eq!(matches[1].pattern_id, 0); // "a"
        assert_eq!(matches[1].start, 3);
        assert_eq!(matches[1].end, 4);
    }

    #[test]
    fn test_leftmost_longest_different_add_order() {
        let mut trie = Trie::new();
        // Add longer pattern first
        trie.add(b"aaa", 0);
        trie.add(b"aa", 1);
        trie.add(b"a", 2);
        trie.build(MatchKind::LeftmostLongest);

        // With different add order, still finds longest match first
        let matches = trie.find(b"aaaa");

        // Should find "aaa" at 0..3, then "a" at 3..4
        assert_eq!(matches.len(), 2);
        assert_eq!(matches[0].pattern_id, 0); // "aaa"
        assert_eq!(matches[0].start, 0);
        assert_eq!(matches[0].end, 3);
        assert_eq!(matches[1].pattern_id, 2); // "a"
        assert_eq!(matches[1].start, 3);
        assert_eq!(matches[1].end, 4);
    }

    #[test]
    fn test_leftmost_longest_find_she_he() {
        let mut trie = Trie::new();
        trie.add(b"he", 0);
        trie.add(b"she", 1);
        trie.build(MatchKind::LeftmostLongest);

        // "she" should only match "she", not "he" (non-overlapping)
        let matches = trie.find(b"she");
        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0].pattern_id, 1); // "she"
        assert_eq!(matches[0].start, 0);
        assert_eq!(matches[0].end, 3);
    }

    #[test]
    fn test_leftmost_longest_find_ushers() {
        let mut trie = Trie::new();
        trie.add(b"he", 0);
        trie.add(b"she", 1);
        trie.add(b"hers", 2);
        trie.build(MatchKind::LeftmostLongest);

        // "ushers" with leftmost-longest
        let matches = trie.find(b"ushers");

        // Should find "she" at 1..4, then nothing overlapping with it
        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0].pattern_id, 1); // "she"
    }

    // ============ LeftmostFirst tests ============

    #[test]
    fn test_leftmost_first_prunes_outputs() {
        let mut trie = Trie::new();
        // Add shorter pattern first - it should block longer ones
        trie.add(b"a", 0);
        trie.add(b"aa", 1);
        trie.add(b"aaa", 2);
        trie.build(MatchKind::LeftmostFirst);

        // Only "a" should have an output (first-added blocks extensions)
        let a = trie.states[0].edges.get(&b'a').copied().unwrap();
        assert!(trie.states[a as usize].outpos.is_some()); // "a" has output

        let aa = trie.states[a as usize].edges.get(&b'a').copied().unwrap();
        assert!(trie.states[aa as usize].outpos.is_none()); // "aa" pruned

        let aaa = trie.states[aa as usize].edges.get(&b'a').copied().unwrap();
        assert!(trie.states[aaa as usize].outpos.is_none()); // "aaa" pruned
    }

    #[test]
    fn test_leftmost_first_short_pattern_wins() {
        let mut trie = Trie::new();
        // Add shorter pattern first
        trie.add(b"a", 0);
        trie.add(b"aa", 1);
        trie.add(b"aaa", 2);
        trie.build(MatchKind::LeftmostFirst);

        // LeftmostFirst: "a" was added first, so it wins at every position
        let matches = trie.find(b"aaaa");

        // Should find "a" at positions 0, 1, 2, 3 (4 matches)
        assert_eq!(matches.len(), 4);
        for (i, m) in matches.iter().enumerate() {
            assert_eq!(m.pattern_id, 0); // "a"
            assert_eq!(m.start, i);
            assert_eq!(m.end, i + 1);
        }
    }

    #[test]
    fn test_leftmost_first_long_pattern_first() {
        let mut trie = Trie::new();
        // Add longer pattern first - no pruning occurs
        trie.add(b"aaa", 0);
        trie.add(b"aa", 1);
        trie.add(b"a", 2);
        trie.build(MatchKind::LeftmostFirst);

        // All outputs should exist (no prefix relationship blocks them)
        let a = trie.states[0].edges.get(&b'a').copied().unwrap();
        assert!(trie.states[a as usize].outpos.is_some()); // "a" has output

        let aa = trie.states[a as usize].edges.get(&b'a').copied().unwrap();
        assert!(trie.states[aa as usize].outpos.is_some()); // "aa" has output

        let aaa = trie.states[aa as usize].edges.get(&b'a').copied().unwrap();
        assert!(trie.states[aaa as usize].outpos.is_some()); // "aaa" has output
    }

    #[test]
    fn test_leftmost_first_long_pattern_first_finds_longest() {
        let mut trie = Trie::new();
        // Add longer pattern first
        trie.add(b"aaa", 0);
        trie.add(b"aa", 1);
        trie.add(b"a", 2);
        trie.build(MatchKind::LeftmostFirst);

        // When longer patterns are added first, they're not pruned
        // Iterator still uses greedy/longest logic, so finds longest
        let matches = trie.find(b"aaaa");

        // Should find "aaa" at 0..3, then "a" at 3..4
        assert_eq!(matches.len(), 2);
        assert_eq!(matches[0].pattern_id, 0); // "aaa"
        assert_eq!(matches[0].start, 0);
        assert_eq!(matches[0].end, 3);
        assert_eq!(matches[1].pattern_id, 2); // "a"
        assert_eq!(matches[1].start, 3);
        assert_eq!(matches[1].end, 4);
    }

    #[test]
    fn test_leftmost_first_non_prefix_patterns() {
        let mut trie = Trie::new();
        // These patterns don't share prefixes, so none get pruned
        trie.add(b"ab", 0);
        trie.add(b"ac", 1);
        trie.build(MatchKind::LeftmostFirst);

        // Both should have outputs
        let a = trie.states[0].edges.get(&b'a').copied().unwrap();
        let ab = trie.states[a as usize].edges.get(&b'b').copied().unwrap();
        let ac = trie.states[a as usize].edges.get(&b'c').copied().unwrap();

        assert!(trie.states[ab as usize].outpos.is_some()); // "ab" has output
        assert!(trie.states[ac as usize].outpos.is_some()); // "ac" has output
    }

    #[test]
    fn test_leftmost_first_stores_match_kind() {
        let mut trie = Trie::new();
        trie.add(b"test", 0);

        // Before build, default is Overlapping
        assert!(matches!(trie.match_kind, MatchKind::Overlapping));

        trie.build(MatchKind::LeftmostFirst);

        // After build with LeftmostFirst, it should be stored
        assert!(matches!(trie.match_kind, MatchKind::LeftmostFirst));
    }

    #[test]
    fn test_leftmost_longest_stores_match_kind() {
        let mut trie = Trie::new();
        trie.add(b"test", 0);
        trie.build(MatchKind::LeftmostLongest);

        assert!(matches!(trie.match_kind, MatchKind::LeftmostLongest));
    }

    // ============ WordPiece tests ============

    #[test]
    fn test_wordpiece_finds_anchor() {
        let mut trie = Trie::new();
        // Add word-start patterns
        trie.add(b"un", 0);
        trie.add(b"break", 1);
        // Add continuation patterns (## prefix)
        trie.add(b"##break", 2);
        trie.add(b"##able", 3);

        trie.build_wordpiece(b"##");

        // Should have found the anchor at "##"
        assert!(trie.anchor.is_some());
        let anchor = trie.anchor.unwrap();

        // Verify anchor is reachable via ## from root
        let hash1 = trie.states[0].edges.get(&b'#').copied().unwrap();
        let hash2 = trie.states[hash1 as usize].edges.get(&b'#').copied().unwrap();
        assert_eq!(anchor, hash2);
    }

    #[test]
    fn test_wordpiece_failure_links_point_to_anchor() {
        let mut trie = Trie::new();
        trie.add(b"un", 0);
        trie.add(b"break", 1);
        trie.add(b"##break", 2);
        trie.add(b"##able", 3);

        trie.build_wordpiece(b"##");

        let anchor = trie.anchor.unwrap();

        // Find "un" state and verify its fail points to anchor
        let u = trie.states[0].edges.get(&b'u').copied().unwrap();
        let un = trie.states[u as usize].edges.get(&b'n').copied().unwrap();
        assert!(trie.states[un as usize].outpos.is_some()); // has output
        assert_eq!(trie.states[un as usize].fail, anchor);

        // Find "break" state and verify its fail points to anchor
        let b = trie.states[0].edges.get(&b'b').copied().unwrap();
        let br = trie.states[b as usize].edges.get(&b'r').copied().unwrap();
        let bre = trie.states[br as usize].edges.get(&b'e').copied().unwrap();
        let brea = trie.states[bre as usize].edges.get(&b'a').copied().unwrap();
        let break_state = trie.states[brea as usize].edges.get(&b'k').copied().unwrap();
        assert!(trie.states[break_state as usize].outpos.is_some());
        assert_eq!(trie.states[break_state as usize].fail, anchor);

        // Find "##able" state and verify its fail points to anchor
        let hash1 = trie.states[0].edges.get(&b'#').copied().unwrap();
        let hash2 = trie.states[hash1 as usize].edges.get(&b'#').copied().unwrap();
        let a = trie.states[hash2 as usize].edges.get(&b'a').copied().unwrap();
        let ab = trie.states[a as usize].edges.get(&b'b').copied().unwrap();
        let abl = trie.states[ab as usize].edges.get(&b'l').copied().unwrap();
        let able = trie.states[abl as usize].edges.get(&b'e').copied().unwrap();
        assert!(trie.states[able as usize].outpos.is_some());
        assert_eq!(trie.states[able as usize].fail, anchor);
    }

    #[test]
    fn test_wordpiece_stores_match_kind() {
        let mut trie = Trie::new();
        trie.add(b"test", 0);
        trie.add(b"##ing", 1);

        trie.build_wordpiece(b"##");

        assert!(matches!(trie.match_kind, MatchKind::WordPiece));
    }

    #[test]
    fn test_wordpiece_no_anchor_prefix() {
        // When ## prefix doesn't exist in vocabulary
        let mut trie = Trie::new();
        trie.add(b"hello", 0);
        trie.add(b"world", 1);

        trie.build_wordpiece(b"##");

        // Anchor should be None
        assert!(trie.anchor.is_none());
    }

    #[test]
    fn test_wordpiece_find_uses_leftmost_longest() {
        let mut trie = Trie::new();
        trie.add(b"a", 0);
        trie.add(b"ab", 1);
        trie.add(b"abc", 2);
        trie.add(b"##d", 3);

        trie.build_wordpiece(b"##");

        // WordPiece find_iter should use leftmost-longest semantics
        let matches = trie.find(b"abc");
        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0].pattern_id, 2); // "abc" (longest)
        assert_eq!(matches[0].start, 0);
        assert_eq!(matches[0].end, 3);
    }
}