nixjit/evaluator/nixjit_lir/src/lib.rs

//! The Low-level Intermediate Representation (LIR) for nixjit.
//!
//! This module defines the LIR, which is a more resolved and explicit representation
//! than the HIR. The key transformation from HIR to LIR is the resolution of variable
//! lookups. In the LIR, variable references are either resolved to a specific expression,
//! a function argument, or are left as-is for dynamic lookup in a `with` environment.
//!
//! Key components:
//! - `Lir`: An enum representing all LIR expression types, generated by the `ir!` macro.
//! - `Resolve`: A trait for converting HIR nodes into LIR expressions.
//! - `ResolveContext`: A trait providing the context for resolution, including scope
//!   management and dependency tracking.

use derive_more::{IsVariant, TryUnwrap, Unwrap};

use hashbrown::HashMap;
use nixjit_error::{Error, Result};
use nixjit_hir as hir;
use nixjit_ir::*;
use nixjit_macros::ir;
use nixjit_value::format_symbol;

// The `ir!` macro generates the `Lir` enum and related structs and traits.
ir! {
    Lir,

    AttrSet,
    List,
    HasAttr,
    BinOp,
    UnOp,
    Select,
    If,
    Call,
    With,
    Assert,
    ConcatStrings,
    Const,
    Str,
    Var,
    Path,
    Arg,
    PrimOp(PrimOpId),
    StackRef(StackIdx),
    ExprRef(ExprId),
    FuncRef(ExprId),
    Thunk(ExprId),
}

/// Represents the result of a variable lookup within the `ResolveContext`.
#[derive(Debug)]
pub enum LookupResult {
    Stack(StackIdx),
    /// The variable was found and resolved to a specific expression.
    PrimOp(ExprId),
    /// The variable could not be resolved statically, likely due to a `with` expression.
    /// The lookup must be performed dynamically at evaluation time.
    Unknown,
    /// The variable was not found in any scope.
    NotFound,
}

/// A context for the HIR-to-LIR resolution process.
///
/// This trait abstracts the environment in which expressions are resolved, managing
/// scopes, dependencies, and the resolution of expressions themselves.
pub trait ResolveContext {
    /// Creates a new function, associating a parameter specification with a body expression.
    fn new_func(&mut self, body: ExprId, param: Param);

    /// Triggers the resolution of a given expression.
    fn resolve(&mut self, expr: ExprId) -> Result<()>;

    fn resolve_root(self, expr: ExprId) -> Result<()>;

    /// Looks up a variable by name in the current scope.
    fn lookup(&mut self, name: &str) -> LookupResult;

    fn lookup_arg(&mut self) -> StackIdx;

    /// Enters a `with` scope for the duration of a closure.
    fn with_with_env<T>(&mut self, f: impl FnOnce(&mut Self) -> T) -> (bool, T);

    /// Enters a `let` scope with a given set of bindings for the duration of a closure.
    fn with_let_env<T>(
        &mut self,
        bindings: HashMap<String, ExprId>,
        f: impl FnOnce(&mut Self) -> T,
    ) -> T;

    /// Enters a function parameter scope for the duration of a closure.
    fn with_closure_env<T>(
        &mut self,
        func: ExprId,
        ident: Option<String>,
        f: impl FnOnce(&mut Self) -> T,
    ) -> T;
}

/// A trait for converting (resolving) an HIR node into an LIR expression.
pub trait Resolve<Ctx: ResolveContext> {
    /// Performs the resolution.
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir>;
}

/// The main entry point for resolving any HIR expression.
impl<Ctx: ResolveContext> Resolve<Ctx> for hir::Hir {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        use hir::Hir::*;
        match self {
            AttrSet(x) => x.resolve(ctx),
            List(x) => x.resolve(ctx),
            HasAttr(x) => x.resolve(ctx),
            BinOp(x) => x.resolve(ctx),
            UnOp(x) => x.resolve(ctx),
            Select(x) => x.resolve(ctx),
            If(x) => x.resolve(ctx),
            Func(x) => x.resolve(ctx),
            Call(x) => x.resolve(ctx),
            With(x) => x.resolve(ctx),
            Assert(x) => x.resolve(ctx),
            ConcatStrings(x) => x.resolve(ctx),
            Const(x) => Ok(Lir::Const(x)),
            Str(x) => Ok(Lir::Str(x)),
            Var(x) => x.resolve(ctx),
            Path(x) => x.resolve(ctx),
            Let(x) => x.resolve(ctx),
            Thunk(x) => {
                ctx.resolve(x)?;
                Ok(Lir::Thunk(x))
            }
            Arg(_) => Ok(Lir::StackRef(ctx.lookup_arg())),
        }
    }
}

/// Resolves an `AttrSet` by resolving all key and value expressions.
impl<Ctx: ResolveContext> Resolve<Ctx> for AttrSet {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        for (_, &v) in self.stcs.iter() {
            ctx.resolve(v)?;
        }
        for &(k, v) in self.dyns.iter() {
            ctx.resolve(k)?;
            ctx.resolve(v)?;
        }
        Ok(self.to_lir())
    }
}

/// Resolves a `List` by resolving each of its items.
impl<Ctx: ResolveContext> Resolve<Ctx> for List {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        for &item in self.items.iter() {
            ctx.resolve(item)?;
        }
        Ok(self.to_lir())
    }
}

/// Resolves a `HasAttr` expression by resolving the LHS and any dynamic attributes in the path.
impl<Ctx: ResolveContext> Resolve<Ctx> for HasAttr {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        ctx.resolve(self.lhs)?;
        for attr in self.rhs.iter() {
            if let &Attr::Dynamic(expr) = attr {
                ctx.resolve(expr)?;
            }
        }
        Ok(self.to_lir())
    }
}

/// Resolves a `BinOp` by resolving its left and right hand sides.
impl<Ctx: ResolveContext> Resolve<Ctx> for BinOp {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        ctx.resolve(self.lhs)?;
        ctx.resolve(self.rhs)?;
        Ok(self.to_lir())
    }
}

/// Resolves a `UnOp` by resolving its right hand side.
impl<Ctx: ResolveContext> Resolve<Ctx> for UnOp {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        ctx.resolve(self.rhs)?;
        Ok(self.to_lir())
    }
}

/// Resolves a `Select` by resolving the expression being selected from, any dynamic
/// attributes in the path, and the default value if it exists.
impl<Ctx: ResolveContext> Resolve<Ctx> for Select {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        ctx.resolve(self.expr)?;
        for attr in self.attrpath.iter() {
            if let &Attr::Dynamic(expr) = attr {
                ctx.resolve(expr)?;
            }
        }
        if let Some(expr) = self.default {
            ctx.resolve(expr)?;
        }
        Ok(self.to_lir())
    }
}

/// Resolves an `If` expression by resolving the condition, consequence, and alternative.
impl<Ctx: ResolveContext> Resolve<Ctx> for If {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        ctx.resolve(self.cond)?;
        ctx.resolve(self.consq)?;
        ctx.resolve(self.alter)?;
        Ok(self.to_lir())
    }
}

/// Resolves a `Func` by resolving its body within a new parameter scope.
/// It then registers the function with the context.
impl<Ctx: ResolveContext> Resolve<Ctx> for Func {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        ctx.with_closure_env(self.body, self.param.ident.clone(), |ctx| {
            ctx.resolve(self.body)
        })?;
        ctx.new_func(self.body, self.param);
        Ok(Lir::FuncRef(self.body))
    }
}

impl<Ctx: ResolveContext> Resolve<Ctx> for Call {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        ctx.resolve(self.func)?;
        ctx.resolve(self.arg)?;
        Ok(self.to_lir())
    }
}

/// Resolves a `With` expression by resolving the namespace and the body.
/// The body is resolved within a special "with" scope.
impl<Ctx: ResolveContext> Resolve<Ctx> for With {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        ctx.resolve(self.namespace)?;
        let (env_used, res) = ctx.with_with_env(|ctx| ctx.resolve(self.expr));
        res?;
        // Optimization: if the `with` environment was not actually used by any variable
        // lookup in the body, we can elide the `With` node entirely.
        if env_used {
            Ok(self.to_lir())
        } else {
            Ok(Lir::ExprRef(self.expr))
        }
    }
}

/// Resolves an `Assert` by resolving the assertion condition and the body.
impl<Ctx: ResolveContext> Resolve<Ctx> for Assert {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        ctx.resolve(self.assertion)?;
        ctx.resolve(self.expr)?;
        Ok(self.to_lir())
    }
}

/// Resolves a `ConcatStrings` by resolving each part.
impl<Ctx: ResolveContext> Resolve<Ctx> for ConcatStrings {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        for &part in self.parts.iter() {
            ctx.resolve(part)?;
        }
        Ok(self.to_lir())
    }
}

/// Resolves a `Var` by looking it up in the current context.
impl<Ctx: ResolveContext> Resolve<Ctx> for Var {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        use LookupResult::*;
        match ctx.lookup(&self.sym) {
            Stack(idx) => Ok(Lir::StackRef(idx)),
            PrimOp(id) => Ok(Lir::ExprRef(id)),
            Unknown => Ok(self.to_lir()),
            NotFound => Err(Error::resolution_error(format!(
                "undefined variable '{}'",
                format_symbol(&self.sym)
            ))),
        }
    }
}

/// Resolves a `Path` by resolving the underlying expression that defines the path's content.
impl<Ctx: ResolveContext> Resolve<Ctx> for Path {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        ctx.resolve(self.expr)?;
        Ok(self.to_lir())
    }
}

/// Resolves a `Let` expression by creating a new scope for the bindings, resolving
/// the bindings and the body, and then returning a reference to the body.
impl<Ctx: ResolveContext> Resolve<Ctx> for hir::Let {
    fn resolve(self, ctx: &mut Ctx) -> Result<Lir> {
        ctx.with_let_env(self.bindings.clone(), |ctx| {
            for &id in self.bindings.values() {
                ctx.resolve(id)?;
            }
            ctx.resolve(self.body)
        })?;
        // The `let` expression itself evaluates to its body.
        Ok(Lir::ExprRef(self.body))
    }
}