nixjit/src/vm/env.rs

use std::fmt::Debug;
use std::hash::Hash;

use hashbrown::{DefaultHashBuilder, HashMap};
use bumpalo::Bump;

use crate::ty::internal::Value;

type Map<'bump, K, V> = HashMap<K, V, DefaultHashBuilder, &'bump Bump>;

#[derive(Clone)]
pub struct Env<'bump, K: Hash + Eq + Clone, V: Clone> {
    let_: &'bump LetEnv<'bump, K, V>,
    with: &'bump With<'bump, K, V>,
    last: Option<&'bump Env<'bump, K, V>>,
}

#[derive(Clone)]
pub struct LetEnv<'bump, K: Hash + Eq + Clone, V: Clone> {
    map: LetNode<'bump, K, V>,
    last: Option<&'bump LetEnv<'bump, K, V>>,
}

impl<K: Debug + Hash + Eq + Clone, V: Debug + Clone> Debug for Env<'_, K, V> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Env")
            .field("let_", &self.let_)
            .field("with", &self.with)
            .field("last", &self.last)
            .finish()
    }
}

impl<K: Debug + Hash + Eq + Clone, V: Debug + Clone> Debug for LetEnv<'_, K, V> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LetEnv")
            .field("map", &self.map)
            .field("last", &self.last)
            .finish()
    }
}

pub type VmEnv<'jit, 'vm> = Env<'vm, usize, Value<'jit, 'vm>>;

#[derive(Debug, Default, Clone)]
pub struct With<'bump, K: Hash + Eq, V> {
    map: Option<&'bump Map<'bump, K, V>>,
    last: Option<&'bump With<'bump, K, V>>,
}

#[derive(Debug, Clone)]
enum LetNode<'bump, K: Hash + Eq + Clone, V: Clone> {
    Let(&'bump Map<'bump, K, V>),
    SingleArg(K, V),
    MultiArg(&'bump Map<'bump, K, V>),
}

#[derive(Debug, Clone, Copy)]
pub enum Type {
    Arg,
    Let,
    With,
}

impl<'bump, K: Hash + Eq + Clone, V: Clone> Env<'bump, K, V> {
    pub fn new(map: &'bump Map<'bump, K, V>, bump: &'bump Bump) -> Self {
        Self {
            let_: &*bump.alloc(LetEnv::new(map)),
            with: &*bump.alloc(With {
                map: None,
                last: None,
            }),
            last: None
        }
    }

    pub fn lookup(&self, symbol: &K) -> Option<&V> {
        if let Some(val) = self.let_.lookup(symbol) {
            return Some(val);
        }
        self.with.lookup(symbol)
    }

    pub fn enter_arg(&'bump self, ident: K, val: V, bump: &'bump Bump) -> &'bump Self {
        &*bump.alloc(Env {
            let_: self.let_.enter_arg(ident, val, bump),
            with: self.with,
            last: Some(self)
        })
    }

    pub fn enter_let(&'bump self, map: &'bump Map<'bump, K, V>, bump: &'bump Bump) -> &'bump Self {
        &*bump.alloc(Env {
            let_: self.let_.enter_let(map, bump),
            with: self.with,
            last: Some(self)
        })
    }

    pub fn enter_with(&'bump self, map: &'bump Map<'bump, K, V>, bump: &'bump Bump) -> &'bump Self {
        &*bump.alloc(Env {
            let_: self.let_,
            with: self.with.enter(map, bump),
            last: Some(self)
        })
    }

    pub fn leave(&self) -> &Self {
        self.last.unwrap()
    }
}

impl<'bump, K: Hash + Eq + Clone, V: Clone> LetEnv<'bump, K, V> {
    pub fn new(map: &'bump HashMap<K, V, DefaultHashBuilder, &Bump>) -> Self {
        Self {
            map: LetNode::Let(map),
            last: None,
        }
    }

    pub fn lookup(&self, symbol: &K) -> Option<&V> {
        use self::LetNode::*;
        match &self.map {
            Let(map) | MultiArg(map) => {
                if let Some(val) = map.get(symbol) {
                    return Some(val);
                }
            }
            SingleArg(sym, val) => {
                if sym == symbol {
                    return Some(val);
                }
            }
        }
        self.last.as_ref().and_then(|env| env.lookup(symbol))
    }

    pub fn enter_arg(&'bump self, ident: K, val: V, bump: &'bump Bump) -> &'bump Self {
        &*bump.alloc(Self {
            map: LetNode::SingleArg(ident, val),
            last: Some(self)
        })
    }

    pub fn enter_let(&'bump self, map: &'bump Map<'bump, K, V>, bump: &'bump Bump) -> &'bump Self {
        &*bump.alloc(Self {
            map: LetNode::Let(map),
            last: Some(self)
        })
    }
}

impl<'bump, K: Hash + Eq + Clone, V: Clone> With<'bump, K, V> {
    pub fn lookup(&'bump self, symbol: &K) -> Option<&'bump V> {
        if let Some(val) = self.map.as_ref()?.get(symbol) {
            return Some(val);
        }
        self.last.as_ref().and_then(|env| env.lookup(symbol))
    }

    pub fn enter(&'bump self, map: &'bump Map<'bump, K, V>, bump: &'bump Bump) -> &'bump Self {
        &*bump.alloc(Self {
            map: Some(map),
            last: Some(self)
        })
    }
}