Solution to "Implement Trie (Prefix Tree)" Leetcode problem

main.rs

#[derive(Default)]
struct AlphabetMap {
    values: [usize; Self::LEN],
    flags: u32,
}

impl AlphabetMap {
    const LEN: usize = (b'z' - b'a' + 1) as usize;

    fn get(&self, idx: u8) -> Option<usize> {
        let idx = usize::from(idx);
        if self.is_present(idx) {
            Some(self.values[idx])
        } else {
            None
        }
    }

    fn is_present(&self, idx: usize) -> bool {
        self.flags & (1 << idx) != 0
    }

    fn has_end(&self) -> bool {
        self.flags & (1 << Self::LEN) != 0
    }

    fn set_end(&mut self) {
        self.flags |= 1 << Self::LEN;
    }

    fn set(&mut self, idx: u8, val: usize) {
        self.flags |= 1 << idx;
        self.values[usize::from(idx)] = val;
    }
}

use std::fmt::{self, Debug};

impl Debug for AlphabetMap {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        struct DebugEntries<'a>(&'a AlphabetMap);

        impl Debug for DebugEntries<'_> {
            fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                f.debug_map()
                    .entries(
                        self.0.values
                            .iter()
                            .zip(0u8..)
                            .filter_map(|(&idx, ch)| {
                                if self.0.is_present(ch.into()) {
                                    Some((char::from(ch + b'a'), idx))
                                } else {
                                    None
                            })
                    )
                    .finish()
            }
        }

        let entries = DebugEntries(self);

        f.debug_struct("AplhabetMap")
            .field("has_end", &self.has_end())
            .field("entries", &entries)
            .finish()
    }
}

#[derive(Default, Debug)]
struct Trie {
    first: AlphabetMap,
    edges: Vec<AlphabetMap>,
}

impl Trie {
    fn new() -> Self {
        <_>::default()
    }

    fn new_edge(&mut self) -> usize {
        let idx = self.edges.len();
        self.edges.push(<_>::default());
        idx
    }

    fn insert(&mut self, word: String) {
        self.insert_(&word)
    }

    fn insert_(&mut self, word: &str) {
        let mut bytes = word.as_bytes().iter().map(|&ch| ch - b'a');
        let first = bytes.next().expect("word is empty");
        let mut idx = match self.first.get(first) {
            Some(idx) => idx,
            None => {
                let new_idx = self.new_edge();
                self.first.set(first, new_idx);
                new_idx
            }
        };

        let not_inserted = loop {
            let ch = match bytes.next() {
                Some(ch) => ch,
                None => {
                    self.edges[idx].set_end();
                    return;
                }
            };
            idx = match self.edges[idx].get(ch) {
                Some(idx) => idx,
                None => break ch,
            };
        };

        let new_idx = self.new_edge();
        self.edges[idx].set(not_inserted, new_idx);
        idx = new_idx;

        self.edges.reserve(bytes.len());
        for ch in bytes {
            let new_idx = self.new_edge();
            self.edges[idx].set(ch, new_idx);
            idx = new_idx;
        }

        self.edges[idx].set_end();
    }

    fn search(&self, word: String) -> bool {
        self.search_(&word)
    }

    fn search_(&self, word: &str) -> bool {
        let mut bytes = word.as_bytes().iter().map(|&ch| ch - b'a');
        let first = bytes.next().expect("word is empty");
        let mut idx = match self.first.get(first) {
            Some(idx) => idx,
            None => return false,
        };

        for ch in bytes {
            idx = match self.edges[idx].get(ch) {
                Some(idx) => idx,
                None => return false,
            }
        }

        self.edges[idx].has_end()
    }

    fn starts_with(&self, word: String) -> bool {
        self.starts_with_(&word)
    }

    fn starts_with_(&self, word: &str) -> bool {
        let mut bytes = word.as_bytes().iter().map(|&ch| ch - b'a');
        let first = bytes.next().expect("word is empty");
        let mut idx = match self.first.get(first) {
            Some(idx) => idx,
            None => return false,
        };

        for ch in bytes {
            idx = match self.edges[idx].get(ch) {
                Some(idx) => idx,
                None => return false,
            }
        }

        true
    }
}

fn main() {
    let mut trie = Trie::new();
    trie.insert_("apple");
    trie.insert_("app");
    println!("{:?}", trie.starts_with_("ap"));
    println!("{:?}", trie.starts_with_("app"));
    println!("{:?}", trie.search_("ap"));
    println!("{:?}", trie.search_("app"));
    println!("{:#?}", trie);
}