nixjit/evaluator/nixjit_hir/src/utils.rs

//! This module provides utility functions for the AST-to-HIR downgrade process.
//! These functions handle common, often complex, patterns in the `rnix` AST,
//! such as parsing parameters, attribute sets, and `inherit` statements.
//! They are helpers to the main `Downgrade` trait implementations.

use hashbrown::HashMap;

use rnix::ast;

use nixjit_error::{Error, Result};
use nixjit_ir::{Attr, AttrSet, ConcatStrings, ExprId, Select, Str, Var};

use super::downgrade::Downgrade;
use super::{Attrs, DowngradeContext, Param, ToHir};

/// Downgrades a function parameter from the AST.
pub fn downgrade_param(param: ast::Param, ctx: &mut impl DowngradeContext) -> Result<Param> {
    match param {
        ast::Param::IdentParam(ident) => Ok(Param::Ident(ident.to_string())),
        ast::Param::Pattern(pattern) => downgrade_pattern(pattern, ctx),
    }
}

/// Downgrades a parameter pattern (formals) from the AST.
/// This handles `{ a, b ? 2, ... }@args` style parameters.
pub fn downgrade_pattern(pattern: ast::Pattern, ctx: &mut impl DowngradeContext) -> Result<Param> {
    // Extract each formal parameter, downgrading its default value if it exists.
    let formals = pattern
        .pat_entries()
        .map(|entry| {
            let ident = entry.ident().unwrap().to_string();
            if entry.default().is_none() {
                Ok((ident, None))
            } else {
                entry
                    .default()
                    .unwrap()
                    .downgrade(ctx)
                    .map(|ok| (ident, Some(ok)))
            }
        })
        .collect::<Result<Vec<_>>>()?;
    let ellipsis = pattern.ellipsis_token().is_some();
    let alias = pattern
        .pat_bind()
        .map(|alias| alias.ident().unwrap().to_string());
    Ok(Param::Formals {
        formals,
        ellipsis,
        alias,
    })
}

/// Downgrades the entries of an attribute set.
/// This handles `inherit` and `attrpath = value;` entries.
pub fn downgrade_attrs(
    attrs: impl ast::HasEntry,
    ctx: &mut impl DowngradeContext,
) -> Result<AttrSet> {
    let entries = attrs.entries();
    let mut attrs = AttrSet {
        stcs: HashMap::new(),
        dyns: Vec::new(),
        rec: false,
    };

    for entry in entries {
        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)
}

/// Downgrades attribute set entries for a `let...in` expression.
/// This is a stricter version of `downgrade_attrs` that disallows dynamic attributes,
/// as `let` bindings must be statically known.
pub fn downgrade_static_attrs(
    attrs: impl ast::HasEntry,
    ctx: &mut impl DowngradeContext,
) -> Result<HashMap<String, ExprId>> {
    let entries = attrs.entries();
    let mut attrs = AttrSet {
        stcs: HashMap::new(),
        dyns: Vec::new(),
        rec: false,
    };

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

    Ok(attrs.stcs)
}

/// Downgrades an `inherit` statement.
/// `inherit (from) a b;` is translated into `a = from.a; b = from.b;`.
/// `inherit a b;` is translated into `a = a; b = b;` (i.e., bringing variables into scope).
pub fn downgrade_inherit(
    inherit: ast::Inherit,
    stcs: &mut HashMap<String, ExprId>,
    ctx: &mut impl DowngradeContext,
) -> Result<()> {
    // Downgrade the `from` expression if it exists.
    let from = if let Some(from) = inherit.from() {
        Some(from.expr().unwrap().downgrade(ctx)?)
    } else {
        None
    };
    for attr in inherit.attrs() {
        let ident = match downgrade_attr(attr, ctx)? {
            Attr::Str(ident) => ident,
            _ => {
                // `inherit` does not allow dynamic attributes.
                return Err(Error::DowngradeError(
                    "dynamic attributes not allowed in inherit".to_string(),
                ));
            }
        };
        let expr = from.as_ref().map_or_else(
            // If `from` is None, `inherit foo;` becomes `foo = foo;`.
            || Var { sym: ident.clone() }.to_hir(),
            // If `from` is Some, `inherit (from) foo;` becomes `foo = from.foo;`.
            |expr| {
                Select {
                    expr: unsafe { expr.clone() },
                    attrpath: vec![Attr::Str(ident.clone())],
                    default: None,
                }
                .to_hir()
            },
        );
        if stcs.insert(ident, ctx.new_expr(expr)).is_some() {
            // TODO: Handle or error on duplicate attribute definitions.
            todo!()
        }
    }
    Ok(())
}

/// Downgrades a single attribute key (part of an attribute path).
/// An attribute can be a static identifier, an interpolated string, or a dynamic expression.
pub fn downgrade_attr(attr: ast::Attr, ctx: &mut impl DowngradeContext) -> Result<Attr> {
    use ast::Attr::*;
    use ast::InterpolPart::*;
    match attr {
        Ident(ident) => Ok(Attr::Str(ident.to_string())),
        Str(string) => {
            let parts = string.normalized_parts();
            if parts.is_empty() {
                Ok(Attr::Str("".into()))
            } else if parts.len() == 1 {
                // If the string has only one part, it's either a literal or a single interpolation.
                match parts.into_iter().next().unwrap() {
                    Literal(ident) => Ok(Attr::Str(ident)),
                    Interpolation(interpol) => {
                        Ok(Attr::Dynamic(interpol.expr().unwrap().downgrade(ctx)?))
                    }
                }
            } else {
                // If the string has multiple parts, it's an interpolated string that must be concatenated.
                let parts = parts
                    .into_iter()
                    .map(|part| match part {
                        Literal(lit) => Ok(ctx.new_expr(self::Str { val: lit }.to_hir())),
                        Interpolation(interpol) => interpol.expr().unwrap().downgrade(ctx),
                    })
                    .collect::<Result<Vec<_>>>()?;
                Ok(Attr::Dynamic(
                    ctx.new_expr(ConcatStrings { parts }.to_hir()),
                ))
            }
        }
        Dynamic(dynamic) => Ok(Attr::Dynamic(dynamic.expr().unwrap().downgrade(ctx)?)),
    }
}

/// Downgrades an attribute path (e.g., `a.b."${c}".d`) into a `Vec<Attr>`.
pub fn downgrade_attrpath(
    attrpath: ast::Attrpath,
    ctx: &mut impl DowngradeContext,
) -> Result<Vec<Attr>> {
    attrpath
        .attrs()
        .map(|attr| downgrade_attr(attr, ctx))
        .collect::<Result<Vec<_>>>()
}

/// Downgrades an `attrpath = value;` expression and inserts it into an `AttrSet`.
pub fn downgrade_attrpathvalue(
    value: ast::AttrpathValue,
    attrs: &mut AttrSet,
    ctx: &mut impl DowngradeContext,
) -> Result<()> {
    let path = downgrade_attrpath(value.attrpath().unwrap(), ctx)?;
    let value = value.value().unwrap().downgrade(ctx)?;
    attrs.insert(path, value, ctx)
}

/// A stricter version of `downgrade_attrpathvalue` for `let...in` bindings.
/// It ensures that the attribute path contains no dynamic parts.
pub fn downgrade_static_attrpathvalue(
    value: ast::AttrpathValue,
    attrs: &mut AttrSet,
    ctx: &mut impl DowngradeContext,
) -> Result<()> {
    let path = downgrade_attrpath(value.attrpath().unwrap(), ctx)?;
    if path.iter().any(|attr| matches!(attr, Attr::Dynamic(_))) {
        return Err(Error::DowngradeError(
            "dynamic attributes not allowed in let bindings".to_string(),
        ));
    }
    let value = value.value().unwrap().downgrade(ctx)?;
    attrs.insert(path, value, ctx)
}