feat(tvix/glue/known_paths): add get_drv_by_output_path r/7539

This allows getting a Derivation struct producing the passed output
path.

Change-Id: I89858d91bffc2ef7f1d86314c16fa4f850f21426
Reviewed-on: https://cl.tvl.fyi/c/depot/+/10791
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
Reviewed-by: Peter Kolloch <info@eigenvalue.net>

1 files changed, 27 insertions, 0 deletions

diff --git a/tvix/glue/src/known_paths.rs b/tvix/glue/src/known_paths.rs
index d3059da4ad7d..fa2cf55fda56 100644
--- a/tvix/glue/src/known_paths.rs
+++ b/tvix/glue/src/known_paths.rs
@@ -25,6 +25,11 @@ pub struct KnownPaths {
     /// 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 {
@@ -42,6 +47,13 @@ impl KnownPaths {
             .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.
@@ -69,6 +81,21 @@ impl KnownPaths {
                 .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