use alloc::collections::BTreeMap;
use alloc::vec::Vec;
use intern_lang::Symbol;
pub struct SymbolTable<T> {
scopes: Vec<BTreeMap<Symbol, T>>,
}
impl<T> SymbolTable<T> {
#[must_use]
pub fn new() -> Self {
Self {
scopes: alloc::vec![BTreeMap::new()],
}
}
fn current_mut(&mut self) -> &mut BTreeMap<Symbol, T> {
let last = self.scopes.len() - 1; &mut self.scopes[last]
}
pub fn enter_scope(&mut self) {
self.scopes.push(BTreeMap::new());
}
pub fn exit_scope(&mut self) -> bool {
if self.scopes.len() > 1 {
let _ = self.scopes.pop();
true
} else {
false
}
}
pub fn scoped<R>(&mut self, f: impl FnOnce(&mut Self) -> R) -> R {
self.enter_scope();
let result = f(self);
let _ = self.exit_scope();
result
}
pub fn define(&mut self, name: Symbol, value: T) -> Option<T> {
self.current_mut().insert(name, value)
}
#[must_use]
pub fn lookup(&self, name: Symbol) -> Option<&T> {
self.scopes.iter().rev().find_map(|scope| scope.get(&name))
}
#[must_use]
pub fn lookup_local(&self, name: Symbol) -> Option<&T> {
self.scopes.last().and_then(|scope| scope.get(&name))
}
#[must_use]
pub fn is_defined(&self, name: Symbol) -> bool {
self.lookup(name).is_some()
}
#[must_use]
pub fn depth(&self) -> usize {
self.scopes.len()
}
}
impl<T> Default for SymbolTable<T> {
#[inline]
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use intern_lang::Interner;
use super::*;
fn names() -> (Interner, Symbol, Symbol, Symbol) {
let mut interner = Interner::new();
let a = interner.intern("a");
let b = interner.intern("b");
let c = interner.intern("c");
(interner, a, b, c)
}
#[test]
fn test_new_has_one_root_scope() {
let table: SymbolTable<()> = SymbolTable::new();
assert_eq!(table.depth(), 1);
}
#[test]
fn test_define_then_lookup() {
let (_i, a, _b, _c) = names();
let mut table: SymbolTable<i32> = SymbolTable::new();
assert_eq!(table.define(a, 7), None);
assert_eq!(table.lookup(a), Some(&7));
}
#[test]
fn test_redefine_returns_previous() {
let (_i, a, _b, _c) = names();
let mut table: SymbolTable<i32> = SymbolTable::new();
table.define(a, 1);
assert_eq!(table.define(a, 2), Some(1));
assert_eq!(table.lookup(a), Some(&2));
}
#[test]
fn test_inner_scope_shadows_then_restores() {
let (_i, a, _b, _c) = names();
let mut table: SymbolTable<i32> = SymbolTable::new();
table.define(a, 1);
table.enter_scope();
table.define(a, 2);
assert_eq!(table.lookup(a), Some(&2));
assert_eq!(table.lookup_local(a), Some(&2));
table.exit_scope();
assert_eq!(table.lookup(a), Some(&1));
}
#[test]
fn test_lookup_local_ignores_outer() {
let (_i, a, _b, _c) = names();
let mut table: SymbolTable<i32> = SymbolTable::new();
table.define(a, 1);
table.enter_scope();
assert_eq!(table.lookup_local(a), None);
assert_eq!(table.lookup(a), Some(&1));
}
#[test]
fn test_exit_root_is_a_noop() {
let mut table: SymbolTable<()> = SymbolTable::new();
assert!(!table.exit_scope());
assert_eq!(table.depth(), 1);
}
#[test]
fn test_define_in_inner_does_not_touch_outer() {
let (_i, a, _b, _c) = names();
let mut table: SymbolTable<i32> = SymbolTable::new();
table.define(a, 1);
table.scoped(|table| {
table.define(a, 99);
});
assert_eq!(table.lookup(a), Some(&1));
}
#[test]
fn test_unbound_lookup_is_none() {
let (_i, a, b, _c) = names();
let mut table: SymbolTable<i32> = SymbolTable::new();
table.define(a, 1);
assert_eq!(table.lookup(b), None);
assert!(!table.is_defined(b));
}
}