Radix 16 Trie Symbol Table

radix16symtable.rs

// top nibble are flags, bottom nibble is the value
type ParentVaL  = u8;
type Trie16Idx       = u32;
type ChildIdx        = Option<core::num::NonZeroU32>;

// This is a proper tree, the parent pointers are used to represent string as a linked list,
struct Trie16Node {
    parent_val : ParentVaL,
    parent     : Trie16Idx,
    children   : [ChildIdx; 16],
}
impl Clone for Trie16Node {
    fn clone(&self) -> Self {
        *self       
    }
}
impl Copy for Trie16Node {}

const fn trie16node_uninit() -> Trie16Node {
    Trie16Node { 
        parent_val : 0,
        parent     : 0, 
        children   : [None;16]
    }
}

struct SymbolTable {
    store : Vec<Trie16Node>,
}

fn debug_symbol_table(symbol_table : &SymbolTable, mut w : impl std::io::Write) -> Result<(),std::io::Error> {
    w.write_all(&*b"SymbolTable {\n")?;
    w.write_all(&*b"    | index      | val        | parent     | children\n")?;
    w.write_all(&*b"    +------------+------------+------------+---------\n")?;
    for each in 0..symbol_table.store.len() {
        w.write_all(&*b"    | ")?;
        let Trie16Node { parent_val: flags_parent_val, parent, children } = symbol_table.store[each];
        write!(w,"0x{each:_>8x} | 0x{flags_parent_val:_>2x}       | 0x{parent:_>8x} | [  ")?;
        for each in 0..0x10 {
            match children[each] {
                Some(v) => write!(w,"0x{each:_>2x}->0x{:_>8x}  ", v.get())?,
                None => {},
            }
        }
        w.write_all(&*b"]\n")?;
    }
    w.write_all(&*b"}\n")?;
    Ok(())

}

fn new_symbol_table() -> SymbolTable {
    let mut store = Vec::with_capacity(1024);
    store.push(trie16node_uninit());
    SymbolTable { store  }
}


struct NibblePos{ byte : usize, end : bool}
fn existing_prefix(sym_table : &SymbolTable, bytes : &[u8]) -> (Trie16Idx, NibblePos) {
    const RHS_BITS: i8 = 0x0F;
    let len = bytes.len() as isize;
    
    let mut cur_forwards           = 0;
    let (mut i, mut cur_mask_bits) = (0, !RHS_BITS);
    

    'find : while (0..len).contains(&i) {
        let f_nibble = ((bytes[i as usize] & cur_mask_bits as u8) >> cur_mask_bits.trailing_zeros())  as usize;
        match sym_table.store[cur_forwards ].children[f_nibble as usize]
        {   None    => break 'find,
            Some(f) => cur_forwards = f.get() as usize,
        };
        i += (cur_mask_bits ==  RHS_BITS) as isize;
        cur_mask_bits = !cur_mask_bits;
    }
    (cur_forwards as u32, NibblePos{ byte : i.clamp(0, len) as usize, end : cur_mask_bits == RHS_BITS})
}

/// this should only ever be used to add a string's suffix that is __not__ already in the trie.
fn allocate_from_existing_prefix(sym_table : &mut SymbolTable, prefix : Trie16Idx, bytes : &[u8], pos : NibblePos) -> Trie16Idx {
    const RHS_BITS: i8 = 0x0F;
    let len   = bytes.len() as isize;
    let mut i = pos.byte as isize;

    let mut cur = prefix;
    
    let mut cur_mask_bits : i8 = (0 - pos.end as i8) ^ !RHS_BITS;

    sym_table.store.reserve((i+1) as usize*2);
    while (0..len).contains(&i) {

        let f_nibble = ((bytes[i as usize] & cur_mask_bits as u8) >> cur_mask_bits.trailing_zeros());
    
        let next_idx = sym_table.store.len();
        sym_table.store[cur as usize].children[ f_nibble as usize ] = Some(unsafe { core::num::NonZero::new_unchecked(next_idx as u32) });
        sym_table.store.push(trie16node_uninit());
        sym_table.store[next_idx].parent = cur as u32;
        sym_table.store[next_idx].parent_val |= f_nibble;

        
        cur = next_idx as u32;

        i += (cur_mask_bits ==  RHS_BITS) as isize;
        
        cur_mask_bits = !cur_mask_bits;
    }

    

    cur
}

fn insert(sym_table : &mut SymbolTable, bytes : &[u8]) -> Trie16Idx {
    let len = bytes.len() as isize;

    const MASK_BITS: i32 = 0x0F;

    
    let cur_forwards = {
        let (prefix, pos) = existing_prefix(sym_table, bytes);
        if pos.byte == len as usize {
            core::debug_assert!(!pos.end);
            prefix
        } else {
            (allocate_from_existing_prefix(sym_table, prefix, bytes, pos))
        }
    };
    cur_forwards as Trie16Idx
}



fn dbg_print_table(table : &SymbolTable) {
    let mut s = String::new();
    debug_symbol_table(table, unsafe { s.as_mut_vec() });

    println!("{}", s);
}

#[test]
fn basics() {
    let mut table = new_symbol_table();

    insert(&mut table, b"hello");
    insert(&mut table, b"h");
    insert(&mut table, b"help");
    insert(&mut table, b"half");
    insert(&mut table, b"beer");
    insert(&mut table, b"beef");
    let mut s = String::new();
    debug_symbol_table(&table, unsafe { s.as_mut_vec() });

    println!("{}", s);    
}

#[test]
fn basics2() {
    let mut table = new_symbol_table();
    

    let mut last = 0;
    for i in  0..WORDS.len()/10 {
        if WORDS[i] == b'\n' {
            insert(&mut table, &WORDS[last..i]);
            last = i+1;

        }
    }

    println!("{}", size_of_val(&table.store[..]));
    println!("{}", size_of_val(&WORDS[..]));

    let t = std::time::Instant::now();
    let (mut prefix, _) = existing_prefix(&table, b"zyloprim");
    let e = t.elapsed();
    println!("{prefix:?}, {e:?}");

    let t = std::time::Instant::now();
    unsafe { str::from_utf8_unchecked(WORDS) }.split_terminator('\n').find(|x|x==&"zyloprim");
    let e = t.elapsed();
    println!("?? {e:?}");
    
    let mut toggle = false;
    let mut v = Vec::new();
    while prefix != 0 {
        let Trie16Node { parent_val: flags_parent_val, parent, children } = table.store[prefix as usize];

        prefix = parent;

        toggle = !toggle;

        if toggle { 
            v.push(flags_parent_val);
        } else {
            *v.last_mut().unwrap() |= flags_parent_val << 4
        }
    }
    v.reverse();
    println!("{:?}", v ); 
    println!("{}", unsafe { str::from_utf8_unchecked(&v) }); 


    let mut s = String::new();
    debug_symbol_table(&table, unsafe { s.as_mut_vec() });

    println!("{}", s);


// https://github.com/kkrypt0nn/wordlists/blob/main/wordlists/languages/english.txt
const WORDS : &[u8] =  &*include_bytes!("../english.txt");

}