nixjit/src/ir.rs

use std::collections::HashMap;

use ecow::EcoString;
use rnix::ast::{self, Expr};

use crate::compile::*;
use crate::error::*;
use crate::ty::internal as i;

pub fn downgrade(expr: Expr) -> Result<Downgraded> {
    let mut ctx = DowngradeContext::new();
    let ir = expr.downgrade(&mut ctx)?;
    Ok(Downgraded {
        top_level: ir,
        consts: ctx.consts.into(),
        symbols: ctx.symbols,
        symmap: ctx.symmap,
        thunks: ctx.thunks.into(),
        funcs: ctx.funcs.into(),
    })
}

trait Downcast<T: Sized>
where
    Self: Sized,
{
    fn downcast_ref(&self) -> Option<&T>;
    fn downcast_mut(&mut self) -> Option<&mut T>;
}

macro_rules! ir {
    (
        $(
            $(#[$($x:tt)*])*
            $ty:ident
            =>
            {$($name:ident : $elemtype:ty),*$(,)?}
        )
        ,*$(,)?
    ) => {
        #[derive(Clone, Debug)]
        pub enum Ir {
            $(
                $ty($ty),
            )*
        }

        impl Ir {
            fn boxed(self) -> Box<Self> {
                Box::new(self)
            }
            fn ok(self) -> Result<Self> {
                Ok(self)
            }
        }

        impl Compile for Ir {
            fn compile(self, ctx: &mut Compiler) {
                match self {
                    $(Ir::$ty(ir) => ir.compile(ctx),)*
                }
            }
        }

        $(
            $(
                #[$($x)*]
            )*
            #[derive(Clone, Debug)]
            pub struct $ty {
                $(
                    pub $name : $elemtype,
                )*
            }

            impl $ty {
                pub fn ir(self) -> Ir {
                    Ir::$ty(self)
                }
            }

            impl Downcast<$ty> for Ir {
                fn downcast_ref(&self) -> Option<&$ty> {
                    match self {
                        Ir::$ty(value) => Some(value),
                        _ => None,
                    }
                }
                fn downcast_mut(&mut self) -> Option<&mut $ty> {
                    match self {
                        Ir::$ty(value) => Some(value),
                        _ => None,
                    }
                }
            }
        )*
    }
}

ir! {
    Attrs => { stcs: HashMap<usize, Ir>, dyns: Vec<DynamicAttrPair>, rec: bool },
    List  => { items: Vec<Ir> },
    HasAttr => { lhs: Box<Ir>, rhs: Vec<Attr> },
    BinOp => { lhs: Box<Ir>, rhs: Box<Ir>, kind: BinOpKind },
    UnOp  => { rhs: Box<Ir>, kind: UnOpKind },
    Select => { expr: Box<Ir>, attrpath: Vec<Attr>, default: Option<Box<Ir>> },
    If => { cond: Box<Ir>, consq: Box<Ir>, alter: Box<Ir> },
    LoadFunc => { idx: usize },
    Call => { func: Box<Ir>, args: Vec<Ir> },

    Let => { attrs: Attrs, expr: Box<Ir> },
    With => { namespace: Box<Ir>, expr: Box<Ir> },
    Assert => { assertion: Box<Ir>, expr: Box<Ir> },
    ConcatStrings => { parts: Vec<Ir> },
    Const => { idx: usize },
    Var => { sym: usize },
    #[derive(Copy)]
    Thunk => { idx: usize },
    Path => { expr: Box<Ir> },
}

#[derive(Clone, Debug)]
pub struct DynamicAttrPair(pub Ir, pub Ir);

#[derive(Default)]
pub struct DowngradeContext {
    thunks: Vec<Ir>,
    funcs: Vec<Func>,
    consts: Vec<i::Const>,
    constmap: HashMap<i::Const, usize>,
    symbols: Vec<EcoString>,
    symmap: HashMap<EcoString, usize>,
}

pub struct Downgraded {
    pub top_level: Ir,
    pub consts: Box<[i::Const]>,
    pub symbols: Vec<EcoString>,
    pub symmap: HashMap<EcoString, usize>,
    pub thunks: Box<[Ir]>,
    pub funcs: Box<[Func]>,
}

impl DowngradeContext {
    fn new() -> DowngradeContext {
        DowngradeContext::default()
    }

    fn new_thunk(&mut self, thunk: Ir) -> Thunk {
        let idx = self.thunks.len();
        self.thunks.push(thunk);
        Thunk { idx }
    }

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

    fn new_const(&mut self, cnst: i::Const) -> Const {
        if let Some(&idx) = self.constmap.get(&cnst) {
            Const { idx }
        } else {
            self.constmap.insert(cnst.clone(), self.consts.len());
            self.consts.push(cnst);
            Const { idx: self.consts.len() - 1 }
        }
    }

    fn new_sym(&mut self, sym: impl Into<EcoString>) -> usize {
        let sym = sym.into();
        if let Some(&idx) = self.symmap.get(&sym) {
            idx
        } else {
            self.symmap.insert(sym.clone(), self.symbols.len());
            self.symbols.push(sym);
            self.symbols.len() - 1
        }
    }
}

impl Attrs {
    fn _insert(&mut self, mut path: std::vec::IntoIter<Attr>, name: Attr, value: Ir) -> Result<()> {
        if let Some(attr) = path.next() {
            match attr {
                Attr::Str(ident) => {
                    if self.stcs.get(&ident).is_some() {
                        self.stcs
                            .get_mut(&ident)
                            .unwrap()
                            .downcast_mut()
                            .ok_or_else(|| {
                                Error::DowngradeError(format!(
                                    r#""{ident}" already exsists in this set"#
                                ))
                            })
                            .and_then(|attrs: &mut Attrs| attrs._insert(path, name, value))
                    } else {
                        let mut attrs = Attrs {
                            rec: false,
                            stcs: HashMap::new(),
                            dyns: Vec::new(),
                        };
                        attrs._insert(path, name, value)?;
                        assert!(self.stcs.insert(ident, attrs.ir()).is_none());
                        Ok(())
                    }
                }
                Attr::Strs(string) => {
                    let mut attrs = Attrs {
                        rec: false,
                        stcs: HashMap::new(),
                        dyns: Vec::new(),
                    };
                    attrs._insert(path, name, value)?;
                    self.dyns.push(DynamicAttrPair(string.ir(), attrs.ir()));
                    Ok(())
                }
                Attr::Dynamic(dynamic) => {
                    let mut attrs = Attrs {
                        rec: false,
                        stcs: HashMap::new(),
                        dyns: Vec::new(),
                    };
                    attrs._insert(path, name, value)?;
                    self.dyns.push(DynamicAttrPair(dynamic, attrs.ir()));
                    Ok(())
                }
            }
        } else {
            match name {
                Attr::Str(ident) => {
                    if self.stcs.get(&ident).is_some() {
                        return Err(Error::DowngradeError(format!(
                            r#""{ident}" already exsists in this set"#
                        )));
                    }
                    self.stcs.insert(ident, value);
                }
                Attr::Strs(string) => {
                    self.dyns.push(DynamicAttrPair(string.ir(), value));
                }
                Attr::Dynamic(dynamic) => {
                    self.dyns.push(DynamicAttrPair(dynamic, value));
                }
            }
            Ok(())
        }
    }

    pub fn insert(&mut self, path: Vec<Attr>, value: Ir) -> Result<()> {
        let mut path = path.into_iter();
        let name = path.next_back().unwrap();
        self._insert(path, name, value)
    }

    fn _has_attr(&self, mut path: std::slice::Iter<Attr>, name: Attr) -> Option<bool> {
        match path.next() {
            Some(Attr::Str(ident)) => self
                .stcs
                .get(ident)
                .and_then(|attrs| attrs.downcast_ref())
                .map_or(Some(false), |attrs: &Attrs| attrs._has_attr(path, name)),
            None => match name {
                Attr::Str(ident) => Some(self.stcs.get(&ident).is_some()),
                _ => None,
            },
            _ => None,
        }
    }

    pub fn has_attr(&self, path: &[Attr]) -> Option<bool> {
        let mut path = path.iter();
        let name = path.next_back().unwrap().clone();
        self._has_attr(path, name)
    }
}

#[derive(Clone, Debug)]
pub enum Attr {
    Dynamic(Ir),
    Strs(ConcatStrings),
    Str(usize),
}

#[derive(Clone, Debug)]
pub enum BinOpKind {
    Add,
    Sub,
    Div,
    Mul,
    Eq,
    Neq,
    Lt,
    Gt,
    Leq,
    Geq,
    And,
    Or,
    Impl,

    Con,
    Upd,

    PipeL,
    PipeR,
}

impl From<ast::BinOpKind> for BinOpKind {
    fn from(op: ast::BinOpKind) -> Self {
        use BinOpKind::*;
        use ast::BinOpKind as astkind;
        match op {
            astkind::Concat => Con,
            astkind::Update => Upd,
            astkind::Add => Add,
            astkind::Sub => Sub,
            astkind::Mul => Mul,
            astkind::Div => Div,
            astkind::And => And,
            astkind::Equal => Eq,
            astkind::Implication => Impl,
            astkind::Less => Lt,
            astkind::LessOrEq => Leq,
            astkind::More => Gt,
            astkind::MoreOrEq => Geq,
            astkind::NotEqual => Neq,
            astkind::Or => Or,
            astkind::PipeLeft => PipeL,
            astkind::PipeRight => PipeR,
        }
    }
}

#[derive(Clone, Debug)]
pub enum UnOpKind {
    Neg,
    Not,
}

impl From<ast::UnaryOpKind> for UnOpKind {
    fn from(value: ast::UnaryOpKind) -> Self {
        match value {
            ast::UnaryOpKind::Invert => UnOpKind::Not,
            ast::UnaryOpKind::Negate => UnOpKind::Neg,
        }
    }
}

pub struct Func {
    pub param: Param,
    pub body: Box<Ir>,
}

#[derive(Clone, Debug)]
pub enum Param {
    Ident(usize),
    Formals {
        formals: Vec<(usize, Option<Thunk>)>,
        ellipsis: bool,
        alias: Option<usize>,
    },
}

trait Downgrade
where
    Self: Sized,
{
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir>;
}

impl Downgrade for Expr {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        match self {
            Expr::Apply(apply) => apply.downgrade(ctx),
            Expr::Assert(assert) => assert.downgrade(ctx),
            Expr::Error(error) => Err(Error::DowngradeError(error.to_string())),
            Expr::IfElse(ifelse) => ifelse.downgrade(ctx),
            Expr::Select(select) => select.downgrade(ctx),
            Expr::Str(str) => str.downgrade(ctx),
            Expr::Path(path) => path.downgrade(ctx),
            Expr::Literal(lit) => lit.downgrade(ctx),
            Expr::Lambda(lambda) => lambda.downgrade(ctx),
            Expr::LegacyLet(let_) => let_.downgrade(ctx),
            Expr::LetIn(letin) => letin.downgrade(ctx),
            Expr::List(list) => list.downgrade(ctx),
            Expr::BinOp(op) => op.downgrade(ctx),
            Expr::Paren(paren) => paren.expr().unwrap().downgrade(ctx),
            Expr::Root(root) => root.expr().unwrap().downgrade(ctx),
            Expr::AttrSet(attrs) => attrs.downgrade(ctx),
            Expr::UnaryOp(op) => op.downgrade(ctx),
            Expr::Ident(ident) => ident.downgrade(ctx),
            Expr::With(with) => with.downgrade(ctx),
            Expr::HasAttr(has) => has.downgrade(ctx),
        }
    }
}

impl Downgrade for ast::Assert {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        Assert {
            assertion: self.condition().unwrap().downgrade(ctx)?.boxed(),
            expr: self.body().unwrap().downgrade(ctx)?.boxed(),
        }
        .ir()
        .ok()
    }
}

impl Downgrade for ast::IfElse {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        If {
            cond: self.condition().unwrap().downgrade(ctx)?.boxed(),
            consq: self.body().unwrap().downgrade(ctx)?.boxed(),
            alter: self.else_body().unwrap().downgrade(ctx)?.boxed(),
        }
        .ir()
        .ok()
    }
}

impl Downgrade for ast::Path {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        let parts = self
            .parts()
            .map(|part| match part {
                ast::InterpolPart::Literal(lit) => ctx.new_const(lit.to_string().into())
                .ir()
                .ok(),
                ast::InterpolPart::Interpolation(interpol) => {
                    interpol.expr().unwrap().downgrade(ctx)
                }
            })
            .collect::<Result<Vec<_>>>()?;
        if parts.len() == 1 {
            Path {
                expr: parts.into_iter().next().unwrap().boxed(),
            }
        } else {
            Path {
                expr: ConcatStrings { parts }.ir().boxed(),
            }
        }
        .ir()
        .ok()
    }
}

impl Downgrade for ast::Str {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        let parts = self
            .normalized_parts()
            .into_iter()
            .map(|part| match part {
                ast::InterpolPart::Literal(lit) => ctx.new_const(lit.into()).ir().ok(),
                ast::InterpolPart::Interpolation(interpol) => {
                    interpol.expr().unwrap().downgrade(ctx)
                }
            })
            .collect::<Result<Vec<_>>>()?;
        if parts.len() == 1 {
            Ok(parts.into_iter().next().unwrap())
        } else {
            ConcatStrings { parts }.ir().ok()
        }
    }
}

impl Downgrade for ast::Literal {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        match self.kind() {
            ast::LiteralKind::Integer(int) => ctx.new_const(int.value().unwrap().into()),
            ast::LiteralKind::Float(float) => ctx.new_const(float.value().unwrap().into()),
            ast::LiteralKind::Uri(uri) => ctx.new_const(uri.to_string().into())
        }
        .ir()
        .ok()
    }
}

impl Downgrade for ast::Ident {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        Var {
            sym: ctx.new_sym(self.to_string()),
        }
        .ir()
        .ok()
    }
}

impl Downgrade for ast::AttrSet {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        let rec = self.rec_token().is_some();
        downgrade_has_entry(self, rec, ctx).map(|attrs| attrs.ir())
    }
}

impl Downgrade for ast::List {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        let mut items = Vec::with_capacity(self.items().size_hint().0);
        for item in self.items() {
            items.push(item.downgrade(ctx)?)
        }
        List { items }.ir().ok()
    }
}

impl Downgrade for ast::BinOp {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        BinOp {
            lhs: self.lhs().unwrap().downgrade(ctx)?.boxed(),
            rhs: self.rhs().unwrap().downgrade(ctx)?.boxed(),
            kind: self.operator().unwrap().into(),
        }
        .ir()
        .ok()
    }
}

impl Downgrade for ast::HasAttr {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        let attrs = self.expr().unwrap().downgrade(ctx)?;
        let path = downgrade_attrpath(self.attrpath().unwrap(), ctx)?;
        HasAttr {
            lhs: attrs.boxed(),
            rhs: path,
        }
        .ir()
        .ok()
    }
}

impl Downgrade for ast::UnaryOp {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        UnOp {
            rhs: self.expr().unwrap().downgrade(ctx)?.boxed(),
            kind: self.operator().unwrap().into(),
        }
        .ir()
        .ok()
    }
}

impl Downgrade for ast::Select {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        Select {
            expr: self.expr().unwrap().downgrade(ctx)?.boxed(),
            attrpath: downgrade_attrpath(self.attrpath().unwrap(), ctx)?,
            default: match self.default_expr() {
                Some(default) => Some(default.downgrade(ctx)?.boxed()),
                None => None,
            },
        }
        .ir()
        .ok()
    }
}

impl Downgrade for ast::LegacyLet {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        let attrs = downgrade_has_entry(self, true, ctx)?;
        Select {
            expr: attrs.ir().boxed(),
            attrpath: vec![Attr::Str(ctx.new_sym("body".to_string()))],
            default: None,
        }
        .ir()
        .ok()
    }
}

impl Downgrade for ast::LetIn {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        let body = self.body().unwrap();
        let attrs = downgrade_has_entry(self, true, ctx)?;
        let expr = body.downgrade(ctx)?.boxed();
        Let { attrs, expr }.ir().ok()
    }
}

impl Downgrade for ast::With {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        let namespace = self.namespace().unwrap().downgrade(ctx)?;
        if let Ir::Attrs(attrs) = namespace {
            let expr = self.body().unwrap().downgrade(ctx)?.boxed();
            Let { attrs, expr }.ir().ok()
        } else {
            let namespace = namespace.boxed();
            let expr = self.body().unwrap().downgrade(ctx)?.boxed();
            With { namespace, expr }.ir().ok()
        }
    }
}

impl Downgrade for ast::Lambda {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        let body = self.body().unwrap();
        let param = downgrade_param(self.param().unwrap(), ctx)?;
        let body = body.downgrade(ctx)?.boxed();
        ctx.new_func(Func { param, body }).ir().ok()
    }
}

impl Downgrade for ast::Apply {
    fn downgrade(self, ctx: &mut DowngradeContext) -> Result<Ir> {
        let mut args = vec![self.argument().unwrap().downgrade(ctx)?];
        let mut func = self.lambda().unwrap();
        while let ast::Expr::Apply(call) = func {
            func = call.lambda().unwrap();
            args.push(call.argument().unwrap().downgrade(ctx)?);
        }
        let func = func.downgrade(ctx)?.boxed();
        args.reverse();
        Call { func, args }.ir().ok()
    }
}

fn downgrade_param(param: ast::Param, ctx: &mut DowngradeContext) -> Result<Param> {
    match param {
        ast::Param::IdentParam(ident) => Ok(Param::Ident(ctx.new_sym(ident.to_string()))),
        ast::Param::Pattern(pattern) => downgrade_pattern(pattern, ctx),
    }
}

fn downgrade_pattern(pattern: ast::Pattern, ctx: &mut DowngradeContext) -> Result<Param> {
    let formals = pattern
        .pat_entries()
        .map(|entry| {
            let ident = ctx.new_sym(entry.ident().unwrap().to_string());
            if entry.default().is_none() {
                Ok((ident, None))
            } else {
                entry
                    .default()
                    .unwrap()
                    .downgrade(ctx)
                    .map(|ok| (ident, Some(ctx.new_thunk(ok))))
            }
        })
        .collect::<Result<Vec<_>>>()?;
    let ellipsis = pattern.ellipsis_token().is_some();
    let alias = pattern
        .pat_bind()
        .map(|alias| ctx.new_sym(alias.ident().unwrap().to_string()));
    Ok(Param::Formals {
        formals,
        ellipsis,
        alias,
    })
}

fn downgrade_has_entry(
    has_entry: impl ast::HasEntry,
    rec: bool,
    ctx: &mut DowngradeContext,
) -> Result<Attrs> {
    let entires = has_entry.entries();
    let mut attrs = Attrs {
        rec,
        stcs: HashMap::new(),
        dyns: Vec::new(),
    };

    for entry in entires {
        match entry {
            ast::Entry::Inherit(inherit) => downgrade_inherit(inherit, &mut attrs.stcs, ctx)?,
            ast::Entry::AttrpathValue(value) => downgrade_attrpathvalue(value, &mut attrs, ctx)?,
        }
    }

    Ok(attrs)
}

fn downgrade_inherit(
    inherit: ast::Inherit,
    stcs: &mut HashMap<usize, Ir>,
    ctx: &mut DowngradeContext,
) -> Result<()> {
    let from = if let Some(from) = inherit.from() {
        let from = from.expr().unwrap().downgrade(ctx)?;
        Some(ctx.new_thunk(from))
    } else {
        None
    };
    for attr in inherit.attrs() {
        let ident = match downgrade_attr(attr, ctx)? {
            Attr::Str(ident) => ctx.new_sym(ident.to_string()),
            _ => {
                return Err(Error::DowngradeError(
                    "dynamic attributes not allowed in inherit".to_string(),
                ));
            }
        };
        let expr = from.map_or_else(
            || Var { sym: ident }.ir().ok(),
            |from| {
                Ok(Select {
                    expr: from.ir().boxed(),
                    attrpath: vec![Attr::Str(ident)],
                    default: None,
                }
                .ir())
            },
        )?;
        assert!(stcs.insert(ident, expr).is_none());
    }
    Ok(())
}

fn downgrade_attr(attr: ast::Attr, ctx: &mut DowngradeContext) -> Result<Attr> {
    use ast::Attr::*;
    use ast::InterpolPart::*;
    match attr {
        Ident(ident) => Ok(Attr::Str(ctx.new_sym(ident.to_string()))),
        Str(string) => {
            let parts = string.normalized_parts();
            if parts.len() == 0 {
                Ok(Attr::Str(ctx.new_sym("")))
            } else if parts.len() == 1 {
                match parts.into_iter().next().unwrap() {
                    Literal(ident) => Ok(Attr::Str(ctx.new_sym(ident))),
                    Interpolation(interpol) => {
                        Ok(Attr::Dynamic(interpol.expr().unwrap().downgrade(ctx)?))
                    }
                }
            } else {
                let parts = parts
                    .into_iter()
                    .map(|part| match part {
                        Literal(lit) => ctx.new_const(lit.into()).ir().ok(),
                        Interpolation(interpol) => interpol.expr().unwrap().downgrade(ctx),
                    })
                    .collect::<Result<Vec<_>>>()?;
                Ok(Attr::Strs(ConcatStrings { parts }))
            }
        }
        Dynamic(dynamic) => Ok(Attr::Dynamic(dynamic.expr().unwrap().downgrade(ctx)?)),
    }
}

fn downgrade_attrpath(attrpath: ast::Attrpath, ctx: &mut DowngradeContext) -> Result<Vec<Attr>> {
    attrpath
        .attrs()
        .map(|attr| downgrade_attr(attr, ctx))
        .collect::<Result<Vec<_>>>()
}

fn downgrade_attrpathvalue(
    value: ast::AttrpathValue,
    attrs: &mut Attrs,
    ctx: &mut DowngradeContext,
) -> Result<()> {
    let path = downgrade_attrpath(value.attrpath().unwrap(), ctx)?;
    let value = value.value().unwrap().downgrade(ctx)?;
    attrs.insert(path, ctx.new_thunk(value).ir())
}