nix-js/fix-ir/src/lib.rs

use std::hash::Hash;
use std::ops::Deref;

use bumpalo::Bump;
use bumpalo::boxed::Box;
use bumpalo::collections::Vec;
use fix_builtins::{BUILTINS, BuiltinId};
use fix_common::StringId;
use ghost_cell::{GhostCell, GhostToken};
use num_enum::TryFromPrimitive as _;
use rnix::{TextRange, ast};
use string_interner::DefaultStringInterner;

pub mod downgrade;

pub type HashMap<'ir, K, V> = hashbrown::HashMap<K, V, hashbrown::DefaultHashBuilder, &'ir Bump>;

#[repr(transparent)]
#[derive(Clone, Copy)]
pub struct IrRef<'id, 'ir>(&'ir GhostCell<'id, Ir<'ir, Self>>);

impl<'id, 'ir> IrRef<'id, 'ir> {
    pub fn new(ir: &'ir GhostCell<'id, Ir<'ir, Self>>) -> Self {
        Self(ir)
    }

    pub fn alloc(bump: &'ir Bump, ir: Ir<'ir, Self>) -> Self {
        Self(bump.alloc(GhostCell::new(ir)))
    }

    pub fn borrow<'a>(&'a self, token: &'a GhostToken<'id>) -> &'a Ir<'ir, Self> {
        self.0.borrow(token)
    }

    /// Freeze a mutable IR reference into a read-only one, consuming the
    /// `GhostToken` to prevent any further mutation.
    ///
    /// # Safety
    /// The transmute is sound because:
    /// - `GhostCell<'id, T>` is `#[repr(transparent)]` over `T`
    /// - `IrRef<'id, 'ir>` is `#[repr(transparent)]` over
    ///   `&'ir GhostCell<'id, Ir<'ir, Self>>`
    /// - `RawIrRef<'ir>` is `#[repr(transparent)]` over `&'ir Ir<'ir, Self>`
    /// - `Ir<'ir, Ref>` is `#[repr(C)]` and both ref types are pointer-sized
    ///
    /// Consuming the `GhostToken` guarantees no `borrow_mut` calls can occur
    /// afterwards, so the shared `&Ir` references from `RawIrRef::Deref` can
    /// never alias with mutable references.
    pub fn freeze(self, _token: GhostToken<'id>) -> RawIrRef<'ir> {
        unsafe { std::mem::transmute(self) }
    }
}

#[repr(transparent)]
#[derive(Clone, Copy, Debug)]
pub struct RawIrRef<'ir>(pub &'ir Ir<'ir, Self>);

impl<'ir> Deref for RawIrRef<'ir> {
    type Target = Ir<'ir, RawIrRef<'ir>>;
    fn deref(&self) -> &Self::Target {
        self.0
    }
}

#[repr(C)]
#[derive(Debug)]
pub enum Ir<'ir, Ref> {
    Int(i64),
    Float(f64),
    Bool(bool),
    Null,
    Str(Box<'ir, String>),
    AttrSet {
        stcs: HashMap<'ir, StringId, (Ref, TextRange)>,
        dyns: Vec<'ir, (Ref, Ref, TextRange)>,
    },
    List {
        items: Vec<'ir, Ref>,
    },
    Path(Ref),
    ConcatStrings {
        parts: Vec<'ir, Ref>,
        force_string: bool,
    },

    // OPs
    UnOp {
        rhs: Ref,
        kind: UnOpKind,
    },
    BinOp {
        lhs: Ref,
        rhs: Ref,
        kind: BinOpKind,
    },
    HasAttr {
        lhs: Ref,
        rhs: Vec<'ir, Attr<Ref>>,
    },
    Select {
        expr: Ref,
        attrpath: Vec<'ir, Attr<Ref>>,
        default: Option<Ref>,
        span: TextRange,
    },

    // Conditionals
    If {
        cond: Ref,
        consq: Ref,
        alter: Ref,
    },
    Assert {
        assertion: Ref,
        expr: Ref,
        assertion_raw: String,
        span: TextRange,
    },

    With {
        namespace: Ref,
        body: Ref,
        thunks: Vec<'ir, (ThunkId, Ref)>,
    },
    WithLookup(StringId),

    // Function related
    Func {
        body: Ref,
        param: Option<Param<'ir>>,
        thunks: Vec<'ir, (ThunkId, Ref)>,
    },
    Arg {
        layer: usize,
    },
    Call {
        func: Ref,
        arg: Ref,
        span: TextRange,
    },

    // Builtins
    Builtins,
    Builtin(BuiltinId),
    BuiltinConst(StringId),

    // Misc
    TopLevel {
        body: Ref,
        thunks: Vec<'ir, (ThunkId, Ref)>,
    },
    Thunk(ThunkId),
    CurPos(TextRange),
    ReplBinding(StringId),
    ScopedImportBinding(StringId),
}

#[repr(transparent)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct ThunkId(pub usize);

#[repr(transparent)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct SpanId(pub u32);

/// Represents a key in an attribute path.
#[allow(unused)]
#[derive(Debug)]
pub enum Attr<Ref> {
    /// A dynamic attribute key, which is an expression that must evaluate to a string.
    /// Example: `attrs.${key}`
    Dynamic(Ref, TextRange),
    /// A static attribute key.
    /// Example: `attrs.key`
    Str(StringId, TextRange),
}

/// The kinds of binary operations supported in Nix.
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
pub enum BinOpKind {
    // Arithmetic
    Add,
    Sub,
    Div,
    Mul,

    // Comparison
    Eq,
    Neq,
    Lt,
    Gt,
    Leq,
    Geq,

    // Logical
    And,
    Or,
    Impl,

    // Set/String/Path operations
    Con, // List concatenation (`++`)
    Upd, // AttrSet update (`//`)

    // Not standard, but part of rnix AST
    PipeL,
    PipeR,
}

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

/// The kinds of unary operations.
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
pub enum UnOpKind {
    Neg, // Negation (`-`)
    Not, // Logical 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,
        }
    }
}

/// Describes the parameters of a function.
#[derive(Debug)]
pub struct Param<'ir> {
    pub required: Vec<'ir, (StringId, TextRange)>,
    pub optional: Vec<'ir, (StringId, TextRange)>,
    pub ellipsis: bool,
}

pub fn new_global_env(
    strings: &mut DefaultStringInterner,
) -> hashbrown::HashMap<StringId, Ir<'static, RawIrRef<'static>>> {
    let mut global_env = hashbrown::HashMap::new();
    let builtins_sym = StringId(strings.get_or_intern("builtins"));
    global_env.insert(builtins_sym, Ir::Builtins);

    for (idx, &(name, _)) in BUILTINS.iter().enumerate() {
        let id = BuiltinId::try_from_primitive(idx as u8).expect("infallible");
        let name = StringId(strings.get_or_intern(name));
        global_env.insert(name, Ir::Builtin(id));
    }

    let consts = [
        (
            "__currentSystem",
            Ir::BuiltinConst(StringId(strings.get_or_intern("currentSystem"))),
        ),
        ("__langVersion", Ir::Int(6)),
        (
            "__nixVersion",
            Ir::BuiltinConst(StringId(strings.get_or_intern("nixVersion"))),
        ),
        (
            "__storeDir",
            Ir::BuiltinConst(StringId(strings.get_or_intern("storeDir"))),
        ),
        (
            "__nixPath",
            Ir::BuiltinConst(StringId(strings.get_or_intern("nixPath"))),
        ),
        ("null", Ir::Null),
        ("true", Ir::Bool(true)),
        ("false", Ir::Bool(false)),
    ];

    for (name, ir) in consts {
        let name = StringId(strings.get_or_intern(name));
        global_env.insert(name, ir);
    }

    global_env
}