env.rs

use std::collections::HashMap;
use std::collections::HashSet;
use std::hash::Hash;

pub struct Env<K, V> {
    table: HashMap<K, Vec<V>>,
    scope: Vec<HashSet<K>>,
    count: usize,
}

impl<K: Eq + Hash + Copy, V> Env<K, V> {
    pub fn new() -> Env<K, V> {
        Env {
            table: HashMap::new(),
            scope: vec![],
            count: 0,
        }
    }

    pub fn lookup(&self, k: &K) -> Option<&V> {
        match self.table.get(k) {
            None => None,
            Some(l) => l.last(),
        }
    }

    pub fn lookup_mut(&mut self, k: &K) -> Option<&mut V> {
        match self.table.get_mut(k) {
            None => None,
            Some(l) => l.last_mut(),
        }
    }

    pub fn insert(&mut self, k: K, v: V) {
        if self.scope[self.scope.len() - 1].contains(&k) {
            match self.table.get_mut(&k) {
                None => panic!("Internal Failure: Environment Insert"),
                Some(r) => {
                    let last = r.len() - 1;
                    r[last] = v;
                }
            }
        } else {
            let last = self.scope.len() - 1;
            self.table.entry(k).or_insert(vec![]).push(v);
            self.scope[last].insert(k);
        }
    }

    pub fn open_scope(&mut self) {
        self.scope.push(HashSet::new());
    }

    pub fn close_scope(&mut self) {
        match self.scope.pop() {
            None => assert!(false, "Internal Failure: Environment no scope to close"),
            Some(to_remove) => {
                for k in to_remove {
                    match self.table.get_mut(&k) {
                        None => {
                            assert!(false, "Internal Failure: Environment Close Scope");
                        }
                        Some(r) => {
                            r.pop();
                        }
                    }
                }
            }
        }
    }

    pub fn uniq(&mut self) -> usize {
        let n = self.count;
        self.count += 1;
        n
    }

    pub fn filter_map<F>(&mut self, mut f: F)
    where
        F: FnMut(V) -> Option<V>,
    {
        // To update the environment we have to associate values in the table with the scopes so
        // that if we delete a value from the table we delete it's note in the scope as well
        let num_scopes = self.scope.len();

        // To do this, we first construct a map from keys to a (sorted) list of the scopes that
        // have bindings of that key.
        let mut scopes: HashMap<K, Vec<usize>> = HashMap::new();
        for (scope, keys) in std::mem::take(&mut self.scope).into_iter().enumerate() {
            for k in keys {
                scopes.entry(k).or_insert(vec![]).push(scope);
            }
        }

        // Now, we can process the actual table and the scopes table in parallel since they have
        // matching structure
        let mut new_table = HashMap::new();
        let mut new_scopes: HashMap<K, Vec<usize>> = HashMap::new();

        for (k, vs) in std::mem::take(&mut self.table) {
            let scope_list = scopes.remove(&k).unwrap();
            assert!(scope_list.len() == vs.len());

            let (new_vals, new_scope_list) = vs
                .into_iter()
                .zip(scope_list.into_iter())
                .filter_map(|(v, s)| f(v).map(|v| (v, s)))
                .unzip();

            new_table.insert(k, new_vals);
            new_scopes.insert(k, new_scope_list);
        }

        // Finally we reconstruct the actual environment
        self.table = new_table;
        let mut scope: Vec<HashSet<K>> = vec![HashSet::new(); num_scopes];
        for (k, scopes) in new_scopes {
            for s in scopes {
                scope[s].insert(k);
            }
        }
        self.scope = scope;
    }
}