tvl-depot/tvix/glue/src/known_paths.rs

//! This module implements logic required for persisting known paths
//! during an evaluation.
//!
//! Tvix needs to be able to keep track of each Nix store path that it
//! knows about during the scope of a single evaluation and its
//! related builds.
//!
//! This data is required to find the derivation needed to actually trigger the
//! build, if necessary.

use nix_compat::{
    derivation::Derivation,
    nixhash::NixHash,
    store_path::{StorePath, StorePathRef},
};
use std::collections::HashMap;

/// Struct keeping track of all known Derivations in the current evaluation.
/// This keeps both the Derivation struct, as well as the "Hash derivation
/// modulo".
#[derive(Debug, Default)]
pub struct KnownPaths {
    /// All known derivation or FOD hashes.
    ///
    /// Keys are derivation paths, values are a tuple of the "hash derivation
    /// modulo" and the Derivation struct itself.
    derivations: HashMap<StorePath, (NixHash, Derivation)>,

    /// A map from output path to (one) drv path.
    /// Note that in the case of FODs, multiple drvs can produce the same output
    /// path. We use one of them.
    outputs_to_drvpath: HashMap<StorePath, StorePath>,
}

impl KnownPaths {
    /// Fetch the opaque "hash derivation modulo" for a given derivation path.
    pub fn get_hash_derivation_modulo(&self, drv_path: &StorePathRef) -> Option<&NixHash> {
        self.derivations
            .get(&drv_path.to_owned())
            .map(|(hash_derivation_modulo, _derivation)| hash_derivation_modulo)
    }

    /// Return a reference to the Derivation for a given drv path.
    pub fn get_drv_by_drvpath(&self, drv_path: &StorePath) -> Option<&Derivation> {
        self.derivations
            .get(drv_path)
            .map(|(_hash_derivation_modulo, derivation)| derivation)
    }

    /// Return the drv path of the derivation producing the passed output path.
    /// Note there can be multiple Derivations producing the same output path in
    /// flight; this function will only return one of them.
    pub fn get_drv_path_for_output_path(&self, output_path: &StorePath) -> Option<&StorePath> {
        self.outputs_to_drvpath.get(output_path)
    }

    /// Insert a new Derivation into this struct.
    /// The Derivation struct must pass validation, and its output paths need to
    /// be fully calculated.
    /// All input derivations this refers to must also be inserted to this
    /// struct.
    pub fn add(&mut self, drv_path: StorePath, drv: Derivation) {
        // check input derivations to have been inserted.
        #[cfg(debug_assertions)]
        {
            // TODO: b/264
            // We assume derivations to be passed validated, so ignoring rest
            // and expecting parsing is ok.
            for input_drv_path_str in drv.input_derivations.keys() {
                let (input_drv_path, _rest) =
                    StorePath::from_absolute_path_full(input_drv_path_str)
                        .expect("parse input drv path");
                debug_assert!(self.derivations.contains_key(&input_drv_path));
            }
        }

        // compute the hash derivation modulo
        let hash_derivation_modulo = drv.derivation_or_fod_hash(|drv_path| {
            self.get_hash_derivation_modulo(drv_path)
                .unwrap_or_else(|| panic!("{} not found", drv_path))
                .to_owned()
        });

        // For all output paths, update our lookup table.
        // We only write into the lookup table once.
        for output in drv.outputs.values() {
            // We assume derivations to be passed validated, so ignoring rest
            // and expecting parsing is ok.
            // TODO: b/264
            let (output_path, _rest) =
                StorePath::from_absolute_path_full(&output.path).expect("parse output path");

            self.outputs_to_drvpath
                .entry(output_path)
                .or_insert(drv_path.to_owned());
        }

        // insert the derivation itself
        #[allow(unused_variables)] // assertions on this only compiled in debug builds
        let old = self
            .derivations
            .insert(drv_path.to_owned(), (hash_derivation_modulo.clone(), drv));

        #[cfg(debug_assertions)]
        {
            if let Some(old) = old {
                debug_assert!(
                    old.0 == hash_derivation_modulo,
                    "hash derivation modulo for a given derivation should always be calculated the same"
                );
            }
        }
    }
}