Aho-Corasick の Leftmost-Longest の償却O(N)

leftmost-longest.rs

use std::collections::HashMap;

#[derive(Debug, Clone, Copy)]
#[repr(C)]
pub struct Match {
    pub start: u32,
    pub end: u32,
    pub pattern: u32,
}

/// Output entry stored in a separate array for compact states.
#[derive(Clone, Copy)]
struct Output {
    byte_len: u32,
    pattern: u32,
}

/// A single state in the Double Array.
///
/// Layout (12 bytes):
///   base:       u32  — XOR base (0 = no children)
///   fail:       u32  — failure link state index
///   check_opos: u32  — packed: check label (8 bits) + output_pos (24 bits)
///
/// XOR addressing: child_idx = base ^ c
/// Transition is valid iff states[base ^ c].check() == c (label byte).
/// False transitions are prevented by ensuring unique base values
/// (used_base_bmp): if base1 == base2 is forbidden, then
/// base1 ^ c == base2 ^ c2 with c == c2 implies base1 == base2, contradiction.
#[derive(Clone, Copy)]
#[repr(C)]
struct State {
    base: u32,
    fail: u32,
    check_opos: u32, // bits [31:24] = check (label), bits [23:0] = output_pos
}

impl State {
    const EMPTY_CHECK: u8 = 0xFF;

    #[inline(always)]
    fn check(&self) -> u8 {
        (self.check_opos >> 24) as u8
    }

    #[inline(always)]
    fn output_pos(&self) -> u32 {
        self.check_opos & 0x00FF_FFFF
    }

    #[inline(always)]
    fn set_check(&mut self, c: u8) {
        self.check_opos = (self.check_opos & 0x00FF_FFFF) | ((c as u32) << 24);
    }

    #[inline(always)]
    fn set_output_pos(&mut self, opos: u32) {
        debug_assert!(opos <= 0x00FF_FFFF);
        self.check_opos = (self.check_opos & 0xFF00_0000) | opos;
    }
}

impl Default for State {
    fn default() -> Self {
        State {
            base: 0,
            fail: 0,
            check_opos: (Self::EMPTY_CHECK as u32) << 24, // check=0xFF, opos=0
        }
    }
}

const ROOT_IDX: u32 = 0;
const NONE: u32 = u32::MAX;

/// Double Array Aho-Corasick with Greedy (Leftmost-Longest) matching.
///
/// Uses XOR-based Double Array addressing (child = base ^ c) with 12-byte
/// states. Outputs are stored in a separate array.
///
/// check stores the label byte (u8), packed with output_pos (24 bits).
/// False transitions are prevented by unique base values (used_base_bmp).
///
/// Greedy optimization: keyword nodes' failure links point to root,
/// skipping suffix matches to guarantee leftmost-longest semantics.
pub struct DaGreedyAhoCorasick {
    states: Vec<State>,
    outputs: Vec<Output>,
}

impl DaGreedyAhoCorasick {
    pub fn new(keywords: &[&str]) -> Self {
        // --- Phase 1: Build temporary Trie ---
        struct TrieNode {
            children: HashMap<u8, usize>,
            output_idx: u32,
        }
        let mut trie: Vec<TrieNode> = vec![TrieNode {
            children: HashMap::new(),
            output_idx: 0,
        }];

        let mut outputs: Vec<Output> = vec![Output { byte_len: 0, pattern: 0 }];

        let mut unique_keywords: Vec<&str> = Vec::new();
        {
            let mut seen = std::collections::HashSet::new();
            for &kw in keywords {
                if seen.insert(kw) {
                    unique_keywords.push(kw);
                }
            }
        }

        for (pat_idx, &kw) in unique_keywords.iter().enumerate() {
            let mut node = 0usize;
            for &b in kw.as_bytes() {
                if !trie[node].children.contains_key(&b) {
                    let new_id = trie.len();
                    trie.push(TrieNode {
                        children: HashMap::new(),
                        output_idx: 0,
                    });
                    trie[node].children.insert(b, new_id);
                }
                node = trie[node].children[&b];
            }
            let out_idx = outputs.len() as u32;
            assert!(out_idx <= 0x00FF_FFFF, "Too many outputs for 24-bit output_pos");
            outputs.push(Output {
                byte_len: kw.len() as u32,
                pattern: pat_idx as u32,
            });
            trie[node].output_idx = out_idx;
        }

        // --- Phase 2: Build Double Array with XOR addressing ---
        let init_size = (trie.len() + 256).max(512);
        let mut states: Vec<State> = vec![State::default(); init_size];

        // u64 bitmap for O(1) occupancy check
        let bmp_words = (init_size + 63) / 64;
        let mut used_bmp: Vec<u64> = vec![0u64; bmp_words];

        // Doubly-linked free list (separate arrays)
        let mut free_next: Vec<u32> = vec![NONE; init_size];
        let mut free_prev: Vec<u32> = vec![NONE; init_size];
        let mut free_head: u32 = 2; // 0=root, 1=reserved
        let mut free_tail: u32 = (init_size - 1) as u32;

        // Initialize free list: 2..init_size
        for i in 2..init_size {
            free_next[i] = if i + 1 < init_size { (i + 1) as u32 } else { NONE };
            free_prev[i] = if i > 2 { (i - 1) as u32 } else { NONE };
        }

        // Mark root (0) and reserved (1) as used
        used_bmp[0] |= 0b11;

        // used_base bitmap — ensures unique base values to prevent false transitions
        let mut used_base_bmp: Vec<u64> = vec![0u64; bmp_words];

        #[inline(always)]
        fn bmp_set(bmp: &mut [u64], idx: usize) {
            if idx / 64 >= bmp.len() { return; }
            bmp[idx / 64] |= 1u64 << (idx % 64);
        }
        #[inline(always)]
        fn bmp_test(bmp: &[u64], idx: usize) -> bool {
            idx / 64 < bmp.len() && (bmp[idx / 64] >> (idx % 64)) & 1 != 0
        }

        fn fl_remove(
            idx: usize,
            free_next: &mut [u32], free_prev: &mut [u32],
            free_head: &mut u32, free_tail: &mut u32,
        ) {
            let next = free_next[idx];
            let prev = free_prev[idx];
            if prev != NONE { free_next[prev as usize] = next; }
            if next != NONE { free_prev[next as usize] = prev; }
            if *free_head == idx as u32 { *free_head = next; }
            if *free_tail == idx as u32 { *free_tail = prev; }
            free_next[idx] = NONE;
            free_prev[idx] = NONE;
        }

        fn expand(
            new_size: usize,
            states: &mut Vec<State>,
            used_bmp: &mut Vec<u64>,
            used_base_bmp: &mut Vec<u64>,
            free_next: &mut Vec<u32>, free_prev: &mut Vec<u32>,
            free_head: &mut u32, free_tail: &mut u32,
        ) {
            let old_len = states.len();
            if new_size <= old_len { return; }
            states.resize(new_size, State::default());
            let new_bmp_words = (new_size + 63) / 64;
            used_bmp.resize(new_bmp_words, 0u64);
            used_base_bmp.resize(new_bmp_words, 0u64);
            free_next.resize(new_size, NONE);
            free_prev.resize(new_size, NONE);

            for i in old_len..new_size {
                free_next[i] = if i + 1 < new_size { (i + 1) as u32 } else { NONE };
                free_prev[i] = if i > old_len { (i - 1) as u32 } else { NONE };
            }
            if *free_tail != NONE {
                free_next[*free_tail as usize] = old_len as u32;
                free_prev[old_len] = *free_tail;
            } else {
                *free_head = old_len as u32;
            }
            *free_tail = (new_size - 1) as u32;
        }

        // BFS
        let mut trie_to_da: Vec<u32> = vec![u32::MAX; trie.len()];
        trie_to_da[0] = ROOT_IDX;

        let mut queue: Vec<usize> = vec![0];
        let mut top = 0;

        while top < queue.len() {
            let trie_node = queue[top];
            top += 1;

            let da_idx = trie_to_da[trie_node];

            let children: &HashMap<u8, usize> = &trie[trie_node].children;
            if children.is_empty() {
                continue;
            }

            let mut labels: Vec<u8> = children.keys().copied().collect();
            labels.sort_unstable();

            // Find base using free list + bitmap occupancy check (XOR addressing)
            let base = {
                let states_len = states.len();
                let first_label = labels[0];
                let mut slot = free_head;
                let mut found_base = None;
                'outer: while slot != NONE {
                    let s = slot as usize;
                    // base = s ^ first_label (so that base ^ first_label == s, a free slot)
                    let b = s as u32 ^ first_label as u32;
                    if b == 0 { slot = free_next[s]; continue; }
                    if bmp_test(&used_base_bmp, b as usize) { slot = free_next[s]; continue; }

                    for &c in &labels[1..] {
                        let ci = (b ^ c as u32) as usize;
                        if ci < states_len && bmp_test(&used_bmp, ci) {
                            slot = free_next[s];
                            continue 'outer;
                        }
                    }
                    found_base = Some(b);
                    break;
                }
                found_base.unwrap_or_else(|| {
                    // Use a base beyond current array
                    let mut b = states_len as u32;
                    if b == 0 { b = 1; }
                    while bmp_test(&used_base_bmp, b as usize) { b += 1; }
                    b
                })
            };

            // Ensure arrays large enough
            // XOR: max child index = base ^ c, worst case is base + 255 (when XOR bits align)
            let need = base as usize + 256;
            if need > states.len() {
                let new_size = need.max(states.len() * 2);
                expand(new_size, &mut states, &mut used_bmp, &mut used_base_bmp,
                       &mut free_next, &mut free_prev, &mut free_head, &mut free_tail);
            }

            // Register base
            states[da_idx as usize].base = base;
            bmp_set(&mut used_base_bmp, base as usize);

            // Set children — XOR addressing, check = label byte (u8)
            for &c in &labels {
                let child_idx = (base ^ c as u32) as usize;
                bmp_set(&mut used_bmp, child_idx);
                states[child_idx].set_check(c);  // label byte
                fl_remove(child_idx, &mut free_next, &mut free_prev,
                          &mut free_head, &mut free_tail);

                let child_trie = children[&c];
                trie_to_da[child_trie] = child_idx as u32;
                queue.push(child_trie);
            }
        }

        // --- Phase 3: Register output_pos ---
        for (trie_id, tnode) in trie.iter().enumerate() {
            if tnode.output_idx > 0 {
                let da_idx = trie_to_da[trie_id] as usize;
                states[da_idx].set_output_pos(tnode.output_idx);
            }
        }

        // --- Phase 4: Build failure links ---
        states[ROOT_IDX as usize].fail = ROOT_IDX;

        let mut fq: Vec<(u32, usize)> = vec![(ROOT_IDX, 0)];
        let mut ftop = 0;
        while ftop < fq.len() {
            let (da_parent, trie_parent) = fq[ftop];
            ftop += 1;

            for (&label, &child_trie) in &trie[trie_parent].children {
                let da_node = trie_to_da[child_trie];

                if trie[child_trie].output_idx > 0 || da_parent == ROOT_IDX {
                    states[da_node as usize].fail = ROOT_IDX;
                } else {
                    let mut f = states[da_parent as usize].fail;
                    let target = loop {
                        let f_base = states[f as usize].base;
                        if f_base != 0 {
                            let next = f_base ^ label as u32;
                            if (next as usize) < states.len()
                                && states[next as usize].check() == label
                            {
                                break next;
                            }
                        }
                        if f == ROOT_IDX { break ROOT_IDX; }
                        f = states[f as usize].fail;
                    };
                    states[da_node as usize].fail = target;

                    let target_opos = states[target as usize].output_pos();
                    if states[da_node as usize].output_pos() == 0 && target_opos != 0 {
                        states[da_node as usize].set_output_pos(target_opos);
                    }
                }

                fq.push((da_node, child_trie));
            }
        }

        // --- Phase 5: Pad states so that base ^ c is always in bounds ---
        // XOR: base ^ c <= base + 255, so max_base + 256 is sufficient.
        {
            let mut max_base: u32 = 0;
            for s in &states {
                if s.base > max_base {
                    max_base = s.base;
                }
            }
            let required = max_base as usize + 256;
            if states.len() < required {
                states.resize(required, State::default());
            }
        }

        DaGreedyAhoCorasick { states, outputs }
    }

    /// Transition: go(state, byte) with integrated failure link traversal.
    /// XOR addressing: child = base ^ c, valid iff check == c (label byte).
    /// Safety: states array is padded so that base ^ 255 is always in bounds.
    #[inline(always)]
    fn go(&self, mut state: u32, c: u8) -> u32 {
        let states = self.states.as_ptr();
        loop {
            let s = unsafe { &*states.add(state as usize) };
            let base = s.base;
            if base != 0 {
                let next = base ^ c as u32;
                if unsafe { &*states.add(next as usize) }.check() == c {
                    return next;
                }
            }
            if state == ROOT_IDX {
                return ROOT_IDX;
            }
            state = s.fail;
        }
    }

    /// Find all leftmost-longest non-overlapping matches.
    pub fn find_leftmost_longest(&self, haystack: &[u8]) -> Vec<Match> {
        let mut candidates: Vec<Match> = Vec::new();
        let mut state: u32 = ROOT_IDX;
        let states = self.states.as_ptr();
        let outputs = self.outputs.as_ptr();
        let hay_ptr = haystack.as_ptr();
        let hay_len = haystack.len();

        let mut i: usize = 0;
        while i < hay_len {
            let byte = unsafe { *hay_ptr.add(i) };
            state = self.go(state, byte);

            let check_opos = unsafe { &*states.add(state as usize) }.check_opos;
            let opos = check_opos & 0x00FF_FFFF;
            if opos != 0 {
                let out = unsafe { &*outputs.add(opos as usize) };
                let len = out.byte_len;
                let end = (i as u32) + 1;
                let start = end - len;
                let pattern = out.pattern;

                loop {
                    if let Some(last) = candidates.last() {
                        if last.end <= start {
                            candidates.push(Match { start, end, pattern });
                            break;
                        } else if start > last.start {
                            break;
                        } else {
                            candidates.pop();
                        }
                    } else {
                        candidates.push(Match { start, end, pattern });
                        break;
                    }
                }
            }
            i += 1;
        }

        candidates
    }
}