1 files changed, 102 insertions, 9 deletions
diff --git a/tvix/glue/src/known_paths.rs b/tvix/glue/src/known_paths.rs
index 9cd9470fa949..873efdcf8a06 100644
--- a/tvix/glue/src/known_paths.rs
+++ b/tvix/glue/src/known_paths.rs
@@ -8,9 +8,14 @@
 //! This data is required to find the derivation needed to actually trigger the
 //! build, if necessary.
 
-use nix_compat::{derivation::Derivation, store_path::StorePath};
+use nix_compat::{
+    derivation::Derivation,
+    store_path::{BuildStorePathError, StorePath, StorePathRef},
+};
 use std::collections::HashMap;
 
+use crate::fetchers::Fetch;
+
 /// Struct keeping track of all known Derivations in the current evaluation.
 /// This keeps both the Derivation struct, as well as the "Hash derivation
 /// modulo".
@@ -26,6 +31,9 @@ pub struct KnownPaths {
     /// 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>,
+
+    /// A map from output path to fetches (and their names).
+    outputs_to_fetches: HashMap<StorePath, (String, Fetch)>,
 }
 
 impl KnownPaths {
@@ -50,12 +58,12 @@ impl KnownPaths {
         self.outputs_to_drvpath.get(output_path)
     }
 
-    /// Insert a new Derivation into this struct.
+    /// 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) {
+    pub fn add_derivation(&mut self, drv_path: StorePath, drv: Derivation) {
         // check input derivations to have been inserted.
         #[cfg(debug_assertions)]
         {
@@ -95,11 +103,39 @@ impl KnownPaths {
             }
         }
     }
+
+    /// Insert a new [Fetch] into this struct, which *must* have an expected
+    /// hash (otherwise we wouldn't be able to calculate the store path).
+    /// Fetches without a known hash need to be fetched inside builtins.
+    pub fn add_fetch<'a>(
+        &mut self,
+        fetch: Fetch,
+        name: &'a str,
+    ) -> Result<StorePathRef<'a>, BuildStorePathError> {
+        let store_path = fetch
+            .store_path(name)?
+            .expect("Tvix bug: fetch must have an expected hash");
+        // insert the fetch.
+        self.outputs_to_fetches
+            .insert(store_path.to_owned(), (name.to_owned(), fetch));
+
+        Ok(store_path)
+    }
+
+    /// Return the name and fetch producing the passed output path.
+    /// Note there can also be (multiple) Derivations producing the same output path.
+    pub fn get_fetch_for_output_path(&self, output_path: &StorePath) -> Option<(String, Fetch)> {
+        self.outputs_to_fetches
+            .get(output_path)
+            .map(|(name, fetch)| (name.to_owned(), fetch.to_owned()))
+    }
 }
 
 #[cfg(test)]
 mod tests {
-    use nix_compat::{derivation::Derivation, store_path::StorePath};
+    use nix_compat::{derivation::Derivation, nixbase32, nixhash::NixHash, store_path::StorePath};
+
+    use crate::fetchers::Fetch;
 
     use super::KnownPaths;
     use hex_literal::hex;
@@ -122,21 +158,33 @@ mod tests {
             StorePath::from_bytes(b"mp57d33657rf34lzvlbpfa1gjfv5gmpg-bar").expect("must parse");
         static ref FOO_OUT_PATH: StorePath =
             StorePath::from_bytes(b"fhaj6gmwns62s6ypkcldbaj2ybvkhx3p-foo").expect("must parse");
+
+        static ref FETCH_URL : Fetch = Fetch::URL(
+            "https://raw.githubusercontent.com/aaptel/notmuch-extract-patch/f732a53e12a7c91a06755ebfab2007adc9b3063b/notmuch-extract-patch".into(),
+            Some(NixHash::Sha256(nixbase32::decode_fixed("0nawkl04sj7psw6ikzay7kydj3dhd0fkwghcsf5rzaw4bmp4kbax").unwrap()))
+        );
+        static ref FETCH_URL_OUT_PATH: StorePath = StorePath::from_bytes(b"06qi00hylriyfm0nl827crgjvbax84mz-notmuch-extract-patch").unwrap();
+
+        static ref FETCH_TARBALL : Fetch = Fetch::Tarball(
+            "https://github.com/NixOS/nixpkgs/archive/91050ea1e57e50388fa87a3302ba12d188ef723a.tar.gz".into(),
+            Some(nixbase32::decode_fixed("1hf6cgaci1n186kkkjq106ryf8mmlq9vnwgfwh625wa8hfgdn4dm").unwrap())
+        );
+        static ref FETCH_TARBALL_OUT_PATH: StorePath = StorePath::from_bytes(b"7adgvk5zdfq4pwrhsm3n9lzypb12gw0g-source").unwrap();
     }
 
     /// ensure we don't allow acdding a Derivation that depends on another,
     /// not-yet-added Derivation.
     #[test]
     #[should_panic]
-    fn reject_if_missing_input_drv() {
+    fn drv_reject_if_missing_input_drv() {
         let mut known_paths = KnownPaths::default();
 
         // FOO_DRV depends on BAR_DRV, which wasn't added.
-        known_paths.add(FOO_DRV_PATH.clone(), FOO_DRV.clone());
+        known_paths.add_derivation(FOO_DRV_PATH.clone(), FOO_DRV.clone());
     }
 
     #[test]
-    fn happy_path() {
+    fn drv_happy_path() {
         let mut known_paths = KnownPaths::default();
 
         // get_drv_by_drvpath should return None for non-existing Derivations,
@@ -149,7 +197,7 @@ mod tests {
         );
 
         // Add BAR_DRV
-        known_paths.add(BAR_DRV_PATH.clone(), BAR_DRV.clone());
+        known_paths.add_derivation(BAR_DRV_PATH.clone(), BAR_DRV.clone());
 
         // We should get it back
         assert_eq!(
@@ -173,7 +221,7 @@ mod tests {
 
         // Now insert FOO_DRV too. It shouldn't panic, as BAR_DRV is already
         // added.
-        known_paths.add(FOO_DRV_PATH.clone(), FOO_DRV.clone());
+        known_paths.add_derivation(FOO_DRV_PATH.clone(), FOO_DRV.clone());
 
         assert_eq!(
             Some(&FOO_DRV.clone()),
@@ -192,4 +240,49 @@ mod tests {
             known_paths.get_drv_path_for_output_path(&FOO_OUT_PATH)
         );
     }
+
+    #[test]
+    fn fetch_happy_path() {
+        let mut known_paths = KnownPaths::default();
+
+        // get_fetch_for_output_path should return None for new fetches.
+        assert!(known_paths
+            .get_fetch_for_output_path(&FETCH_TARBALL_OUT_PATH)
+            .is_none());
+
+        // add_fetch should return the properly calculated store paths.
+        assert_eq!(
+            *FETCH_TARBALL_OUT_PATH,
+            known_paths
+                .add_fetch(FETCH_TARBALL.clone(), "source")
+                .unwrap()
+                .to_owned()
+        );
+
+        assert_eq!(
+            *FETCH_URL_OUT_PATH,
+            known_paths
+                .add_fetch(FETCH_URL.clone(), "notmuch-extract-patch")
+                .unwrap()
+                .to_owned()
+        );
+
+        // We should be able to get these fetches out, when asking for their out path.
+        let (got_name, got_fetch) = known_paths
+            .get_fetch_for_output_path(&FETCH_URL_OUT_PATH)
+            .expect("must be some");
+
+        assert_eq!("notmuch-extract-patch", got_name);
+        assert_eq!(FETCH_URL.clone(), got_fetch);
+
+        // … multiple times.
+        let (got_name, got_fetch) = known_paths
+            .get_fetch_for_output_path(&FETCH_URL_OUT_PATH)
+            .expect("must be some");
+
+        assert_eq!("notmuch-extract-patch", got_name);
+        assert_eq!(FETCH_URL.clone(), got_fetch);
+    }
+
+    // TODO: add test panicking about missing digest
 }