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

use radix_trie::{Trie, TrieCommon, TrieKey};
use super::{Scope, FQSN};
use std::hash::{Hasher, Hash};
use std::collections::hash_map::DefaultHasher;

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

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() {
      if let Scope::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) {
    self.0.insert(fqsn.clone(), ());
  }

  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).map(|fqsn| fqsn.clone()).collect();
    output
  }
}

#[cfg(test)]
mod test {
  use super::*;
  use crate::symbol_table::{Scope, FQSN};

  fn make_fqsn(strs: &[&str]) -> FQSN {
    let mut scopes = vec![];
    for s in strs {
      scopes.push(Scope::Name(s.to_string()));
    }
    FQSN {
      scopes
    }
  }

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

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

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