evix 0.3.0

use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::atomic::AtomicBool;
use std::sync::{Arc, Mutex};

use anyhow::Context as _;
use serde::{Deserialize, Serialize};
use tracing::debug;

mod eval;
mod master;
mod worker;

/// Environment variable used to distinguish worker subprocesses spawned by
/// [`evaluate`]. A binary that re-executes itself to host workers should check
/// this variable and call [`run_worker`] when it is set.
pub const WORKER_ENV: &str = "EVIX_WORKER";

/// Input source for a Nix evaluation.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Input {
    Flake(String),
    Expr(String),
    File(PathBuf),
}

/// Argument passed to a Nix function parameter.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum AutoArg {
    Expr(String),
    Str(String),
}

/// Configuration for an evaluation run.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Config {
    pub input: Input,
    pub auto_args: Vec<(String, AutoArg)>,
    pub force_recurse: bool,
    pub gc_roots_dir: Option<PathBuf>,
    pub workers: usize,
    pub max_memory_size: usize,
    /// Attach each derivation's `meta` attribute (description, license,
    /// homepage, maintainers, ...) to the emitted [`Derivation`]. Off by
    /// default because forcing `meta` deeply costs extra evaluation.
    #[serde(default)]
    pub meta: bool,
    /// Read each derivation's input derivations from the store and attach them
    /// as [`Derivation::input_drvs`]. Off by default because it reads the
    /// `.drv` file for every job.
    #[serde(default)]
    pub show_input_drvs: bool,
    /// Flake input overrides applied while locking, as `(input_path, ref)`
    /// pairs (e.g., `("nixpkgs", "github:NixOS/nixpkgs/nixos-unstable")`). Only
    /// meaningful for [`Input::Flake`].
    #[serde(default)]
    pub override_inputs: Vec<(String, String)>,
    /// Nix settings applied to the evaluation context before the eval state is
    /// built, as `(key, value)` pairs (e.g.,
    /// `("allow-import-from-derivation", "false")`). Equivalent to `nix`'s
    /// `--option KEY VALUE`.
    #[serde(default)]
    pub nix_options: Vec<(String, String)>,
}

/// A derivation emitted by evaluation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Derivation {
    pub attr: String,
    pub attr_path: Vec<String>,
    pub name: String,
    pub system: String,
    pub drv_path: String,
    pub outputs: BTreeMap<String, Option<String>>,
    /// The derivation's `meta` attribute as freeform JSON, present only when
    /// [`Config::meta`] is set and the attribute exists.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub meta: Option<serde_json::Value>,
    /// Input derivations keyed by `.drv` store path, present only when
    /// [`Config::show_input_drvs`] is set. The value is the output-name list for
    /// that derivation input (e.g., `["out"]`).
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub input_drvs: BTreeMap<String, serde_json::Value>,
    /// Constituent attribute names for an aggregate job (Hydra
    /// `constituents`), present only when the derivation declares them.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub constituents: Option<Vec<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub gc_root_error: Option<String>,
}

/// An evaluation error associated with a specific attribute path.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EvalError {
    pub attr: String,
    pub attr_path: Vec<String>,
    pub error: String,
    pub fatal: bool,
}

/// Event produced while traversing a Nix expression.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Event {
    Derivation(Derivation),
    AttrSet {
        attr: String,
        attr_path: Vec<String>,
        attrs: Vec<String>,
    },
    Error(EvalError),
}

impl Event {
    /// Attribute path rendered with dots.
    pub fn attr(&self) -> &str {
        match self {
            Event::Derivation(d) => &d.attr,
            Event::AttrSet { attr, .. } => attr,
            Event::Error(e) => &e.attr,
        }
    }

    /// Attribute path as a list of names.
    pub fn attr_path(&self) -> &[String] {
        match self {
            Event::Derivation(d) => &d.attr_path,
            Event::AttrSet { attr_path, .. } => attr_path,
            Event::Error(e) => &e.attr_path,
        }
    }
}

/// Run an evaluation and deliver each event to `sink`.
///
/// The implementation uses worker subprocesses to isolate evaluation memory.
/// Each worker re-executes the current binary; the binary must call
/// [`run_worker`] when the [`WORKER_ENV`] environment variable is set.
///
/// ```no_run
/// use evix::{Config, Event, Input};
///
/// let config = Config {
///     input: Input::Expr("import <nixpkgs> {}".into()),
///     auto_args: vec![],
///     force_recurse: false,
///     gc_roots_dir: None,
///     workers: 4,
///     max_memory_size: 4096,
///     meta: false,
///     show_input_drvs: false,
///     override_inputs: vec![],
///     nix_options: vec![],
/// };
///
/// evix::evaluate(&config, |event| {
///     println!("{:?}", event);
///     Ok(())
/// }).unwrap();
/// ```
pub fn evaluate<F>(config: &Config, sink: F) -> anyhow::Result<()>
where
    F: FnMut(&Event) -> anyhow::Result<()> + Send + 'static,
{
    evaluate_cancellable(config, &Arc::new(AtomicBool::new(false)), sink)
}

/// Like [`evaluate`], but observes a cancellation flag.
///
/// Setting `cancel` makes the master stop dispatching work and tell its workers
/// to exit. Cancellation is cooperative: a worker already evaluating an
/// attribute finishes it before observing the request, so a caller can enforce a
/// wall-clock timeout without leaking worker processes.
///
/// # Errors
///
/// Returns an error if a worker reports a fatal evaluation error, if a worker
/// process fails unexpectedly, or if `sink` returns an error.
pub fn evaluate_cancellable<F>(
    config: &Config,
    cancel: &Arc<AtomicBool>,
    sink: F,
) -> anyhow::Result<()>
where
    F: FnMut(&Event) -> anyhow::Result<()> + Send + 'static,
{
    debug!("evaluating input, {} workers", config.workers);

    if let Some(dir) = &config.gc_roots_dir {
        std::fs::create_dir_all(dir).with_context(|| format!("creating gc-roots dir {dir:?}"))?;
        debug!("ensured gc-roots directory exists");
    }

    let sink = Arc::new(Mutex::new(sink));
    master::run(config, cancel, sink)
}

/// Worker entrypoint.
///
/// Reads the [`Config`] as a JSON line from stdin, then processes attribute
/// paths requested by the master process.
pub fn run_worker() -> anyhow::Result<()> {
    worker::run()
}