nixjit/src/ir/ctx.rs

use derive_more::Unwrap;
use ecow::EcoString;
use hashbrown::{HashMap, HashSet};

use crate::{
    error::{Error, Result},
    ty::common::Const,
};

use super::{Dep, Func, Ir, LoadFunc, MaybeThunk, SccAnalyzer, SccNode, Thunk};

#[derive(Clone, Copy, Unwrap, Debug)]
pub enum Index {
    Thunk(usize),
    Func(usize),
}

pub struct DowngradeContext {
    pub thunks: Vec<(Ir, bool)>,
    pub thunk_deps: Vec<HashSet<usize>>,
    pub func_deps: Vec<HashSet<Dep>>,
    pub func_arg_dep: Vec<bool>,
    pub funcs: Vec<Func>,
}

pub struct Env<'a, 'env> {
    env: EnvNode<'a>,
    prev: Option<&'env Env<'a, 'env>>,
}

enum EnvNode<'a> {
    Builtins(&'a HashMap<EcoString, Ir>),
    Let(&'a Vec<(EcoString, MaybeThunk)>),
    SingleArg(EcoString),
    MultiArg(HashMap<EcoString, Option<Ir>>, Option<EcoString>),
    With,
}

pub enum LookupResult {
    Builtin(Ir),
    MaybeThunk(MaybeThunk),
    Let { level: usize, idx: usize },
    SingleArg { level: usize },
    MultiArg { level: usize, default: Option<Ir> },
    With,
}

impl<'a, 'env> Env<'a, 'env> {
    pub fn new(base: &'a HashMap<EcoString, Ir>) -> Self {
        Self {
            env: EnvNode::Builtins(base),
            prev: None,
        }
    }

    pub fn enter_let(&'env self, map: &'a Vec<(EcoString, MaybeThunk)>) -> Self {
        Self {
            env: EnvNode::Let(map),
            prev: Some(self),
        }
    }

    pub fn enter_single_arg(&'env self, ident: EcoString) -> Self {
        Self {
            env: EnvNode::SingleArg(ident),
            prev: Some(self),
        }
    }

    pub fn enter_multi_arg(
        &'env self,
        map: HashMap<EcoString, Option<Ir>>,
        alias: Option<EcoString>,
    ) -> Self {
        Self {
            env: EnvNode::MultiArg(map, alias),
            prev: Some(self),
        }
    }

    pub fn enter_with(&'env self) -> Self {
        Self {
            env: EnvNode::With,
            prev: Some(self),
        }
    }

    fn _lookup(
        &self,
        ident: &EcoString,
        mut arg_level: usize,
        has_with: bool,
    ) -> core::result::Result<LookupResult, ()> {
        use EnvNode::*;
        let mut has_with = has_with;
        match &self.env {
            Builtins(map) => {
                return if let Some(ir) = map.get(ident) {
                    Ok(LookupResult::Builtin(ir.clone()))
                } else if has_with {
                    Ok(LookupResult::With)
                } else {
                    Err(())
                };
            }
            Let(map) => {
                if let Ok(idx) = map.binary_search_by(|(k, _)| k.cmp(ident)) {
                    return Ok(LookupResult::MaybeThunk(map[idx].1.clone()));
                }
            }
            SingleArg(arg) => {
                if arg == ident {
                    return Ok(LookupResult::SingleArg { level: arg_level });
                } else {
                    arg_level += 1;
                }
            }
            MultiArg(set, alias) => {
                if let Some(default) = set.get(ident) {
                    return Ok(LookupResult::MultiArg {
                        level: arg_level,
                        default: default.clone(),
                    });
                } else if alias.as_ref() == Some(ident) {
                    return Ok(LookupResult::SingleArg { level: arg_level });
                } else {
                    arg_level += 1;
                }
            }
            With => has_with = true,
        }
        self.prev
            .map(|prev| prev._lookup(ident, arg_level, has_with))
            .map_or_else(|| unreachable!(), |x| x)
    }

    pub fn lookup(&self, ident: &EcoString) -> core::result::Result<LookupResult, ()> {
        self._lookup(ident, 0, false)
    }
}

impl DowngradeContext {
    pub fn new() -> Self {
        DowngradeContext {
            thunks: Vec::new(),
            thunk_deps: Vec::new(),
            func_deps: Vec::new(),
            func_arg_dep: Vec::new(),
            funcs: Vec::new(),
        }
    }
}

impl DowngradeContext {
    pub fn new_thunk(&mut self, thunk: Ir) -> Thunk {
        let idx = self.thunks.len();
        self.thunks.push((thunk, false));
        self.thunk_deps.push(HashSet::new());
        Thunk { idx }
    }

    pub fn maybe_thunk(&mut self, ir: Ir) -> MaybeThunk {
        match ir {
            Ir::Const(cnst) => MaybeThunk::Const(cnst),
            Ir::String(string) => MaybeThunk::String(string),
            Ir::Thunk(thunk) => MaybeThunk::Thunk(thunk),
            ir => MaybeThunk::Thunk(self.new_thunk(ir)),
        }
    }

    pub fn new_dep(&mut self, this: Index, dep: Dep) -> Result<()> {
        #[cfg(debug_assertions)]
        println!("{this:?} => {dep:?}");
        match this {
            Index::Thunk(idx) => {
                if dep == Dep::Thunk(idx) && !matches!(self.thunks[idx].0, Ir::List(_)) {
                    return Err(Error::DowngradeError(
                        "infinite recursion encountered".into(),
                    ));
                }
                self.thunk_deps[idx].insert(dep.unwrap_thunk())
            }
            Index::Func(idx) => self.func_deps[idx].insert(dep),
        };
        Ok(())
    }

    pub fn new_func(&mut self, func: Func) -> LoadFunc {
        let idx = self.funcs.len();
        self.funcs.push(func);
        self.func_deps.push(HashSet::new());
        self.func_arg_dep.push(false);
        LoadFunc { idx }
    }

    pub fn resolve_func(&mut self, idx: usize, env: &Env) -> Result<()> {
        let self_ptr = self as *mut Self;
        self.funcs.get_mut(idx).map_or_else(
            || unreachable!(),
            |func| {
                unsafe {
                    let old = std::ptr::read(func);
                    match old.resolve(Index::Func(idx), self_ptr.as_mut().unwrap(), env) {
                        Ok(ok) => std::ptr::write(func, ok),
                        Err(err) => {
                            std::ptr::write(
                                func,
                                Func {
                                    param: crate::ir::Param::Ident(EcoString::new()),
                                    body: super::Const { val: Const::Null }.ir().boxed(),
                                },
                            );
                            return Err(err);
                        }
                    }
                }
                Ok(())
            },
        )
    }

    pub fn resolve_thunk(&mut self, idx: usize, env: &Env) -> Result<()> {
        if self.thunks[idx].1 {
            return Ok(());
        }
        self.thunks[idx].1 = true;
        let self_ptr = self as *mut Self;
        self.thunks.get_mut(idx).map_or_else(
            || unreachable!(),
            |thunk| {
                unsafe {
                    let (old, _) = std::ptr::read(thunk);
                    match old.resolve(Index::Thunk(idx), self_ptr.as_mut().unwrap(), env) {
                        Ok(ok) => std::ptr::write(&mut thunk.0, ok),
                        Err(err) => {
                            std::ptr::write(
                                &mut thunk.0,
                                Ir::Const(super::Const { val: Const::Null }),
                            );
                            return Err(err);
                        }
                    }
                }
                Ok(())
            },
        )
    }
}

#[derive(Debug)]
pub struct Downgraded {
    pub thunks: Box<[Ir]>,
    pub func_offset: usize,
    pub func_deps: Vec<HashSet<Dep>>,
    pub graph: Vec<SccNode>,
}

impl Downgraded {
    pub fn new(ctx: DowngradeContext) -> Self {
        Self {
            graph: SccAnalyzer::new(&ctx).analyze(),
            func_offset: ctx.thunks.len(),
            func_deps: ctx.func_deps,
            thunks: ctx
                .thunks
                .into_iter()
                .map(|(ir, _)| ir)
                .chain(ctx.funcs.into_iter().map(|Func { body, .. }| *body))
                .collect(),
        }
    }
}