5 files changed, 104 insertions, 0 deletions
diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock
index 01b92baccf65..430d369a1bee 100644
--- a/tvix/Cargo.lock
+++ b/tvix/Cargo.lock
@@ -3377,6 +3377,7 @@ dependencies = [
  "bytes",
  "criterion",
  "data-encoding",
+ "hex-literal",
  "lazy_static",
  "nix-compat",
  "pretty_assertions",
diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index b53ae3477615..4e6c037d0f52 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -10682,6 +10682,10 @@ rec {
             features = [ "html_reports" ];
           }
           {
+            name = "hex-literal";
+            packageId = "hex-literal";
+          }
+          {
             name = "lazy_static";
             packageId = "lazy_static";
           }
diff --git a/tvix/glue/Cargo.toml b/tvix/glue/Cargo.toml
index 451b28f1f0a0..748635f3940d 100644
--- a/tvix/glue/Cargo.toml
+++ b/tvix/glue/Cargo.toml
@@ -24,6 +24,7 @@ git = "https://github.com/tvlfyi/wu-manber.git"
 
 [dev-dependencies]
 criterion = { version = "0.5", features = ["html_reports"] }
+hex-literal = "0.4.1"
 lazy_static = "1.4.0"
 pretty_assertions = "1.4.0"
 rstest = "0.18.2"
diff --git a/tvix/glue/src/known_paths.rs b/tvix/glue/src/known_paths.rs
index 1e3fe85a0f7c..03b04d4fc85f 100644
--- a/tvix/glue/src/known_paths.rs
+++ b/tvix/glue/src/known_paths.rs
@@ -108,3 +108,100 @@ impl KnownPaths {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use nix_compat::{derivation::Derivation, nixhash::NixHash, store_path::StorePath};
+
+    use super::KnownPaths;
+    use hex_literal::hex;
+    use lazy_static::lazy_static;
+
+    lazy_static! {
+        static ref BAR_DRV: Derivation = Derivation::from_aterm_bytes(include_bytes!(
+            "tests/ss2p4wmxijn652haqyd7dckxwl4c7hxx-bar.drv"
+        ))
+        .expect("must parse");
+        static ref FOO_DRV: Derivation = Derivation::from_aterm_bytes(include_bytes!(
+            "tests/ch49594n9avinrf8ip0aslidkc4lxkqv-foo.drv"
+        ))
+        .expect("must parse");
+        static ref BAR_DRV_PATH: StorePath =
+            StorePath::from_bytes(b"ss2p4wmxijn652haqyd7dckxwl4c7hxx-bar.drv").expect("must parse");
+        static ref FOO_DRV_PATH: StorePath =
+            StorePath::from_bytes(b"ch49594n9avinrf8ip0aslidkc4lxkqv-foo.drv").expect("must parse");
+        static ref BAR_OUT_PATH: StorePath =
+            StorePath::from_bytes(b"mp57d33657rf34lzvlbpfa1gjfv5gmpg-bar").expect("must parse");
+        static ref FOO_OUT_PATH: StorePath =
+            StorePath::from_bytes(b"fhaj6gmwns62s6ypkcldbaj2ybvkhx3p-foo").expect("must parse");
+    }
+
+    /// 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() {
+        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());
+    }
+
+    #[test]
+    fn happy_path() {
+        let mut known_paths = KnownPaths::default();
+
+        // get_drv_by_drvpath should return None for non-existing Derivations,
+        // same as get_hash_derivation_modulo and get_drv_path_for_output_path
+        assert_eq!(None, known_paths.get_drv_by_drvpath(&BAR_DRV_PATH));
+        assert_eq!(None, known_paths.get_hash_derivation_modulo(&BAR_DRV_PATH));
+        assert_eq!(
+            None,
+            known_paths.get_drv_path_for_output_path(&BAR_OUT_PATH)
+        );
+
+        // Add BAR_DRV
+        known_paths.add(BAR_DRV_PATH.clone(), BAR_DRV.clone());
+
+        // We should get it back
+        assert_eq!(
+            Some(&BAR_DRV.clone()),
+            known_paths.get_drv_by_drvpath(&BAR_DRV_PATH)
+        );
+
+        // Test get_drv_path_for_output_path
+        assert_eq!(
+            Some(&BAR_DRV_PATH.clone()),
+            known_paths.get_drv_path_for_output_path(&BAR_OUT_PATH)
+        );
+
+        // It should be possible to get the hash derivation modulo.
+        assert_eq!(
+            Some(&NixHash::Sha256(hex!(
+                "c79aebd0ce3269393d4a1fde2cbd1d975d879b40f0bf40a48f550edc107fd5df"
+            ))),
+            known_paths.get_hash_derivation_modulo(&BAR_DRV_PATH.clone())
+        );
+
+        // 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());
+
+        assert_eq!(
+            Some(&FOO_DRV.clone()),
+            known_paths.get_drv_by_drvpath(&FOO_DRV_PATH)
+        );
+        assert_eq!(
+            Some(&NixHash::Sha256(hex!(
+                "af030d36d63d3d7f56a71adaba26b36f5fa1f9847da5eed953ed62e18192762f"
+            ))),
+            known_paths.get_hash_derivation_modulo(&FOO_DRV_PATH.clone())
+        );
+
+        // Test get_drv_path_for_output_path
+        assert_eq!(
+            Some(&FOO_DRV_PATH.clone()),
+            known_paths.get_drv_path_for_output_path(&FOO_OUT_PATH)
+        );
+    }
+}
diff --git a/tvix/glue/src/tests/ss2p4wmxijn652haqyd7dckxwl4c7hxx-bar.drv b/tvix/glue/src/tests/ss2p4wmxijn652haqyd7dckxwl4c7hxx-bar.drv
new file mode 100644
index 000000000000..559e93ed0ed6
--- /dev/null
+++ b/tvix/glue/src/tests/ss2p4wmxijn652haqyd7dckxwl4c7hxx-bar.drv
@@ -0,0 +1 @@
+Derive([("out","/nix/store/mp57d33657rf34lzvlbpfa1gjfv5gmpg-bar","r:sha1","0beec7b5ea3f0fdbc95d0dd47f3c5bc275da8a33")],[],[],":",":",[],[("builder",":"),("name","bar"),("out","/nix/store/mp57d33657rf34lzvlbpfa1gjfv5gmpg-bar"),("outputHash","0beec7b5ea3f0fdbc95d0dd47f3c5bc275da8a33"),("outputHashAlgo","sha1"),("outputHashMode","recursive"),("system",":")])
\ No newline at end of file