sozu_lib/router/

pattern_trie.rs

use std::{collections::HashMap, fmt::Debug, iter, str};

use regex::bytes::Regex;

pub type Key = Vec<u8>;
pub type KeyValue<K, V> = (K, V);

#[derive(Debug, PartialEq, Eq)]
pub enum InsertResult {
    Ok,
    Existing,
    Failed,
}

#[derive(Debug, PartialEq, Eq)]
pub enum RemoveResult {
    Ok,
    NotFound,
}

fn find_last_dot(input: &[u8]) -> Option<usize> {
    //println!("find_last_dot: input = {}", from_utf8(input).unwrap());
    (0..input.len()).rev().find(|&i| input[i] == b'.')
}

fn find_last_slash(input: &[u8]) -> Option<usize> {
    //println!("find_last_dot: input = {}", from_utf8(input).unwrap());
    (0..input.len()).rev().find(|&i| input[i] == b'/')
}

/// Implementation of a trie tree structure.
/// In Sozu this is used to store and lookup domains recursively.
/// Each node represents a "level domain".
/// A leaf node (leftmost label) can be a wildcard, a regex pattern or a plain string.
/// Leaves also store a value associated with the complete domain.
/// For Sozu it is a list of (PathRule, MethodRule, ClusterId). See the Router strucure.
#[derive(Debug, Default)]
pub struct TrieNode<V> {
    key_value: Option<KeyValue<Key, V>>,
    wildcard: Option<KeyValue<Key, V>>,
    children: HashMap<Key, TrieNode<V>>,
    regexps: Vec<(Regex, TrieNode<V>)>,
}

/// One step of a trie traversal where a non-literal segment matched.
///
/// `Wildcard` carries the actual segment bytes consumed by a `*` wildcard
/// (so a router that wants to capture them can splice them into a rewrite
/// template). `Regexp` carries both the matched bytes and the regex itself
/// so the caller can re-run `Regex::captures` to pull explicit groups.
#[derive(Debug)]
pub enum TrieSubMatch<'a, 'b> {
    Wildcard(&'a [u8]),
    Regexp(&'a [u8], &'b Regex),
}

/// Ordered list of non-literal trie segments visited during a successful
/// `lookup_with_path` traversal. Routers feed the entries into rewrite
/// templates (`$HOST[n]`) so frontend rewrites can reach into the matched
/// segments. Empty when only literal segments matched.
pub type TrieMatches<'a, 'b> = Vec<TrieSubMatch<'a, 'b>>;

impl<V: PartialEq> std::cmp::PartialEq for TrieNode<V> {
    fn eq(&self, other: &Self) -> bool {
        self.key_value == other.key_value
            && self.wildcard == other.wildcard
            && self.children == other.children
            && self.regexps.len() == other.regexps.len()
            && self
                .regexps
                .iter()
                .zip(other.regexps.iter())
                .fold(true, |b, (left, right)| {
                    b && left.0.as_str() == right.0.as_str() && left.1 == right.1
                })
    }
}

impl<V: Debug + Clone> TrieNode<V> {
    pub fn new(key: Key, value: V) -> TrieNode<V> {
        TrieNode {
            key_value: Some((key, value)),
            wildcard: None,
            children: HashMap::new(),
            regexps: Vec::new(),
        }
    }

    pub fn wildcard(key: Key, value: V) -> TrieNode<V> {
        TrieNode {
            key_value: None,
            wildcard: Some((key, value)),
            children: HashMap::new(),
            regexps: Vec::new(),
        }
    }

    pub fn root() -> TrieNode<V> {
        TrieNode {
            key_value: None,
            wildcard: None,
            children: HashMap::new(),
            regexps: Vec::new(),
        }
    }

    pub fn is_empty(&self) -> bool {
        self.key_value.is_none()
            && self.wildcard.is_none()
            && self.regexps.is_empty()
            && self.children.is_empty()
    }

    pub fn insert(&mut self, key: Key, value: V) -> InsertResult {
        //println!("insert: key == {}", std::str::from_utf8(&key).unwrap());
        if key.is_empty() {
            return InsertResult::Failed;
        }
        if key[..] == b"."[..] {
            return InsertResult::Failed;
        }

        #[cfg(debug_assertions)]
        let before = self.count_values();

        let insert_result = self.insert_recursive(&key, &key, value);
        assert_ne!(insert_result, InsertResult::Failed);

        // Post: the value count grows by exactly one on a fresh insert
        // and is unchanged when the key already existed. `Failed` is
        // ruled out above, so these two are the only reachable cases.
        #[cfg(debug_assertions)]
        {
            let after = self.count_values();
            match insert_result {
                InsertResult::Ok => debug_assert_eq!(
                    after,
                    before + 1,
                    "a fresh insert must add exactly one value to the trie",
                ),
                InsertResult::Existing => debug_assert_eq!(
                    after, before,
                    "an Existing insert must not change the trie value count",
                ),
                InsertResult::Failed => {
                    unreachable!("insert_recursive returned Failed after key validation")
                }
            }
            self.check_invariants();
        }

        insert_result
    }

    pub fn insert_recursive(&mut self, partial_key: &[u8], key: &Key, value: V) -> InsertResult {
        //println!("insert_rec: key == {}", std::str::from_utf8(partial_key).unwrap());
        assert_ne!(partial_key, &b""[..]);
        // `partial_key` is always a suffix of the full `key` being
        // inserted — the recursion only ever shrinks the head, never
        // rewrites the tail.
        debug_assert!(
            partial_key.len() <= key.len(),
            "insert recursion must consume the key, never grow past it",
        );

        if partial_key[partial_key.len() - 1] == b'/' {
            let pos = find_last_slash(&partial_key[..partial_key.len() - 1]);

            if let Some(pos) = pos {
                if pos > 0 && partial_key[pos - 1] != b'.' {
                    return InsertResult::Failed;
                }

                if let Ok(s) = str::from_utf8(&partial_key[pos + 1..partial_key.len() - 1]) {
                    let anchored_s = format!("\\A{s}\\z");
                    debug_assert!(
                        anchored_s.starts_with("\\A") && anchored_s.ends_with("\\z"),
                        "segment regex must be fully anchored so it matches the whole segment only",
                    );
                    for t in self.regexps.iter_mut() {
                        if t.0.as_str() == anchored_s {
                            // `pos > 0`: there is a `.`-separated prefix
                            // before this regex segment; recurse on it
                            // (dropping the leading `.` via `pos - 1`).
                            // `pos == 0`: the regex is the leftmost/only
                            // segment, so its subtree is already a
                            // value-bearing leaf (the create-path below
                            // built it via `TrieNode::new`); re-inserting
                            // the same host is `Existing`. The pre-fix code
                            // did `partial_key[..pos - 1]` unconditionally,
                            // underflowing to `usize::MAX` and panicking on
                            // `pos == 0` (same latent bug as `lookup_mut`).
                            if pos > 0 {
                                return t.1.insert_recursive(&partial_key[..pos - 1], key, value);
                            } else {
                                return InsertResult::Existing;
                            }
                        }
                    }

                    // Anchor segment regexes so they only match the entire
                    // segment, not partial overlaps. Without `\A...\z`, a
                    // pattern like `cdn[0-9]+` would match `cdn123xxx`,
                    // which silently widens the routing surface.
                    let anchored = format!("\\A{s}\\z");
                    if let Ok(r) = Regex::new(&anchored) {
                        if pos > 0 {
                            let mut node = TrieNode::root();
                            let pos = pos - 1;

                            let res = node.insert_recursive(&partial_key[..pos], key, value);

                            if res == InsertResult::Ok {
                                self.regexps.push((r, node));
                            }

                            return res;
                        } else {
                            let node = TrieNode::new(key.to_vec(), value);
                            self.regexps.push((r, node));
                            return InsertResult::Ok;
                        }
                    }
                }
            }

            return InsertResult::Failed;
        }

        let pos = find_last_dot(partial_key);
        match pos {
            None => {
                if self.children.contains_key(partial_key) {
                    InsertResult::Existing
                } else if partial_key == &b"*"[..] {
                    if self.wildcard.is_some() {
                        InsertResult::Existing
                    } else {
                        self.wildcard = Some((key.to_vec(), value));
                        InsertResult::Ok
                    }
                } else {
                    let node = TrieNode::new(key.to_vec(), value);
                    self.children.insert(partial_key.to_vec(), node);
                    InsertResult::Ok
                }
            }
            Some(pos) => {
                // The dot at `pos` is kept on the child key (suffix) and
                // stripped from the recursive prefix; the two slices
                // partition `partial_key` exactly.
                debug_assert_eq!(
                    partial_key[..pos].len() + partial_key[pos..].len(),
                    partial_key.len(),
                    "dot-split must partition partial_key without losing bytes",
                );
                debug_assert_eq!(
                    partial_key[pos], b'.',
                    "find_last_dot must point at a '.' byte",
                );
                if let Some(child) = self.children.get_mut(&partial_key[pos..]) {
                    return child.insert_recursive(&partial_key[..pos], key, value);
                }

                let mut node = TrieNode::root();
                let res = node.insert_recursive(&partial_key[..pos], key, value);

                if res == InsertResult::Ok {
                    self.children.insert(partial_key[pos..].to_vec(), node);
                }

                res
            }
        }
    }

    pub fn remove(&mut self, key: &Key) -> RemoveResult {
        #[cfg(debug_assertions)]
        let before = self.count_values();

        let remove_result = self.remove_recursive(key);

        // Post: a successful remove drops exactly one value; a NotFound
        // is a no-op on the value count. The structural invariants then
        // guarantee no emptied subtree was stranded by the prune.
        #[cfg(debug_assertions)]
        {
            let after = self.count_values();
            match remove_result {
                RemoveResult::Ok => debug_assert_eq!(
                    after + 1,
                    before,
                    "a successful remove must drop exactly one value from the trie",
                ),
                RemoveResult::NotFound => debug_assert_eq!(
                    after, before,
                    "a NotFound remove must not change the trie value count",
                ),
            }
            self.check_invariants();
        }

        remove_result
    }

    pub fn remove_recursive(&mut self, partial_key: &[u8]) -> RemoveResult {
        //println!("remove: key == {}", std::str::from_utf8(partial_key).unwrap());

        if partial_key.is_empty() {
            if self.key_value.is_some() {
                self.key_value = None;
                return RemoveResult::Ok;
            } else {
                return RemoveResult::NotFound;
            }
        }

        if partial_key == &b"*"[..] {
            if self.wildcard.is_some() {
                self.wildcard = None;
                return RemoveResult::Ok;
            } else {
                return RemoveResult::NotFound;
            }
        }

        if partial_key[partial_key.len() - 1] == b'/' {
            let pos = find_last_slash(&partial_key[..partial_key.len() - 1]);

            if let Some(pos) = pos {
                if pos > 0 && partial_key[pos - 1] != b'.' {
                    return RemoveResult::NotFound;
                }

                if let Ok(s) = str::from_utf8(&partial_key[pos + 1..partial_key.len() - 1]) {
                    let anchored_s = format!("\\A{s}\\z");
                    if pos > 0 {
                        let mut remove_result = RemoveResult::NotFound;
                        for t in self.regexps.iter_mut() {
                            if t.0.as_str() == anchored_s
                                && t.1.remove_recursive(&partial_key[..pos - 1]) == RemoveResult::Ok
                            {
                                remove_result = RemoveResult::Ok;
                            }
                        }
                        return remove_result;
                    } else {
                        let len = self.regexps.len();
                        self.regexps.retain(|(r, _)| r.as_str() != anchored_s);
                        if len > self.regexps.len() {
                            return RemoveResult::Ok;
                        }
                    }
                }
            }

            return RemoveResult::NotFound;
        }

        let pos = find_last_dot(partial_key);
        let (prefix, suffix) = match pos {
            None => (&b""[..], partial_key),
            Some(pos) => (&partial_key[..pos], &partial_key[pos..]),
        };
        //println!("remove: prefix|suffix: {} | {}", std::str::from_utf8(prefix).unwrap(), std::str::from_utf8(suffix).unwrap());
        debug_assert_eq!(
            prefix.len() + suffix.len(),
            partial_key.len(),
            "dot-split must partition the key without losing or duplicating bytes",
        );

        match self.children.get_mut(suffix) {
            Some(child) => match child.remove_recursive(prefix) {
                RemoveResult::NotFound => RemoveResult::NotFound,
                RemoveResult::Ok => {
                    // An emptied child subtree MUST be pruned here so the
                    // parent never strands a node with no value. After the
                    // prune the suffix key is gone from `children`.
                    if child.is_empty() {
                        self.children.remove(suffix);
                        debug_assert!(
                            !self.children.contains_key(suffix),
                            "an emptied child subtree must be removed from the parent",
                        );
                    } else {
                        // `count_values` is debug-only; gate the whole
                        // assert so the call does not have to compile in
                        // release (HARD RULE 2 — E0425 guard).
                        #[cfg(debug_assertions)]
                        debug_assert!(
                            child.count_values() > 0,
                            "a retained child subtree must still hold at least one value",
                        );
                    }
                    RemoveResult::Ok
                }
            },
            None => RemoveResult::NotFound,
        }
    }

    /// Look up `partial_key` and additionally collect the non-literal segments
    /// that matched along the way (`TrieMatches`).
    ///
    /// Equivalent to `lookup` for callers that don't need the captures, but
    /// frontends with `$HOST[n]` rewrite templates need the matched segments
    /// to fill the placeholders. The accumulator is passed in by value so
    /// callers can pre-size it (`Vec::with_capacity`) and we own the path
    /// returned alongside the value.
    pub fn lookup_with_path<'a, 'b>(
        &'b self,
        partial_key: &'a [u8],
        accept_wildcard: bool,
        mut trace: TrieMatches<'a, 'b>,
    ) -> Option<(&'b KeyValue<Key, V>, TrieMatches<'a, 'b>)> {
        if partial_key.is_empty() {
            return self.key_value.as_ref().map(|kv| (kv, trace));
        }

        let pos = find_last_dot(partial_key);
        let (prefix, suffix) = match pos {
            None => (&b""[..], partial_key),
            Some(pos) => (&partial_key[..pos], &partial_key[pos..]),
        };
        // The dot-split partitions the key exactly: prefix ++ suffix is
        // the whole input, and a dotted split puts the `.` at the head
        // of the suffix (this is the byte the wildcard/regex arms strip).
        debug_assert_eq!(
            prefix.len() + suffix.len(),
            partial_key.len(),
            "dot-split must partition the key without losing or duplicating bytes",
        );
        debug_assert!(
            pos.is_none() || suffix.first() == Some(&b'.'),
            "a dotted split must place the separator at the head of the suffix",
        );

        match self.children.get(suffix) {
            Some(child) => child.lookup_with_path(prefix, accept_wildcard, trace),
            None => {
                if prefix.is_empty() && self.wildcard.is_some() && accept_wildcard {
                    let segment = if !suffix.is_empty() && suffix[0] == b'.' {
                        &suffix[1..]
                    } else {
                        suffix
                    };
                    trace.push(TrieSubMatch::Wildcard(segment));
                    self.wildcard.as_ref().map(|kv| (kv, trace))
                } else {
                    for (regexp, child) in self.regexps.iter() {
                        let segment = if !suffix.is_empty() && suffix[0] == b'.' {
                            &suffix[1..]
                        } else {
                            suffix
                        };
                        if regexp.is_match(segment) {
                            let mut next = trace;
                            next.push(TrieSubMatch::Regexp(segment, regexp));
                            return child.lookup_with_path(prefix, accept_wildcard, next);
                        }
                    }
                    None
                }
            }
        }
    }

    pub fn lookup(&self, partial_key: &[u8], accept_wildcard: bool) -> Option<&KeyValue<Key, V>> {
        //println!("lookup: key == {}", std::str::from_utf8(partial_key).unwrap());

        if partial_key.is_empty() {
            return self.key_value.as_ref();
        }

        let pos = find_last_dot(partial_key);
        let (prefix, suffix) = match pos {
            None => (&b""[..], partial_key),
            Some(pos) => (&partial_key[..pos], &partial_key[pos..]),
        };
        //println!("lookup: prefix|suffix: {} | {}", std::str::from_utf8(prefix).unwrap(), std::str::from_utf8(suffix).unwrap());
        debug_assert_eq!(
            prefix.len() + suffix.len(),
            partial_key.len(),
            "dot-split must partition the key without losing or duplicating bytes",
        );
        debug_assert!(
            !suffix.is_empty(),
            "the suffix the trie matches children against must be non-empty",
        );

        match self.children.get(suffix) {
            Some(child) => child.lookup(prefix, accept_wildcard),
            None => {
                //println!("no child found, testing wildcard and regexps");

                if prefix.is_empty() && self.wildcard.is_some() && accept_wildcard {
                    //println!("no dot, wildcard applies");
                    self.wildcard.as_ref()
                } else {
                    //println!("there's still a subdomain, wildcard does not apply");

                    for (regexp, child) in self.regexps.iter() {
                        let suffix = if suffix[0] == b'.' {
                            &suffix[1..]
                        } else {
                            suffix
                        };
                        //println!("testing regexp: {} on suffix {}", r.as_str(), str::from_utf8(s).unwrap());

                        if regexp.is_match(suffix) {
                            //println!("matched");
                            return child.lookup(prefix, accept_wildcard);
                        }
                    }

                    None
                }
            }
        }
    }

    pub fn lookup_mut(
        &mut self,
        partial_key: &[u8],
        accept_wildcard: bool,
    ) -> Option<&mut KeyValue<Key, V>> {
        //println!("lookup: key == {}", std::str::from_utf8(partial_key).unwrap());

        if partial_key.is_empty() {
            return self.key_value.as_mut();
        }

        if partial_key == &b"*"[..] {
            return self.wildcard.as_mut();
        }

        if partial_key[partial_key.len() - 1] == b'/' {
            let pos = find_last_slash(&partial_key[..partial_key.len() - 1]);

            if let Some(pos) = pos {
                if pos > 0 && partial_key[pos - 1] != b'.' {
                    return None;
                }

                if let Ok(s) = str::from_utf8(&partial_key[pos + 1..partial_key.len() - 1]) {
                    let anchored_s = format!("\\A{s}\\z");
                    for t in self.regexps.iter_mut() {
                        if t.0.as_str() == anchored_s {
                            // `pos == 0` means the regex is the leftmost
                            // segment and its subtree is a value-bearing
                            // leaf, reachable via the empty-prefix recursion
                            // (`lookup_mut(b"")` returns `key_value`). The
                            // pre-fix `partial_key[..pos - 1]` underflowed to
                            // `usize::MAX` and panicked on `pos == 0` — the
                            // same latent bug as the insert dedup loop. Drop
                            // the leading `.` only when there is one.
                            let rest = if pos > 0 {
                                &partial_key[..pos - 1]
                            } else {
                                &partial_key[..0]
                            };
                            return t.1.lookup_mut(rest, accept_wildcard);
                        }
                    }
                }
            }

            return None;
        }

        let pos = find_last_dot(partial_key);
        let (prefix, suffix) = match pos {
            None => (&b""[..], partial_key),
            Some(pos) => (&partial_key[..pos], &partial_key[pos..]),
        };
        //println!("lookup: prefix|suffix: {} | {}", std::str::from_utf8(prefix).unwrap(), std::str::from_utf8(suffix).unwrap());
        debug_assert_eq!(
            prefix.len() + suffix.len(),
            partial_key.len(),
            "dot-split must partition the key without losing or duplicating bytes",
        );
        debug_assert!(
            !suffix.is_empty(),
            "the suffix the trie matches children against must be non-empty",
        );

        match self.children.get_mut(suffix) {
            Some(child) => child.lookup_mut(prefix, accept_wildcard),
            None => {
                //println!("no child found, testing wildcard and regexps");

                if prefix.is_empty() && self.wildcard.is_some() && accept_wildcard {
                    //println!("no dot, wildcard applies");
                    self.wildcard.as_mut()
                } else {
                    //println!("there's still a subdomain, wildcard does not apply");

                    for &mut (ref regexp, ref mut child) in self.regexps.iter_mut() {
                        let suffix = if suffix[0] == b'.' {
                            &suffix[1..]
                        } else {
                            suffix
                        };
                        //println!("testing regexp: {} on suffix {}", r.as_str(), str::from_utf8(s).unwrap());

                        if regexp.is_match(suffix) {
                            //println!("matched");
                            return child.lookup_mut(prefix, accept_wildcard);
                        }
                    }

                    None
                }
            }
        }
    }

    pub fn print(&self) {
        self.print_recursive(b"", 0)
    }

    pub fn print_recursive(&self, partial_key: &[u8], indent: u8) {
        let raw_prefix: Vec<u8> = iter::repeat_n(b' ', 2 * indent as usize).collect();
        let prefix = str::from_utf8(&raw_prefix).unwrap();

        print!("{}{}: ", prefix, str::from_utf8(partial_key).unwrap());
        if let Some((ref key, ref value)) = self.key_value {
            print!("({}, {:?}) | ", str::from_utf8(key).unwrap(), value);
        } else {
            print!("None | ");
        }

        if let Some((key, value)) = &self.wildcard {
            println!("({}, {:?})", str::from_utf8(key).unwrap(), value);
        } else {
            println!("None");
        }

        for (child_key, child) in self.children.iter() {
            child.print_recursive(child_key, indent + 1);
        }

        for (regexp, child) in self.regexps.iter() {
            //print!("{}{}:", prefix, regexp.as_str());
            child.print_recursive(regexp.as_str().as_bytes(), indent + 1);
        }
    }

    /// Visit every stored value in the trie (the literal `key_value` and
    /// the leftmost `wildcard` slot of every node, plus all
    /// regex-subtree leaves) and invoke `f` on each. Used by the router
    /// to walk all routes for cross-cutting refreshes (e.g. listener-
    /// default HSTS reflow) without rebuilding the trie.
    pub fn for_each_value_mut<F: FnMut(&mut V)>(&mut self, f: &mut F) {
        if let Some((_, ref mut value)) = self.key_value {
            f(value);
        }
        if let Some((_, ref mut value)) = self.wildcard {
            f(value);
        }
        for child in self.children.values_mut() {
            child.for_each_value_mut(f);
        }
        for (_, child) in self.regexps.iter_mut() {
            child.for_each_value_mut(f);
        }
    }

    /// Count every value slot reachable from this node: the literal
    /// `key_value`, the leftmost `wildcard`, plus all values stored in
    /// child subtrees and regex subtrees. Used only by the
    /// `#[cfg(debug_assertions)]` invariant checks as the leaf-count
    /// accounting (`inserts − removes`); never called in release.
    #[cfg(debug_assertions)]
    fn count_values(&self) -> usize {
        let local = self.key_value.is_some() as usize + self.wildcard.is_some() as usize;
        let in_children: usize = self.children.values().map(TrieNode::count_values).sum();
        let in_regexps: usize = self.regexps.iter().map(|(_, c)| c.count_values()).sum();
        local + in_children + in_regexps
    }

    /// Full structural invariant sweep for the trie, asserted as a
    /// run-to-completion postcondition at the end of every mutating
    /// public operation. Encodes the cross-field invariants that the
    /// recursive insert/remove logic must preserve:
    ///
    /// - **No stranded interior node**: every non-root node reachable
    ///   through `children` / `regexps` must hold a value somewhere in
    ///   its subtree (`!is_empty()` and `count_values() > 0`). A node
    ///   that holds neither a value nor any descendant value is a leak —
    ///   `remove_recursive` is supposed to prune it via `is_empty()`.
    /// - **Unique regex segments**: the anchored pattern strings stored
    ///   in `regexps` are unique within a node (insert dedups by
    ///   `as_str()` before pushing a new subtree).
    /// - **Child-key invariant**: no child is keyed by the empty slice.
    ///
    /// `debug_assertions`-only; compiled out of release builds.
    #[cfg(debug_assertions)]
    fn check_invariants(&self) {
        // Regex segment patterns are unique per node.
        for i in 0..self.regexps.len() {
            for j in (i + 1)..self.regexps.len() {
                debug_assert_ne!(
                    self.regexps[i].0.as_str(),
                    self.regexps[j].0.as_str(),
                    "trie node must not hold two subtrees for the same regex segment",
                );
            }
        }

        for (child_key, child) in self.children.iter() {
            debug_assert!(
                !child_key.is_empty(),
                "trie child must not be keyed by the empty segment",
            );
            // A child subtree that has been fully emptied must have been
            // pruned by remove_recursive; reaching it here means a
            // subtree was stranded.
            debug_assert!(
                !child.is_empty(),
                "trie must not strand an empty child subtree (remove must prune)",
            );
            debug_assert!(
                child.count_values() > 0,
                "trie child subtree must lead to at least one value",
            );
            child.check_invariants();
        }

        for (_, child) in self.regexps.iter() {
            debug_assert!(
                !child.is_empty(),
                "trie must not strand an empty regex subtree (remove must prune)",
            );
            debug_assert!(
                child.count_values() > 0,
                "trie regex subtree must lead to at least one value",
            );
            child.check_invariants();
        }
    }

    pub fn domain_insert(&mut self, key: Key, value: V) -> InsertResult {
        self.insert(key, value)
    }

    pub fn domain_remove(&mut self, key: &Key) -> RemoveResult {
        self.remove(key)
    }

    pub fn domain_lookup(&self, key: &[u8], accept_wildcard: bool) -> Option<&KeyValue<Key, V>> {
        self.lookup(key, accept_wildcard)
    }

    pub fn domain_lookup_mut(
        &mut self,
        key: &[u8],
        accept_wildcard: bool,
    ) -> Option<&mut KeyValue<Key, V>> {
        self.lookup_mut(key, accept_wildcard)
    }

    pub fn size(&self) -> usize {
        ::std::mem::size_of::<TrieNode<V>>()
            + ::std::mem::size_of::<Option<KeyValue<Key, V>>>() * 2
            + self
                .children
                .iter()
                .fold(0, |acc, c| acc + c.0.len() + c.1.size())
    }

    pub fn to_hashmap(&self) -> HashMap<Key, V> {
        let mut h = HashMap::new();

        self.to_hashmap_recursive(&mut h);

        h
    }

    pub fn to_hashmap_recursive(&self, h: &mut HashMap<Key, V>) {
        if let Some((key, value)) = &self.key_value {
            h.insert(key.clone(), value.clone());
        }

        if let Some((key, value)) = &self.wildcard {
            h.insert(key.clone(), value.clone());
        }

        for child in self.children.values() {
            child.to_hashmap_recursive(h);
        }
    }
}

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

    #[test]
    fn insert() {
        let mut root: TrieNode<u8> = TrieNode::root();
        root.print();

        assert_eq!(
            root.domain_insert(Vec::from(&b"abcd"[..]), 1),
            InsertResult::Ok
        );
        root.print();
        assert_eq!(
            root.domain_insert(Vec::from(&b"abce"[..]), 2),
            InsertResult::Ok
        );
        root.print();
        assert_eq!(
            root.domain_insert(Vec::from(&b"abgh"[..]), 3),
            InsertResult::Ok
        );
        root.print();

        assert_eq!(
            root.domain_lookup(&b"abce"[..], true),
            Some(&(b"abce"[..].to_vec(), 2))
        );
        //assert!(false);
    }

    #[test]
    fn remove() {
        let mut root: TrieNode<u8> = TrieNode::root();
        println!("creating root:");
        root.print();

        println!("adding (abcd, 1)");
        assert_eq!(root.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
        root.print();
        println!("adding (abce, 2)");
        assert_eq!(root.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);
        root.print();
        println!("adding (abgh, 3)");
        assert_eq!(root.insert(Vec::from(&b"abgh"[..]), 3), InsertResult::Ok);
        root.print();

        let mut root2: TrieNode<u8> = TrieNode::root();

        assert_eq!(root2.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
        assert_eq!(root2.insert(Vec::from(&b"abgh"[..]), 3), InsertResult::Ok);

        println!("before remove");
        root.print();
        assert_eq!(root.remove(&Vec::from(&b"abce"[..])), RemoveResult::Ok);
        println!("after remove");
        root.print();

        println!("expected");
        root2.print();
        assert_eq!(root, root2);

        assert_eq!(root.remove(&Vec::from(&b"abgh"[..])), RemoveResult::Ok);
        println!("after remove");
        root.print();
        println!("expected");
        let mut root3: TrieNode<u8> = TrieNode::root();
        assert_eq!(root3.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
        root3.print();
        assert_eq!(root, root3);
    }

    #[test]
    fn insert_remove_through_regex() {
        let mut root: TrieNode<u8> = TrieNode::root();
        println!("creating root:");
        root.print();

        println!("adding (www./.*/.com, 1)");
        assert_eq!(
            root.insert(Vec::from(&b"www./.*/.com"[..]), 1),
            InsertResult::Ok
        );
        root.print();
        println!("adding (www.doc./.*/.com, 2)");
        assert_eq!(
            root.insert(Vec::from(&b"www.doc./.*/.com"[..]), 2),
            InsertResult::Ok
        );
        root.print();
        assert_eq!(
            root.domain_lookup(b"www.sozu.com".as_ref(), false),
            Some(&(b"www./.*/.com".to_vec(), 1))
        );
        assert_eq!(
            root.domain_lookup(b"www.doc.sozu.com".as_ref(), false),
            Some(&(b"www.doc./.*/.com".to_vec(), 2))
        );

        assert_eq!(
            root.domain_remove(&b"www./.*/.com".to_vec()),
            RemoveResult::Ok
        );
        root.print();
        assert_eq!(root.domain_lookup(b"www.sozu.com".as_ref(), false), None);
        assert_eq!(
            root.domain_lookup(b"www.doc.sozu.com".as_ref(), false),
            Some(&(b"www.doc./.*/.com".to_vec(), 2))
        );
    }

    /// Segment regexes must match the entire segment, not just a prefix.
    /// Without `\A...\z` anchoring the previous behaviour matched any
    /// segment whose prefix satisfied the pattern, silently widening the
    /// routing surface. This regression test exercises the exact-match
    /// invariant that anchoring guarantees.
    #[test]
    fn segment_regex_rejects_partial_matches() {
        let mut root: TrieNode<u8> = TrieNode::root();
        // The regex segment `cdn[0-9]+` must match `cdn1`, `cdn99`, etc.
        // exactly — never `cdn1xxx` or `xxxcdn1` as a prefix/suffix.
        assert_eq!(
            root.insert(Vec::from(&b"/cdn[0-9]+/.example.com"[..]), 7),
            InsertResult::Ok
        );

        // Exact-match cases still resolve.
        assert_eq!(
            root.domain_lookup(b"cdn1.example.com".as_ref(), false),
            Some(&(b"/cdn[0-9]+/.example.com".to_vec(), 7))
        );
        assert_eq!(
            root.domain_lookup(b"cdn123.example.com".as_ref(), false),
            Some(&(b"/cdn[0-9]+/.example.com".to_vec(), 7))
        );

        // Trailing characters past the digit run must fail. Pre-anchoring
        // the trie would have matched `cdn1xxx` because `cdn[0-9]+` ate
        // the `cdn1` prefix; with `\A...\z` the segment is rejected.
        assert_eq!(
            root.domain_lookup(b"cdn1xxx.example.com".as_ref(), false),
            None
        );
        // Leading characters likewise must fail.
        assert_eq!(
            root.domain_lookup(b"xxxcdn1.example.com".as_ref(), false),
            None
        );
        // Non-digit middle bytes break the digit run and the segment.
        assert_eq!(
            root.domain_lookup(b"cdnabc.example.com".as_ref(), false),
            None
        );
    }

    #[test]
    fn add_child_to_leaf() {
        let mut root1: TrieNode<u8> = TrieNode::root();

        println!("creating root1:");
        root1.print();
        println!("adding (abcd, 1)");
        assert_eq!(root1.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
        root1.print();
        println!("adding (abce, 2)");
        assert_eq!(root1.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);
        root1.print();
        println!("adding (abc, 3)");
        assert_eq!(root1.insert(Vec::from(&b"abc"[..]), 3), InsertResult::Ok);

        println!("root1:");
        root1.print();

        let mut root2: TrieNode<u8> = TrieNode::root();

        assert_eq!(root2.insert(Vec::from(&b"abc"[..]), 3), InsertResult::Ok);
        assert_eq!(root2.insert(Vec::from(&b"abcd"[..]), 1), InsertResult::Ok);
        assert_eq!(root2.insert(Vec::from(&b"abce"[..]), 2), InsertResult::Ok);

        println!("root2:");
        root2.print();
        assert_eq!(root2.remove(&Vec::from(&b"abc"[..])), RemoveResult::Ok);

        println!("root2 after,remove:");
        root2.print();
        let mut expected: TrieNode<u8> = TrieNode::root();

        assert_eq!(
            expected.insert(Vec::from(&b"abcd"[..]), 1),
            InsertResult::Ok
        );
        assert_eq!(
            expected.insert(Vec::from(&b"abce"[..]), 2),
            InsertResult::Ok
        );

        println!("root2 after insert");
        root2.print();
        println!("expected");
        expected.print();
        assert_eq!(root2, expected);
    }

    #[test]
    fn domains() {
        let mut root: TrieNode<u8> = TrieNode::root();
        root.print();

        assert_eq!(
            root.domain_insert(Vec::from(&b"www.example.com"[..]), 1),
            InsertResult::Ok
        );
        root.print();
        assert_eq!(
            root.domain_insert(Vec::from(&b"test.example.com"[..]), 2),
            InsertResult::Ok
        );
        root.print();
        assert_eq!(
            root.domain_insert(Vec::from(&b"*.alldomains.org"[..]), 3),
            InsertResult::Ok
        );
        root.print();
        assert_eq!(
            root.domain_insert(Vec::from(&b"alldomains.org"[..]), 4),
            InsertResult::Ok
        );
        assert_eq!(
            root.domain_insert(Vec::from(&b"pouet.alldomains.org"[..]), 5),
            InsertResult::Ok
        );
        root.print();
        assert_eq!(
            root.domain_insert(Vec::from(&b"hello.com"[..]), 6),
            InsertResult::Ok
        );
        assert_eq!(
            root.domain_insert(Vec::from(&b"*.hello.com"[..]), 7),
            InsertResult::Ok
        );
        assert_eq!(
            root.domain_insert(Vec::from(&b"images./cdn[0-9]+/.hello.com"[..]), 8),
            InsertResult::Ok
        );
        root.print();
        assert_eq!(
            root.domain_insert(Vec::from(&b"/test[0-9]+/.www.hello.com"[..]), 9),
            InsertResult::Ok
        );
        root.print();

        assert_eq!(root.domain_lookup(&b"example.com"[..], true), None);
        assert_eq!(
            root.domain_lookup(&b"blah.test.example.com"[..], true),
            None
        );
        assert_eq!(
            root.domain_lookup(&b"www.example.com"[..], true),
            Some(&(b"www.example.com"[..].to_vec(), 1))
        );
        assert_eq!(
            root.domain_lookup(&b"alldomains.org"[..], true),
            Some(&(b"alldomains.org"[..].to_vec(), 4))
        );
        assert_eq!(
            root.domain_lookup(&b"test.hello.com"[..], true),
            Some(&(b"*.hello.com"[..].to_vec(), 7))
        );
        assert_eq!(
            root.domain_lookup(&b"images.cdn10.hello.com"[..], true),
            Some(&(b"images./cdn[0-9]+/.hello.com"[..].to_vec(), 8))
        );
        assert_eq!(
            root.domain_lookup(&b"test42.www.hello.com"[..], true),
            Some(&(b"/test[0-9]+/.www.hello.com"[..].to_vec(), 9))
        );
        assert_eq!(
            root.domain_lookup(&b"test.alldomains.org"[..], true),
            Some(&(b"*.alldomains.org"[..].to_vec(), 3))
        );
        assert_eq!(
            root.domain_lookup(&b"hello.alldomains.org"[..], true),
            Some(&(b"*.alldomains.org"[..].to_vec(), 3))
        );
        assert_eq!(
            root.domain_lookup(&b"pouet.alldomains.org"[..], true),
            Some(&(b"pouet.alldomains.org"[..].to_vec(), 5))
        );
        assert_eq!(
            root.domain_lookup(&b"blah.test.alldomains.org"[..], true),
            None
        );

        assert_eq!(
            root.domain_remove(&Vec::from(&b"alldomains.org"[..])),
            RemoveResult::Ok
        );
        println!("after remove");
        root.print();
        assert_eq!(root.domain_lookup(&b"alldomains.org"[..], true), None);
        assert_eq!(
            root.domain_lookup(&b"test.alldomains.org"[..], true),
            Some(&(b"*.alldomains.org"[..].to_vec(), 3))
        );
        assert_eq!(
            root.domain_lookup(&b"hello.alldomains.org"[..], true),
            Some(&(b"*.alldomains.org"[..].to_vec(), 3))
        );
        assert_eq!(
            root.domain_lookup(&b"pouet.alldomains.org"[..], true),
            Some(&(b"pouet.alldomains.org"[..].to_vec(), 5))
        );
        assert_eq!(
            root.domain_lookup(&b"test.hello.com"[..], true),
            Some(&(b"*.hello.com"[..].to_vec(), 7))
        );
        assert_eq!(
            root.domain_lookup(&b"blah.test.alldomains.org"[..], true),
            None
        );
    }

    #[test]
    fn wildcard() {
        let mut root: TrieNode<u8> = TrieNode::root();
        root.print();
        root.domain_insert("*.clever-cloud.com".as_bytes().to_vec(), 2u8);
        root.domain_insert("services.clever-cloud.com".as_bytes().to_vec(), 0u8);
        root.domain_insert("*.services.clever-cloud.com".as_bytes().to_vec(), 1u8);

        let res = root.domain_lookup(b"test.services.clever-cloud.com", true);
        println!("query result: {res:?}");

        assert_eq!(
            root.domain_lookup(b"pgstudio.services.clever-cloud.com", true),
            Some(&("*.services.clever-cloud.com".as_bytes().to_vec(), 1u8))
        );
    }

    fn hm_insert(h: std::collections::HashMap<String, u32>) -> bool {
        let mut root: TrieNode<u32> = TrieNode::root();

        for (k, v) in h.iter() {
            if k.is_empty() {
                continue;
            }

            if k.as_bytes()[0] == b'.' {
                continue;
            }

            if k.contains('/') {
                continue;
            }

            if k == "*" {
                continue;
            }

            //println!("inserting key: '{}', value: '{}'", k, v);
            //assert_eq!(root.domain_insert(Vec::from(k.as_bytes()), *v), InsertResult::Ok);
            assert_eq!(
                root.insert(Vec::from(k.as_bytes()), *v),
                InsertResult::Ok,
                "could not insert ({k}, {v})"
            );
            //root.print();
        }

        //root.print();
        for (k, v) in h.iter() {
            if k.is_empty() {
                continue;
            }

            if k.as_bytes()[0] == b'.' {
                continue;
            }

            if k.contains('/') {
                continue;
            }

            if k == "*" {
                continue;
            }

            //match root.domain_lookup(k.as_bytes()) {
            match root.lookup(k.as_bytes(), false) {
                None => {
                    println!("did not find key '{k}'");
                    return false;
                }
                Some(&(ref k1, v1)) => {
                    if k.as_bytes() != &k1[..] || *v != v1 {
                        println!(
                            "request ({}, {}), got ({}, {})",
                            k,
                            v,
                            str::from_utf8(&k1[..]).unwrap(),
                            v1
                        );
                        return false;
                    }
                }
            }
        }

        true
    }

    /* FIXME: randomly fails
    quickcheck! {
      fn qc_insert(h: std::collections::HashMap<String, u32>) -> bool {
        hm_insert(h)
      }
    }
    */

    #[test]
    fn insert_disappearing_tree() {
        let h: std::collections::HashMap<String, u32> = [
            (String::from("\n\u{3}"), 0),
            (String::from("\n\u{0}"), 1),
            (String::from("\n"), 2),
        ]
        .iter()
        .cloned()
        .collect();
        assert!(hm_insert(h));
    }

    #[test]
    fn size() {
        assert_size!(TrieNode<u32>, 136);
    }

    /// Regression: a hostname whose LEFTMOST segment is a regex
    /// (`/test[0-9]/.example.com`) used to underflow `pos - 1` (to
    /// `usize::MAX`) and panic on the second insert (the dedup loop) and
    /// on any `lookup_mut`. Both paths now special-case `pos == 0`
    /// (regex is the leftmost/only segment → value-bearing leaf). This
    /// asserts the panic is gone and the entry resolves correctly.
    #[test]
    fn leftmost_regex_segment_reinsert_and_lookup_mut_do_not_panic() {
        let mut root: TrieNode<u8> = TrieNode::root();

        assert_eq!(
            root.insert(Vec::from(&b"/test[0-9]/.example.com"[..]), 7),
            InsertResult::Ok
        );
        // Second insert of the SAME leftmost-regex host: dedup loop with
        // pos == 0. Previously panicked; must now report Existing.
        assert_eq!(
            root.insert(Vec::from(&b"/test[0-9]/.example.com"[..]), 8),
            InsertResult::Existing
        );

        // lookup_mut on the existing leftmost-regex host: previously
        // panicked at `partial_key[..pos - 1]`; must now resolve the leaf
        // (value unchanged from the first insert — Existing did not
        // overwrite).
        let resolved = root.domain_lookup_mut(b"test4.example.com", false);
        assert_eq!(
            resolved.map(|(_, v)| *v),
            Some(7),
            "leftmost-regex host must resolve via lookup_mut without panicking",
        );

        // The immutable lookup path (never buggy) agrees.
        assert_eq!(
            root.domain_lookup(b"test4.example.com", false),
            Some(&(b"/test[0-9]/.example.com"[..].to_vec(), 7))
        );

        // Removing the last rule clears the host.
        assert_eq!(
            root.domain_remove(&Vec::from(&b"/test[0-9]/.example.com"[..])),
            RemoveResult::Ok
        );
        assert_eq!(root.domain_lookup(b"test4.example.com", false), None);
    }
}