ferrix/crates/primitives/src/poptrie/mod.rs

//! A Poptrie (Asai and Ohara, SIGCOMM 2015): a compressed multiway trie for
//! longest-prefix matching.
//!
//! Work in progress. Nothing in the daemon uses it yet, and the construction
//! side has no public entry point, hence the lint expectation below.
#![expect(dead_code, reason = "not yet wired into a RIB")]

use std::collections::BTreeMap;

#[derive(Debug)]
pub struct Poptrie<T> {
    pub inners: Vec<Inner>,
    pub leaves: Vec<Leaf<T>>,
}

#[derive(Debug)]
pub struct Inner {
    /// A bit vector that describes the child internal nodes.
    pub internal_v: u64,
    /// A bit vector that describes the child leaf nodes.
    pub leaf_v: u64,
    /// Starting offset where the child interior nodes of this node begins.
    pub inner_offset: u64,
    /// Starting offset where the child leaf nodes of this node begins.
    pub leaf_offset: u64,
}

#[derive(Debug)]
pub struct Leaf<T> {
    pub data: T,
}

#[inline]
fn extract(width: u8, offset: u8, v: u32, bits: u8) -> u8 {
    let shift = bits.saturating_sub(width * (offset + 1));
    let mask = 0b111111 << shift;
    ((v & mask) >> shift) as u8
}

fn extract_32(width: u8, offset: u8, v: u32) -> u8 {
    extract(width, offset, v, 32)
}

fn extract_128(width: u8, offset: u8, v: u32) -> u8 {
    extract(width, offset, v, 128)
}

impl<T: Clone> Poptrie<T> {
    pub fn match_u32(&self, addr: u32) -> Option<&T> {
        let mut i = 0_u64;
        let mut v = self.inners.get(i as usize)?.internal_v;
        let mut offset = 0;
        let mut n = extract(6, offset, addr, 32);

        let mut result = None;

        while (v & (1 << n)) != 0 {
            let base = self.inners.get(i as usize)?.inner_offset;
            let arg = v & ((2 << n) - 1);
            let bc = arg.count_ones() as u64;
            i = base + bc - 1;
            v = self.inners.get(i as usize)?.internal_v;

            offset += 1;
            n = extract(6, offset, addr, 32);

            let base = self.inners.get(i as usize)?.leaf_offset;
            let v = self.inners.get(i as usize)?.leaf_v;
            if (v & (1 << n)) != 0 {
                let i = base - 1;
                result = Some(&self.leaves.get(i as usize)?.data);
            }
        }

        let base = self.inners.get(i as usize)?.leaf_offset;
        let v = self.inners.get(i as usize)?.leaf_v;
        if (v & (1 << n)) != 0 {
            i = base - 1;
            result = Some(&self.leaves.get(i as usize)?.data);
        }

        result
    }

    /// Fills the trie from a routing table keyed by (address, prefix length).
    fn build(&mut self, rt: BTreeMap<(u32, u8), T>) {
        const DEPTH: u8 = 6;
        let mut forest = vec![(0, rt)];

        let mut i_offset = 1;
        for _ in 0..DEPTH {
            let mut subforest = Vec::<(u8, BTreeMap<(u32, u8), T>)>::new();
            let mut children = 0;
            let mut siblings = 0;
            for (_, tree) in &forest {
                let mut iv = 0_u64;
                let mut lv = 0_u64;

                let mut subsubforest = Vec::<(u8, BTreeMap<(u32, u8), T>)>::new();
                for (r, e) in tree {
                    let k = extract(6, DEPTH, r.0, 32);
                    let consumed = core::cmp::min((DEPTH + 1) * 6, 32);
                    if r.1 < consumed {
                        if ((1 << k) & iv) == 0 {
                            lv |= 1 << k;
                            self.leaves.push(Leaf { data: e.clone() });
                        }

                        // If the prefix of the route entry is less than but not equal
                        // to the consumed number of bits, we need to add those bits
                        // to the bitvec.
                        if r.1 != consumed {
                            // Shift the extra bits and add to the bitvec for this
                            // internal node.
                            let extra = 1_u32 << (consumed - r.1);
                            for i in 1..(extra).min(64 - k as u32) {
                                lv |= 1_u64 << (k as u32 + i);
                                self.leaves.push(Leaf { data: e.clone() });
                            }
                        }
                        continue;
                    }

                    iv |= 1 << k;
                    match subsubforest.iter_mut().find(|x| x.0 == k) {
                        Some(ref mut entry) => {
                            entry.1.insert(*r, e.clone());
                        }
                        None => {
                            let mut tbl = BTreeMap::new();
                            tbl.insert(*r, e.clone());
                            subsubforest.push((k, tbl));
                            if iv > 0 {
                                children += 1;
                            }
                        }
                    }
                }

                if iv > 0 || lv > 0 {
                    self.inners.push(Inner {
                        internal_v: iv,
                        leaf_v: lv,
                        inner_offset: if iv > 0 { i_offset + siblings } else { 0 },
                        leaf_offset: self.leaves.len() as u64,
                    });
                    if iv > 0 {
                        siblings += 1;
                    }
                }

                subforest.extend_from_slice(subsubforest.as_slice());
            }
            i_offset += children;
            forest = subforest;
        }
    }
}

#[cfg(test)]
mod tests {
    use std::{collections::BTreeMap, net::Ipv4Addr};

    use crate::poptrie::Poptrie;

    #[test]
    fn test_poptrie() {
        let mut map: BTreeMap<(u32, u8), Ipv4Addr> = Default::default();
        map.insert(
            (Ipv4Addr::new(127, 0, 0, 0).into(), 8),
            Ipv4Addr::new(10, 0, 0, 1),
        );

        let mut poptrie: Poptrie<Ipv4Addr> = Poptrie {
            inners: vec![],
            leaves: vec![],
        };

        poptrie.build(map);
        assert!(
            !poptrie.inners.is_empty(),
            "building from a non-empty table yields nodes"
        );
    }
}