daggrs 0.1.0

mod allocator;
mod iter;

pub use iter::{FindIter, OutputIter};

use alloc::vec::Vec;
use allocator::Allocator;

use crate::trie::Trie;
use crate::types::{Match, MatchKind, Output};

/// Block size for double-array allocation.
pub(crate) const BLOCK_LEN: u32 = 256;
/// Index of the root state in the double-array.
pub(crate) const ROOT_IDX: u32 = 1;

/// A state in the double-array structure.
#[derive(Debug, Clone, Copy)]
pub struct State {
    /// Base value for computing child indices: `child_idx = base ^ label`.
    pub base: u32,
    /// Parent state index for validation: transition is valid if `states[child].check == parent`.
    pub check: u32,
    /// Failure link for Aho-Corasick matching.
    pub fail: u32,
    /// Output index if this state completes a pattern, `u32::MAX` otherwise.
    pub outpos: u32,
}

/// A double-array Aho-Corasick automaton for fast multi-pattern matching.
///
/// This is a compact representation of the trie that provides O(1) state transitions.
/// Build from a [`Trie`] using [`from_trie()`](Self::from_trie) or [`Trie::compile()`].
#[derive(Debug, Clone)]
pub struct DoubleArrayAhoCorasick {
    pub states: Vec<State>,
    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 DoubleArrayAhoCorasick {
    fn default() -> Self {
        Self::new()
    }
}

impl DoubleArrayAhoCorasick {
    /// Creates a new empty double-array with dead and root states.
    pub fn new() -> Self {
        let dead = State {
            base: 0,
            check: u32::MAX,
            fail: 0,
            outpos: u32::MAX,
        };
        let root = State {
            base: 0,
            check: u32::MAX,
            fail: ROOT_IDX,
            outpos: u32::MAX,
        };
        Self {
            states: vec![dead, root],
            outputs: Vec::new(),
            match_kind: MatchKind::default(),
            anchor: None,
        }
    }

    /// Checks if an index is available for use.
    fn is_vacant(&self, idx: u32) -> bool {
        idx as usize >= self.states.len() || self.states[idx as usize].check == 0
    }

    /// Checks if a base value works for all given labels.
    fn is_valid_base(&self, base: u32, labels: &[u8]) -> bool {
        labels.iter().all(|&c| self.is_vacant(base ^ (c as u32)))
    }

    /// Finds a valid base value for the given labels.
    fn find_base(&mut self, labels: &[u8], allocator: &mut Allocator) -> u32 {
        for idx in allocator.iter() {
            let base = idx ^ (labels[0] as u32);
            if self.is_valid_base(base, labels) {
                return base;
            }
        }

        // Extend array with another block
        let old_size = self.states.len() as u32;
        let num_blocks = (old_size + BLOCK_LEN - 1) / BLOCK_LEN;
        let new_block_end = (num_blocks + 1) * BLOCK_LEN;

        self.states.resize(
            new_block_end as usize,
            State {
                base: 0,
                check: 0,
                fail: 0,
                outpos: u32::MAX,
            },
        );
        allocator.extend(old_size, new_block_end);

        for idx in allocator.iter() {
            let base = idx ^ (labels[0] as u32);
            if self.is_valid_base(base, labels) {
                return base;
            }
        }

        panic!("Could not find valid base after extending");
    }

    /// Recursively builds the double-array from a trie using DFS.
    fn build_recursive(
        &mut self,
        trie: &Trie,
        trie_state: usize,
        daac_state: u32,
        mapping: &mut Vec<u32>,
        allocator: &mut Allocator,
    ) {
        let edges: Vec<(u8, u32)> = trie.states[trie_state]
            .edges
            .iter()
            .map(|(&label, &child)| (label, child))
            .collect();

        if edges.is_empty() {
            return;
        }

        let labels: Vec<u8> = edges.iter().map(|(l, _)| *l).collect();
        let base = self.find_base(&labels, allocator);
        self.states[daac_state as usize].base = base;

        // First pass: allocate all children to prevent sibling conflicts
        for &(label, trie_child) in &edges {
            let daac_child = base ^ (label as u32);

            if daac_child as usize >= self.states.len() {
                let old_size = self.states.len() as u32;
                self.states.resize(
                    daac_child as usize + 1,
                    State {
                        base: 0,
                        check: 0,
                        fail: 0,
                        outpos: u32::MAX,
                    },
                );
                allocator.extend(old_size, self.states.len() as u32);
            }

            self.states[daac_child as usize].check = daac_state;
            allocator.delete(daac_child);
            mapping[trie_child as usize] = daac_child;
        }

        // Second pass: recurse into children
        for (label, trie_child) in edges {
            let daac_child = base ^ (label as u32);
            self.build_recursive(trie, trie_child as usize, daac_child, mapping, allocator);
        }
    }

    /// Converts a trie into a double-array automaton.
    pub fn from_trie(trie: Trie) -> Self {
        use crate::trie::DEAD_STATE;

        let mut daac = Self::new();
        daac.match_kind = trie.match_kind;

        let mut allocator = Allocator::new(daac.states.len());
        let mut mapping = vec![0u32; trie.num_states()];
        mapping[0] = ROOT_IDX;

        daac.build_recursive(&trie, 0, ROOT_IDX, &mut mapping, &mut allocator);

        // Map failure links and outputs from trie to double-array indices
        for trie_id in 0..trie.num_states() {
            let daac_id = mapping[trie_id];
            let trie_fail = trie.states[trie_id].fail;

            daac.states[daac_id as usize].fail = if trie_fail == DEAD_STATE {
                DEAD_STATE
            } else {
                mapping[trie_fail as usize]
            };
            daac.states[daac_id as usize].outpos = trie.states[trie_id].outpos.unwrap_or(u32::MAX);
        }

        // Map anchor state for WordPiece
        daac.anchor = trie.anchor.map(|a| mapping[a as usize]);

        daac.outputs = trie.outputs;
        daac
    }

    /// Returns an iterator over all matches in the given text.
    pub fn find_iter<'a>(&'a self, text: &'a [u8]) -> FindIter<'a> {
        FindIter::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 start state for manual state machine traversal.
    ///
    /// Use with [`next_state`](Self::next_state) and [`outputs`](Self::outputs)
    /// for fine-grained control over pattern matching, e.g., for BPE tokenization.
    #[inline]
    pub fn start_state(&self) -> u32 {
        ROOT_IDX
    }

    /// Transitions to the next state given the current state and input byte.
    ///
    /// Follows failure links as needed until a valid transition is found
    /// or the root state is reached.
    ///
    /// # Arguments
    /// * `state` - The current state (from `start_state()` or previous `next_state()`)
    /// * `byte` - The input byte to process
    ///
    /// # Returns
    /// The next state after processing the byte.
    #[inline]
    pub fn next_state(&self, mut state: u32, byte: u8) -> u32 {
        use crate::trie::DEAD_STATE;

        let states = &self.states;
        let states_len = states.len();

        loop {
            let current = &states[state as usize];
            let child = current.base ^ (byte as u32);

            // Check if transition is valid
            if (child as usize) < states_len && states[child as usize].check == state {
                return child;
            }

            // No valid transition, follow failure link
            if state == ROOT_IDX {
                return ROOT_IDX;
            }

            let fail = current.fail;
            if fail == DEAD_STATE {
                return ROOT_IDX;
            }
            state = fail;
        }
    }

    /// Returns an iterator over all outputs (pattern matches) at the given state.
    ///
    /// This walks the output chain, yielding each pattern that matches when
    /// the automaton reaches this state. Useful for BPE tokenization where
    /// you need all tokens ending at each position.
    ///
    /// # Arguments
    /// * `state` - A state from `start_state()` or `next_state()`
    ///
    /// # Returns
    /// An iterator yielding `&Output` for each matching pattern.
    #[inline]
    pub fn outputs(&self, state: u32) -> iter::OutputIter<'_> {
        let outpos = self.states[state as usize].outpos;
        iter::OutputIter::new(&self.outputs, outpos)
    }

    /// Consumes a byte and returns the new state along with an iterator over outputs.
    ///
    /// This is a combined operation that's more efficient than calling `next_state()`
    /// followed by `outputs()` separately, as it avoids redundant state lookups.
    ///
    /// # Arguments
    /// * `state` - Current state
    /// * `byte` - Input byte to consume
    ///
    /// # Returns
    /// A tuple of (new_state, output_iterator)
    #[inline]
    pub fn consume(&self, state: u32, byte: u8) -> (u32, iter::OutputIter<'_>) {
        let next = self.next_state(state, byte);
        let outpos = self.states[next as usize].outpos;
        (next, iter::OutputIter::new(&self.outputs, outpos))
    }

    /// Serializes the automaton to bytes.
    ///
    /// Format:
    /// - num_states: u32
    /// - states: [State; num_states] (each 16 bytes)
    /// - num_outputs: u32
    /// - outputs: [Output; num_outputs] (each 12 bytes)
    /// - match_kind: u8
    /// - anchor: u32 (u32::MAX if None)
    pub fn serialize(&self) -> Vec<u8> {
        use core::mem::size_of;

        let state_bytes = self.states.len() * size_of::<State>();
        let output_bytes = self.outputs.len() * size_of::<Output>();
        let total = 4 + state_bytes + 4 + output_bytes + 1 + 4;

        let mut buf = Vec::with_capacity(total);

        // num_states + states
        buf.extend_from_slice(&(self.states.len() as u32).to_le_bytes());
        for state in &self.states {
            buf.extend_from_slice(&state.base.to_le_bytes());
            buf.extend_from_slice(&state.check.to_le_bytes());
            buf.extend_from_slice(&state.fail.to_le_bytes());
            buf.extend_from_slice(&state.outpos.to_le_bytes());
        }

        // num_outputs + outputs
        buf.extend_from_slice(&(self.outputs.len() as u32).to_le_bytes());
        for output in &self.outputs {
            buf.extend_from_slice(&output.pattern_id.to_le_bytes());
            buf.extend_from_slice(&output.length.to_le_bytes());
            buf.extend_from_slice(&output.parent.to_le_bytes());
        }

        // match_kind
        buf.push(match self.match_kind {
            MatchKind::Overlapping => 0,
            MatchKind::LeftmostFirst => 1,
            MatchKind::LeftmostLongest => 2,
            MatchKind::WordPiece => 3,
        });

        // anchor (u32::MAX = None)
        buf.extend_from_slice(&self.anchor.unwrap_or(u32::MAX).to_le_bytes());

        buf
    }

    /// Deserializes an automaton from bytes.
    ///
    /// Returns the automaton and the remaining unused bytes.
    /// Note: Reads data manually to handle potentially unaligned slices from file reads.
    pub fn deserialize(data: &[u8]) -> Option<(Self, &[u8])> {
        use core::mem::size_of;

        if data.len() < 4 {
            return None;
        }

        let mut pos = 0;

        // num_states
        let num_states = u32::from_le_bytes(data[pos..pos + 4].try_into().ok()?) as usize;
        pos += 4;

        // states - read manually for alignment safety
        let state_bytes = num_states * size_of::<State>();
        if data.len() < pos + state_bytes {
            return None;
        }
        let mut states = Vec::with_capacity(num_states);
        for i in 0..num_states {
            let start = pos + i * size_of::<State>();
            states.push(State {
                base: u32::from_le_bytes(data[start..start + 4].try_into().ok()?),
                check: u32::from_le_bytes(data[start + 4..start + 8].try_into().ok()?),
                fail: u32::from_le_bytes(data[start + 8..start + 12].try_into().ok()?),
                outpos: u32::from_le_bytes(data[start + 12..start + 16].try_into().ok()?),
            });
        }
        pos += state_bytes;

        // num_outputs
        if data.len() < pos + 4 {
            return None;
        }
        let num_outputs = u32::from_le_bytes(data[pos..pos + 4].try_into().ok()?) as usize;
        pos += 4;

        // outputs - read manually for alignment safety
        let output_bytes = num_outputs * size_of::<Output>();
        if data.len() < pos + output_bytes {
            return None;
        }
        let mut outputs = Vec::with_capacity(num_outputs);
        for i in 0..num_outputs {
            let start = pos + i * size_of::<Output>();
            outputs.push(Output {
                pattern_id: u32::from_le_bytes(data[start..start + 4].try_into().ok()?),
                length: u32::from_le_bytes(data[start + 4..start + 8].try_into().ok()?),
                parent: u32::from_le_bytes(data[start + 8..start + 12].try_into().ok()?),
            });
        }
        pos += output_bytes;

        // match_kind
        if data.len() < pos + 1 {
            return None;
        }
        let match_kind = match data[pos] {
            0 => MatchKind::Overlapping,
            1 => MatchKind::LeftmostFirst,
            2 => MatchKind::LeftmostLongest,
            3 => MatchKind::WordPiece,
            _ => return None,
        };
        pos += 1;

        // anchor
        if data.len() < pos + 4 {
            return None;
        }
        let anchor_val = u32::from_le_bytes(data[pos..pos + 4].try_into().ok()?);
        let anchor = if anchor_val == u32::MAX { None } else { Some(anchor_val) };
        pos += 4;

        Some((
            Self { states, outputs, match_kind, anchor },
            &data[pos..],
        ))
    }
}

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

    #[test]
    fn test_new_creates_dead_and_root_states() {
        let daac = DoubleArrayAhoCorasick::new();
        assert_eq!(daac.states.len(), 2);

        // Root state at index 1
        assert_eq!(daac.states[ROOT_IDX as usize].check, u32::MAX);
        assert_eq!(daac.states[ROOT_IDX as usize].fail, ROOT_IDX);
    }

    #[test]
    fn test_is_vacant() {
        let daac = DoubleArrayAhoCorasick::new();

        // Index beyond array is vacant
        assert!(daac.is_vacant(100));

        // Root state is not vacant (check != 0)
        assert!(!daac.is_vacant(ROOT_IDX));
    }

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

        let daac = DoubleArrayAhoCorasick::from_trie(trie);

        // Should have outputs
        assert_eq!(daac.outputs.len(), 1);
        assert_eq!(daac.outputs[0].pattern_id, 0);
        assert_eq!(daac.outputs[0].length, 2);

        // Should find "he" in text
        let matches = daac.find(b"she");
        assert_eq!(matches.len(), 1);
        assert_eq!(matches[0].pattern_id, 0);
    }

    #[test]
    fn test_from_trie_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);

        let daac = DoubleArrayAhoCorasick::from_trie(trie);

        // Should have 3 outputs
        assert_eq!(daac.outputs.len(), 3);

        // Verify all patterns are found in appropriate text
        let matches = daac.find(b"ushers");
        assert_eq!(matches.len(), 3);

        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_no_match() {
        let mut trie = Trie::new();
        trie.add(b"he", 0);
        trie.build(MatchKind::Overlapping);

        let daac = DoubleArrayAhoCorasick::from_trie(trie);

        // No matches for text without pattern
        let matches = daac.find(b"xyz");
        assert_eq!(matches.len(), 0);
    }

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

        let daac = DoubleArrayAhoCorasick::from_trie(trie);

        // "she" should match both "she" and "he" via fail links
        let matches = daac.find(b"she");
        assert_eq!(matches.len(), 2);

        let match_ids: Vec<u32> = matches.iter().map(|m| m.pattern_id).collect();
        assert!(match_ids.contains(&0)); // "he"
        assert!(match_ids.contains(&1)); // "she"
    }

    #[test]
    fn test_empty_trie() {
        let trie = Trie::new();
        let daac = DoubleArrayAhoCorasick::from_trie(trie);

        // Should still have dead and root states
        assert!(daac.states.len() >= 2);
        assert_eq!(daac.outputs.len(), 0);
    }

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

        let daac = DoubleArrayAhoCorasick::from_trie(trie);

        // All single-byte patterns should be found
        let matches = daac.find(b"abc");
        assert_eq!(matches.len(), 3);

        let match_ids: Vec<u32> = matches.iter().map(|m| m.pattern_id).collect();
        assert!(match_ids.contains(&0)); // "a"
        assert!(match_ids.contains(&1)); // "b"
        assert!(match_ids.contains(&2)); // "c"
    }

    #[test]
    fn test_start_state() {
        let daac = DoubleArrayAhoCorasick::new();
        assert_eq!(daac.start_state(), ROOT_IDX);
    }

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

        let daac = DoubleArrayAhoCorasick::from_trie(trie);

        // Manually walk through "she"
        let mut state = daac.start_state();

        // Process 's'
        state = daac.next_state(state, b's');
        let outputs: Vec<u32> = daac.outputs(state).map(|o| o.pattern_id).collect();
        assert!(outputs.is_empty()); // No match at 's'

        // Process 'h'
        state = daac.next_state(state, b'h');
        let outputs: Vec<u32> = daac.outputs(state).map(|o| o.pattern_id).collect();
        assert!(outputs.is_empty()); // No match at 'sh'

        // Process 'e'
        state = daac.next_state(state, b'e');
        let outputs: Vec<u32> = daac.outputs(state).map(|o| o.pattern_id).collect();
        // Should have both "she" and "he" (via output chain)
        assert_eq!(outputs.len(), 2);
        assert!(outputs.contains(&0)); // "he"
        assert!(outputs.contains(&1)); // "she"
    }

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

        let daac = DoubleArrayAhoCorasick::from_trie(trie);

        // At start state, no outputs
        let state = daac.start_state();
        let outputs: Vec<_> = daac.outputs(state).collect();
        assert!(outputs.is_empty());
    }

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

        let daac = DoubleArrayAhoCorasick::from_trie(trie);
        let text = b"abc";

        // Collect matches using manual traversal
        let mut manual_matches = Vec::new();
        let mut state = daac.start_state();

        for (pos, &byte) in text.iter().enumerate() {
            state = daac.next_state(state, byte);
            for output in daac.outputs(state) {
                manual_matches.push((output.pattern_id, pos + 1 - output.length as usize, pos + 1));
            }
        }

        // Collect matches using find()
        let find_matches: Vec<_> = daac
            .find(text)
            .iter()
            .map(|m| (m.pattern_id, m.start, m.end))
            .collect();

        // Should be the same
        assert_eq!(manual_matches.len(), find_matches.len());
        for m in &manual_matches {
            assert!(find_matches.contains(m));
        }
    }

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

        let daac = DoubleArrayAhoCorasick::from_trie(trie);

        // Serialize
        let bytes = daac.serialize();

        // Deserialize
        let (restored, remainder) = DoubleArrayAhoCorasick::deserialize(&bytes).unwrap();
        assert!(remainder.is_empty());

        // Verify same state
        assert_eq!(daac.states.len(), restored.states.len());
        assert_eq!(daac.outputs.len(), restored.outputs.len());
        assert_eq!(daac.match_kind, restored.match_kind);

        // Verify same behavior
        let original_matches = daac.find(b"ushers");
        let restored_matches = restored.find(b"ushers");
        assert_eq!(original_matches, restored_matches);
    }

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

        let daac = DoubleArrayAhoCorasick::from_trie(trie);

        let bytes = daac.serialize();
        let (restored, _) = DoubleArrayAhoCorasick::deserialize(&bytes).unwrap();

        assert_eq!(restored.match_kind, MatchKind::LeftmostFirst);

        let original_matches = daac.find(b"aaaa");
        let restored_matches = restored.find(b"aaaa");
        assert_eq!(original_matches, restored_matches);
    }

    #[test]
    fn test_serialize_deserialize_empty() {
        let daac = DoubleArrayAhoCorasick::new();

        let bytes = daac.serialize();
        let (restored, _) = DoubleArrayAhoCorasick::deserialize(&bytes).unwrap();

        assert_eq!(daac.states.len(), restored.states.len());
        assert_eq!(daac.outputs.len(), restored.outputs.len());
    }

    #[test]
    fn test_deserialize_invalid_data() {
        // Too short
        assert!(DoubleArrayAhoCorasick::deserialize(&[0, 1]).is_none());

        // Invalid match_kind (match_kind is 5 bytes before end: 1 byte match_kind + 4 bytes anchor)
        let mut trie = Trie::new();
        trie.add(b"a", 0);
        trie.build(MatchKind::Overlapping);
        let daac = DoubleArrayAhoCorasick::from_trie(trie);
        let mut bytes = daac.serialize();
        let match_kind_idx = bytes.len() - 5;
        bytes[match_kind_idx] = 99; // Invalid match_kind
        assert!(DoubleArrayAhoCorasick::deserialize(&bytes).is_none());
    }
}