kiri_engine/dictionary/

trie.rs

//! Double-array trie for efficient common-prefix search.
//! Ported from Sudachi's DARTSCLONE implementation.
//!
//! The trie is stored as a flat i32 slice. Each element is a single unit
//! encoding offset, hasLeaf, label, and value in different bit fields.

/// Result of a trie lookup.
#[derive(Debug, Clone, Copy)]
pub struct TrieMatch {
    /// The value stored at the trie node.
    pub value: i32,
    /// The byte length of the matched key.
    pub length: usize,
}

// ---- DARTSCLONE unit bit-field accessors ----

#[inline]
fn unit_has_leaf(unit: i32) -> bool {
    ((unit as u32) >> 8) & 1 != 0
}

#[inline]
fn unit_value(unit: i32) -> i32 {
    unit & 0x7FFF_FFFFi32
}

#[inline]
fn unit_label(unit: i32) -> i32 {
    unit & (i32::MIN | 0xff)
}

#[inline]
fn unit_offset(unit: i32) -> u32 {
    let u = unit as u32;
    (u >> 10) << ((u & (1 << 9)) >> 6)
}

/// Read a trie from raw bytes. Returns (array, size, bytes_read).
pub fn read_trie(data: &[u8], offset: usize) -> (Vec<i32>, usize, usize) {
    let size = i32::from_le_bytes([
        data[offset],
        data[offset + 1],
        data[offset + 2],
        data[offset + 3],
    ]) as usize;

    let mut array = Vec::with_capacity(size);
    let mut pos = offset + 4;
    for _ in 0..size {
        let val = i32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);
        array.push(val);
        pos += 4;
    }

    let bytes_read = 4 + size * 4;
    (array, size, bytes_read)
}

/// Double-array trie backed by a pre-built i32 array.
pub struct DoubleArrayTrie {
    array: Vec<i32>,
    size: usize,
}

impl DoubleArrayTrie {
    /// Create from a pre-built array.
    pub fn new(array: Vec<i32>, size: usize) -> Self {
        Self { array, size }
    }

    /// Read a trie from raw bytes at the given offset.
    pub fn from_bytes(data: &[u8], offset: usize) -> (Self, usize) {
        let (array, size, bytes_read) = read_trie(data, offset);
        (Self { array, size }, bytes_read)
    }

    /// Common prefix search: find all keys that are prefixes of the input.
    /// This is the core operation for lattice-based tokenization.
    #[allow(clippy::needless_range_loop)]
    pub fn common_prefix_search(&self, key: &[u8], offset: usize, limit: usize) -> Vec<TrieMatch> {
        let mut results = Vec::new();
        let arr = &self.array;
        let end = std::cmp::min(offset + limit, key.len());

        let mut node_pos = 0u32;

        // XOR with root offset
        let root = match arr.get(node_pos as usize) {
            Some(&v) => v,
            None => return results,
        };
        node_pos ^= unit_offset(root);

        for i in offset..end {
            let b = key[i] as u32;

            // Follow edge labeled `b`
            node_pos ^= b;

            let unit = match arr.get(node_pos as usize) {
                Some(&v) => v,
                None => return results,
            };

            // Check if this node exists (label must match the byte we followed)
            if unit_label(unit) != b as i32 {
                return results;
            }

            // Advance nodePos to next level
            node_pos ^= unit_offset(unit);

            // If this unit has a leaf, read the value from the NEW nodePos
            if unit_has_leaf(unit) {
                if let Some(&leaf) = arr.get(node_pos as usize) {
                    results.push(TrieMatch {
                        value: unit_value(leaf),
                        length: i - offset + 1,
                    });
                } else {
                    return results;
                }
            }
        }

        results
    }

    /// Exact match search: look up a single key.
    #[allow(clippy::needless_range_loop)]
    pub fn exact_match_search(&self, key: &[u8], offset: usize, length: usize) -> i32 {
        let arr = &self.array;
        let end = offset + length;

        let mut node_pos = 0u32;
        let mut last_unit = 0i32;

        let root = match arr.get(node_pos as usize) {
            Some(&v) => v,
            None => return -1,
        };
        node_pos ^= unit_offset(root);

        for i in offset..end {
            let b = key[i] as u32;
            node_pos ^= b;

            let unit = match arr.get(node_pos as usize) {
                Some(&v) => v,
                None => return -1,
            };
            if unit_label(unit) != b as i32 {
                return -1;
            }

            node_pos ^= unit_offset(unit);
            last_unit = unit;
        }

        if !unit_has_leaf(last_unit) {
            return -1;
        }

        match arr.get(node_pos as usize) {
            Some(&v) => unit_value(v),
            None => -1,
        }
    }

    /// Total number of units in the trie.
    pub fn total_size(&self) -> usize {
        self.size
    }
}

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

    #[test]
    fn test_bit_field_accessors() {
        // Test with a known unit value that has a leaf
        // Bit 8 set = hasLeaf
        let unit: i32 = 0x100; // bit 8 set
        assert!(unit_has_leaf(unit));
        assert!(!unit_has_leaf(0));

        // unitLabel: extracts bits 31 and 0..7
        let unit_with_label = 0x42i32; // label = 0x42
        assert_eq!(unit_label(unit_with_label) & 0xff, 0x42);

        // unitValue: lower 31 bits
        let unit_val = 0x7FFF_FFFFi32;
        assert_eq!(unit_value(unit_val), 0x7FFF_FFFF);
    }

    #[test]
    fn test_unit_offset_basic() {
        // unitOffset: (u >> 10) << ((u & (1 << 9)) >> 6)
        // When bit 9 is 0: shift = 0, so offset = u >> 10
        let unit = 0x400i32; // bit 10 set
        assert_eq!(unit_offset(unit), 1);

        let unit2 = 0x800i32; // bit 11 set
        assert_eq!(unit_offset(unit2), 2);
    }
}