schala/schala-lang/src/symbol_table/symbol_trie.rs

use std::{
    collections::hash_map::DefaultHasher,
    hash::{Hash, Hasher},
};

use radix_trie::{Trie, TrieCommon, TrieKey};

use super::{DefId, Fqsn, ScopeSegment};

#[derive(Debug)]
pub struct SymbolTrie(Trie<Fqsn, DefId>);

impl TrieKey for Fqsn {
    fn encode_bytes(&self) -> Vec<u8> {
        let mut hasher = DefaultHasher::new();
        let mut output = vec![];
        for segment in self.scopes.iter() {
            let ScopeSegment::Name(s) = segment;
            s.as_bytes().hash(&mut hasher);
            output.extend_from_slice(&hasher.finish().to_be_bytes());
        }
        output
    }
}

impl SymbolTrie {
    pub fn new() -> SymbolTrie {
        SymbolTrie(Trie::new())
    }

    pub fn insert(&mut self, fqsn: &Fqsn, def_id: DefId) {
        self.0.insert(fqsn.clone(), def_id);
    }

    pub fn lookup(&self, fqsn: &Fqsn) -> Option<DefId> {
        self.0.get(fqsn).cloned()
    }

    pub fn get_children(&self, fqsn: &Fqsn) -> Vec<Fqsn> {
        let subtrie = match self.0.subtrie(fqsn) {
            Some(s) => s,
            None => return vec![],
        };
        let output: Vec<Fqsn> = subtrie.keys().filter(|cur_key| **cur_key != *fqsn).cloned().collect();
        output
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::symbol_table::Fqsn;

    fn make_fqsn(strs: &[&str]) -> Fqsn {
        Fqsn::from_strs(strs)
    }

    #[test]
    fn test_trie_insertion() {
        let id = DefId::default();
        let mut trie = SymbolTrie::new();

        trie.insert(&make_fqsn(&["unrelated", "thing"]), id);
        trie.insert(&make_fqsn(&["outer", "inner"]), id);
        trie.insert(&make_fqsn(&["outer", "inner", "still_inner"]), id);

        let children = trie.get_children(&make_fqsn(&["outer", "inner"]));
        assert_eq!(children.len(), 1);
    }
}