nixjit/evaluator/nixjit_eval/src/value/attrset.rs

//! Defines the runtime representation of an attribute set (a map).

use core::ops::Deref;
use std::fmt::Debug;
use std::rc::Rc;

use derive_more::Constructor;
use hashbrown::HashMap;
use hashbrown::hash_map::Entry;
use itertools::Itertools;

use nixjit_error::{Error, Result};
use nixjit_ir::{ExprId, SymId};
use nixjit_value::{self as p, format_symbol};

use crate::EvalContext;

use super::Value;

/// A wrapper around a `HashMap` representing a Nix attribute set.
///
/// It uses `#[repr(transparent)]` to ensure it has the same memory layout
/// as `HashMap<String, Value>`.
#[repr(transparent)]
#[derive(Clone, Constructor)]
pub struct AttrSet {
    data: HashMap<SymId, Value>,
}

impl Debug for AttrSet {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        use Value::*;
        write!(f, "{{ ")?;
        for (k, v) in self.data.iter() {
            match v {
                List(_) => write!(f, "{:?} = [ ... ]; ", k)?,
                AttrSet(_) => write!(f, "{:?} = {{ ... }}; ", k)?,
                v => write!(f, "{:?} = {v:?}; ", k)?,
            }
        }
        write!(f, "}}")
    }
}

impl From<HashMap<SymId, Value>> for AttrSet {
    fn from(data: HashMap<SymId, Value>) -> Self {
        Self { data }
    }
}

impl Deref for AttrSet {
    type Target = HashMap<SymId, Value>;
    fn deref(&self) -> &Self::Target {
        &self.data
    }
}

impl AttrSet {
    /// Creates a new `AttrSet` with a specified initial capacity.
    pub fn with_capacity(cap: usize) -> Self {
        AttrSet {
            data: HashMap::with_capacity(cap),
        }
    }

    /// Inserts an attribute, overwriting any existing attribute with the same name.
    pub fn push_attr_force(&mut self, sym: SymId, val: Value) {
        self.data.insert(sym, val);
    }

    /// Inserts an attribute, returns an error if the attribute is already defined.
    pub fn push_attr(&mut self, sym: SymId, val: Value, ctx: &mut impl EvalContext) -> Result<()> {
        match self.data.entry(sym) {
            Entry::Occupied(occupied) => Err(Error::eval_error(format!(
                "attribute '{}' already defined",
                format_symbol(ctx.get_sym(*occupied.key()))
            ))),
            Entry::Vacant(vacant) => {
                vacant.insert(val);
                Ok(())
            }
        }
    }

    pub fn select(&self, name: SymId, ctx: &mut impl EvalContext) -> Result<Value> {
        self.data
            .get(&name)
            .cloned()
            .map(|attr| match attr {
                Value::Thunk(id) => ctx.force(id),
                val => Ok(val),
            })
            .ok_or_else(|| {
                Error::eval_error(format!("attribute '{}' not found", format_symbol(ctx.get_sym(name))))
            })?
    }

    pub fn select_or(
        &self,
        name: SymId,
        default: ExprId,
        ctx: &mut impl EvalContext,
    ) -> Result<Value> {
        self.data
            .get(&name)
            .map(|attr| match attr {
                &Value::Thunk(id) => ctx.force(id),
                val => Ok(val.clone()),
            })
            .unwrap_or_else(|| ctx.eval(default))
    }

    /// Checks if an attribute path exists within the set.
    pub fn has_attr(
        &self,
        mut path: impl DoubleEndedIterator<Item = Result<SymId>>,
    ) -> Result<Value> {
        let mut data = &self.data;
        let last = path.nth_back(0).unwrap();
        for item in path {
            let Some(Value::AttrSet(attrs)) = data.get(&item?)
            else {
                return Ok(Value::Bool(false));
            };
            data = attrs.as_inner();
        }
        Ok(Value::Bool(
            data.get(&last?).is_some(),
        ))
    }

    /// Merges another `AttrSet` into this one, with attributes from `other`
    /// overwriting existing ones. This corresponds to the `//` operator in Nix.
    pub fn update(&mut self, other: &Self) {
        for (k, v) in other.data.iter() {
            self.push_attr_force(k.clone(), v.clone())
        }
    }

    /// Returns a reference to the inner `HashMap`.
    pub fn as_inner(&self) -> &HashMap<SymId, Value> {
        &self.data
    }

    /// Converts an `Rc<AttrSet>` to an `Rc<HashMap<String, Value>>` without allocation.
    pub fn into_inner(self: Rc<Self>) -> Rc<HashMap<String, Value>> {
        // SAFETY: This is safe because `AttrSet` is `#[repr(transparent)]` over
        // `HashMap<String, Value>`, so `Rc<Self>` has the same layout as
        // `Rc<HashMap<String, Value>>`.
        unsafe { core::mem::transmute(self) }
    }

    /// Performs a deep equality comparison between two `AttrSet`s.
    ///
    /// It recursively compares the contents of both sets, ensuring that both keys
    /// and values are identical. The attributes are sorted before comparison to
    /// ensure a consistent result.
    pub fn eq_impl(&self, other: &Self) -> bool {
        self.data.iter().len() == other.data.iter().len()
            && std::iter::zip(
                self.data.iter().sorted_by(|(a, _), (b, _)| a.cmp(b)),
                other.data.iter().sorted_by(|(a, _), (b, _)| a.cmp(b)),
            )
            .all(|((k1, v1), (k2, v2))| k1 == k2 && v1.eq_impl(v2))
    }

    /// Converts the `AttrSet` to its public-facing representation.
    pub fn to_public(self, ctx: &mut impl EvalContext) -> p::Value {
        p::Value::AttrSet(p::AttrSet::new(
            self.data
                .into_iter()
                .map(|(sym, value)| (ctx.get_sym(sym).into(), value.to_public(ctx)))
                .collect(),
        ))
    }
}