nixjit/evaluator/nixjit_eval/src/lib.rs

//! This module defines the core traits and logic for evaluating the LIR.
//!
//! The central components are:
//! - `EvalContext`: A trait that defines the environment and operations needed for evaluation.
//!   It manages the evaluation stack, scopes, and primop calls.
//! - `Evaluate`: A trait implemented by LIR nodes to define how they are evaluated.
//! - `Value`: An enum representing all possible values during evaluation. This is an
//!   internal representation, distinct from the public-facing `nixjit_value::Value`.

use std::rc::Rc;

use hashbrown::HashMap;

use nixjit_error::{Error, Result};
use nixjit_ir::{self as ir, ExprId, PrimOpId, StackIdx};
use nixjit_lir as lir;
use nixjit_value::{Const, format_symbol};

pub use crate::value::*;

mod value;

/// A trait defining the context in which LIR expressions are evaluated.
pub trait EvalContext {
    fn eval_root(self, expr: ExprId) -> Result<Value>;


    /// Evaluates an expression by its ID.
    fn eval(&mut self, expr: ExprId) -> Result<Value>;

    fn call(&mut self, func: ExprId, arg: Option<Value>, frame: StackFrame) -> Result<Value>;
    /// Enters a `with` scope for the duration of a closure's execution.
    fn with_with_env<T>(
        &mut self,
        namespace: Rc<HashMap<String, Value>>,
        f: impl FnOnce(&mut Self) -> T,
    ) -> T;

    /// Looks up a stack slot on the current stack frame.
    fn lookup_stack(&self, idx: StackIdx) -> &Value;

    fn capture_stack(&self) -> &StackFrame;

    /// Looks up an identifier in the current `with` scope chain.
    fn lookup_with<'a>(&'a self, ident: &str) -> Option<&'a Value>;

    /// Calls a primitive operation (builtin) by its ID.
    fn call_primop(&mut self, id: PrimOpId, args: Args) -> Result<Value>;

    fn get_primop_arity(&self, id: PrimOpId) -> usize;
}

/// A trait for types that can be evaluated within an `EvalContext`.
pub trait Evaluate<Ctx: EvalContext> {
    /// Performs the evaluation.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value>;
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ExprId {
    /// Evaluating an `ExprId` simply delegates to the context.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        ctx.eval(*self)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for lir::Lir {
    /// Evaluates an LIR node by dispatching to the specific implementation for its variant.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        use lir::Lir::*;
        match self {
            AttrSet(x) => x.eval(ctx),
            List(x) => x.eval(ctx),
            HasAttr(x) => x.eval(ctx),
            BinOp(x) => x.eval(ctx),
            UnOp(x) => x.eval(ctx),
            Select(x) => x.eval(ctx),
            If(x) => x.eval(ctx),
            Call(x) => x.eval(ctx),
            With(x) => x.eval(ctx),
            Assert(x) => x.eval(ctx),
            ConcatStrings(x) => x.eval(ctx),
            Const(x) => x.eval(ctx),
            Str(x) => x.eval(ctx),
            Var(x) => x.eval(ctx),
            Path(x) => x.eval(ctx),
            &StackRef(idx) => Ok(ctx.lookup_stack(idx).clone()),
            &ExprRef(expr) => ctx.eval(expr),
            &FuncRef(body) => Ok(Value::Closure(Closure::new(body, ctx.capture_stack().clone()).into())),
            &Arg(_) => unreachable!(),
            &PrimOp(primop) => Ok(Value::PrimOp(primop)),
            &Thunk(id) => Ok(Value::Thunk(id)),
        }
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::AttrSet {
    /// Evaluates an `AttrSet` by evaluating all its static and dynamic attributes.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let mut attrs = AttrSet::new(
            self.stcs
                .iter()
                .map(|(k, v)| {
                    let eval_result = v.eval(ctx);
                    Ok((k.clone(), eval_result?))
                })
                .collect::<Result<_>>()?,
        );
        for (k, v) in self.dyns.iter() {
            let v = v.eval(ctx)?;
            attrs.push_attr(k.eval(ctx)?.force_string_no_ctx()?, v)?;
        }
        let result = Value::AttrSet(attrs.into());
        Ok(result)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::List {
    /// Evaluates a `List` by evaluating all its items.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let items = self
            .items
            .iter()
            .map(|val| val.eval(ctx))
            .collect::<Result<Vec<_>>>()?;
        let result = Value::List(List::from(items).into());
        Ok(result)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::HasAttr {
    /// Evaluates a `HasAttr` by evaluating the LHS and the attribute path, then performing the check.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        use ir::Attr::*;
        let mut val = self.lhs.eval(ctx)?;
        val.has_attr(self.rhs.iter().map(|attr| match attr {
            Str(ident) => Ok(Value::String(ident.clone())),
            Dynamic(expr) => expr.eval(ctx),
        }))?;
        Ok(val)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::BinOp {
    /// Evaluates a `BinOp` by evaluating the LHS and RHS, then performing the operation.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        use ir::BinOpKind::*;
        let mut lhs = self.lhs.eval(ctx)?;
        if matches!((&self.kind, &lhs), (And, Value::Bool(false))) {
            return Ok(Value::Bool(false));
        } else if matches!((&self.kind, &lhs), (Or, Value::Bool(true))) {
            return Ok(Value::Bool(true));
        }
        let mut rhs = self.rhs.eval(ctx)?;
        match self.kind {
            Add => lhs.add(rhs)?,
            Sub => {
                rhs.neg()?;
                lhs.add(rhs)?;
            }
            Mul => lhs.mul(rhs)?,
            Div => lhs.div(rhs)?,
            Eq => Value::eq(&mut lhs, rhs),
            Neq => {
                Value::eq(&mut lhs, rhs);
                let _ = lhs.not();
            }
            Lt => lhs.lt(rhs)?,
            Gt => {
                rhs.lt(lhs)?;
                lhs = rhs;
            }
            Leq => {
                rhs.lt(lhs)?;
                let _ = rhs.not();
                lhs = rhs;
            }
            Geq => {
                lhs.lt(rhs)?;
                let _ = lhs.not()?;
            }
            And => lhs.and(rhs)?,
            Or => lhs.or(rhs)?,
            Impl => {
                let _ = lhs.not();
                lhs.or(rhs)?;
            }
            Con => lhs.concat(rhs)?,
            Upd => lhs.update(rhs)?,
            PipeL => lhs.call(rhs, ctx)?,
            PipeR => {
                rhs.call(lhs, ctx)?;
                lhs = rhs;
            }
        }
        Ok(lhs)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::UnOp {
    /// Evaluates a `UnOp` by evaluating the RHS and performing the operation.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        use ir::UnOpKind::*;
        let mut rhs = self.rhs.eval(ctx)?;
        match self.kind {
            Neg => {
                rhs.neg()?;
            }
            Not => {
                rhs.not()?;
            }
        };
        Ok(rhs)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Select {
    /// Evaluates a `Select` by evaluating the expression, the path, and the default value (if any),
    /// then performing the selection.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        use ir::Attr::*;
        let mut val = self.expr.eval(ctx)?;
        for attr in self.attrpath.iter() {
            let name_val;
            let name = match attr {
                Str(name) => name,
                Dynamic(expr) => {
                    name_val = expr.eval(ctx)?;
                    &*name_val.force_string_no_ctx()?
                }
            };
            if let Some(default) = self.default {
                val.select_or(name, default, ctx)
            } else {
                val.select(name, ctx)
            }?
        }
        Ok(val)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::If {
    /// Evaluates an `If` by evaluating the condition and then either the consequence or the alternative.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let cond = &self.cond.eval(ctx)?;
        let &cond = cond.try_into().map_err(|_| {
            Error::eval_error(format!(
                "if-condition must be a boolean, but got {}",
                cond.typename()
            ))
        })?;

        if cond {
            self.consq.eval(ctx)
        } else {
            self.alter.eval(ctx)
        }
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Call {
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let mut func = self.func.eval(ctx)?;
        func.call(self.arg.eval(ctx)?, ctx)?;
        Ok(func)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::With {
    /// Evaluates a `With` by evaluating the namespace, entering a `with` scope,
    /// and then evaluating the body.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let namespace = self.namespace.eval(ctx)?;
        let typename = namespace.typename();
        ctx.with_with_env(
            namespace
                .try_unwrap_attr_set()
                .map_err(|_| {
                    Error::eval_error(format!("'with' expects a set, but got {}", typename))
                })?
                .into_inner(),
            |ctx| self.expr.eval(ctx),
        )
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Assert {
    /// Evaluates an `Assert` by evaluating the condition. If true, it evaluates and
    /// returns the body; otherwise, it returns an error.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let cond = &self.assertion.eval(ctx)?;
        let &cond = cond.try_into().map_err(|_| {
            Error::eval_error(format!(
                "assertion condition must be a boolean, but got {}",
                cond.typename()
            ))
        })?;
        if cond {
            self.expr.eval(ctx)
        } else {
            Err(Error::catchable("assertion failed".into()))
        }
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::ConcatStrings {
    /// Evaluates a `ConcatStrings` by evaluating each part, coercing it to a string,
    /// and then concatenating the results.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        let mut buf = String::new();
        for part in self.parts.iter() {
            buf.push_str(&part.eval(ctx)?.force_string_no_ctx()?);
        }
        Ok(Value::String(buf))
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Str {
    /// Evaluates a `Str` literal into a `Value::String`.
    fn eval(&self, _: &mut Ctx) -> Result<Value> {
        Ok(Value::String(self.val.clone()))
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Const {
    /// Evaluates a `Const` literal into its corresponding `Value` variant.
    fn eval(&self, _: &mut Ctx) -> Result<Value> {
        let result = match self.val {
            Const::Null => Value::Null,
            Const::Int(x) => Value::Int(x),
            Const::Float(x) => Value::Float(x),
            Const::Bool(x) => Value::Bool(x),
        };
        Ok(result)
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Var {
    /// Evaluates a `Var` by looking it up in the `with` scope chain.
    /// This is for variables that could not be resolved statically.
    fn eval(&self, ctx: &mut Ctx) -> Result<Value> {
        ctx.lookup_with(&self.sym)
            .ok_or_else(|| {
                Error::eval_error(format!("undefined variable '{}'", format_symbol(&self.sym)))
            })
            .map(|val| val.clone())
    }
}

impl<Ctx: EvalContext> Evaluate<Ctx> for ir::Path {
    /// Evaluates a `Path`. (Currently a TODO).
    fn eval(&self, _ctx: &mut Ctx) -> Result<Value> {
        todo!()
    }
}