feat(tvix/nix-compat): input_derivations with StorePaths r/7583

...in `Derivation`.

This is more type-safe and should consume less memory.

This also removes some allocations in the potentially hot path of output hash calculation.

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

5 files changed, 48 insertions, 47 deletions

diff --git a/tvix/nix-compat/src/derivation/mod.rs b/tvix/nix-compat/src/derivation/mod.rs
index 3a09c9d56632..3cf9a5ba7f22 100644
--- a/tvix/nix-compat/src/derivation/mod.rs
+++ b/tvix/nix-compat/src/derivation/mod.rs
@@ -36,7 +36,7 @@ pub struct Derivation {
 
     /// Map from drv path to output names used from this derivation.
     #[serde(rename = "inputDrvs")]
-    pub input_derivations: BTreeMap<String, BTreeSet<String>>,
+    pub input_derivations: BTreeMap<StorePath, BTreeSet<String>>,
 
     /// Plain store paths of additional inputs.
     #[serde(rename = "inputSrcs")]
@@ -70,19 +70,7 @@ impl Derivation {
         write_outputs(writer, &self.outputs)?;
         write_char(writer, COMMA)?;
 
-        // To reproduce the exact hashes of original nix, we need to sort
-        // these by their aterm representation.
-        // StorePath has a different sort order, so we convert it here.
-        let input_derivations: BTreeMap<BString, &BTreeSet<String>> = input_derivations
-            .iter()
-            .map(|(k, v)| {
-                let mut aterm_k = Vec::new();
-                k.aterm_write(&mut aterm_k).expect("no write error to Vec");
-                (BString::new(aterm_k), v)
-            })
-            .collect();
-
-        write_input_derivations(writer, &input_derivations)?;
+        write_input_derivations(writer, input_derivations)?;
         write_char(writer, COMMA)?;
 
         write_input_sources(writer, &self.input_sources)?;
@@ -145,8 +133,11 @@ impl Derivation {
         // into a (sorted, guaranteed by BTreeSet) list of references
         let references: BTreeSet<String> = {
             let mut inputs = self.input_sources.clone();
-            let input_derivation_keys: Vec<String> =
-                self.input_derivations.keys().cloned().collect();
+            let input_derivation_keys: Vec<String> = self
+                .input_derivations
+                .keys()
+                .map(|k| k.to_absolute_path())
+                .collect();
             inputs.extend(input_derivation_keys);
             inputs
         };
@@ -211,12 +202,8 @@ impl Derivation {
             // derivation_or_fod_hash, and replace the derivation path with
             // it's HEXLOWER digest.
             let input_derivations = BTreeMap::from_iter(self.input_derivations.iter().map(
-                |(drv_path_str, output_names)| {
-                    // parse drv_path to StorePathRef
-                    let drv_path = StorePathRef::from_absolute_path(drv_path_str.as_bytes())
-                        .expect("invalid input derivation path");
-
-                    let hash = fn_get_derivation_or_fod_hash(&drv_path);
+                |(drv_path, output_names)| {
+                    let hash = fn_get_derivation_or_fod_hash(&drv_path.into());
 
                     (hash, output_names.to_owned())
                 },
diff --git a/tvix/nix-compat/src/derivation/parser.rs b/tvix/nix-compat/src/derivation/parser.rs
index 7ffa6fd46eb6..2cfcf2eae377 100644
--- a/tvix/nix-compat/src/derivation/parser.rs
+++ b/tvix/nix-compat/src/derivation/parser.rs
@@ -14,6 +14,7 @@ use thiserror;
 
 use crate::derivation::parse_error::{into_nomerror, ErrorKind, NomError, NomResult};
 use crate::derivation::{write, CAHash, Derivation, Output};
+use crate::store_path::{self, StorePath, StorePathRef};
 use crate::{aterm, nixhash};
 
 #[derive(Debug, thiserror::Error)]
@@ -142,11 +143,11 @@ fn parse_outputs(i: &[u8]) -> NomResult<&[u8], BTreeMap<String, Output>> {
     }
 }
 
-fn parse_input_derivations(i: &[u8]) -> NomResult<&[u8], BTreeMap<String, BTreeSet<String>>> {
+fn parse_input_derivations(i: &[u8]) -> NomResult<&[u8], BTreeMap<StorePath, BTreeSet<String>>> {
     let (i, input_derivations_list) = parse_kv::<Vec<String>, _>(aterm::parse_str_list)(i)?;
 
     // This is a HashMap of drv paths to a list of output names.
-    let mut input_derivations: BTreeMap<String, BTreeSet<String>> = BTreeMap::new();
+    let mut input_derivations: BTreeMap<StorePath, BTreeSet<String>> = BTreeMap::new();
 
     for (input_derivation, output_names) in input_derivations_list {
         let mut new_output_names = BTreeSet::new();
@@ -159,10 +160,26 @@ fn parse_input_derivations(i: &[u8]) -> NomResult<&[u8], BTreeMap<String, BTreeS
                         output_name.to_string(),
                     ),
                 }));
-            } else {
-                new_output_names.insert(output_name);
             }
+            new_output_names.insert(output_name);
         }
+
+        #[cfg(debug_assertions)]
+        let input_derivation_str = input_derivation.clone();
+
+        let input_derivation: StorePath =
+            StorePathRef::from_absolute_path(input_derivation.as_bytes())
+                .map_err(|e: store_path::Error| {
+                    nom::Err::Failure(NomError {
+                        input: i,
+                        code: e.into(),
+                    })
+                })?
+                .to_owned();
+
+        #[cfg(debug_assertions)]
+        assert_eq!(input_derivation_str, input_derivation.to_absolute_path());
+
         input_derivations.insert(input_derivation, new_output_names);
     }
 
@@ -297,8 +314,11 @@ where
 mod tests {
     use std::collections::{BTreeMap, BTreeSet};
 
-    use crate::derivation::{
-        parse_error::ErrorKind, parser::from_algo_and_mode_and_digest, CAHash, NixHash, Output,
+    use crate::{
+        derivation::{
+            parse_error::ErrorKind, parser::from_algo_and_mode_and_digest, CAHash, NixHash, Output,
+        },
+        store_path::StorePath,
     };
     use bstr::{BString, ByteSlice};
     use hex_literal::hex;
@@ -334,10 +354,11 @@ mod tests {
             b.insert("b".to_string(), b"2".as_bstr().to_owned());
             b
         };
-        static ref EXP_INPUT_DERIVATIONS_SIMPLE: BTreeMap<String, BTreeSet<String>> = {
+        static ref EXP_INPUT_DERIVATIONS_SIMPLE: BTreeMap<StorePath, BTreeSet<String>> = {
             let mut b = BTreeMap::new();
             b.insert(
-                "/nix/store/8bjm87p310sb7r2r0sg4xrynlvg86j8k-hello-2.12.1.tar.gz.drv".to_string(),
+                StorePath::from_bytes(b"8bjm87p310sb7r2r0sg4xrynlvg86j8k-hello-2.12.1.tar.gz.drv")
+                    .unwrap(),
                 {
                     let mut output_names = BTreeSet::new();
                     output_names.insert("out".to_string());
@@ -345,7 +366,8 @@ mod tests {
                 },
             );
             b.insert(
-                "/nix/store/p3jc8aw45dza6h52v81j7lk69khckmcj-bash-5.2-p15.drv".to_string(),
+                StorePath::from_bytes(b"p3jc8aw45dza6h52v81j7lk69khckmcj-bash-5.2-p15.drv")
+                    .unwrap(),
                 {
                     let mut output_names = BTreeSet::new();
                     output_names.insert("out".to_string());
@@ -400,7 +422,7 @@ mod tests {
     #[test_case(EXP_INPUT_DERIVATIONS_SIMPLE_ATERM.as_bytes(), &EXP_INPUT_DERIVATIONS_SIMPLE; "simple")]
     fn parse_input_derivations(
         input: &'static [u8],
-        expected: &BTreeMap<String, BTreeSet<String>>,
+        expected: &BTreeMap<StorePath, BTreeSet<String>>,
     ) {
         let (rest, parsed) = super::parse_input_derivations(input).expect("must parse");
 
diff --git a/tvix/nix-compat/src/derivation/tests/mod.rs b/tvix/nix-compat/src/derivation/tests/mod.rs
index 168e11d46f1e..fcfea2047e96 100644
--- a/tvix/nix-compat/src/derivation/tests/mod.rs
+++ b/tvix/nix-compat/src/derivation/tests/mod.rs
@@ -104,7 +104,7 @@ fn from_aterm_bytes(path_to_drv_file: &str) {
 
     assert_eq!(
         &aterm_bytes,
-        &parsed_drv.to_aterm_bytes(),
+        &BString::new(parsed_drv.to_aterm_bytes()),
         "expected serialized ATerm to match initial input"
     );
 }
@@ -381,10 +381,9 @@ fn output_path_construction() {
     );
 
     // assemble foo input_derivations
-    foo_drv.input_derivations.insert(
-        bar_drv_path.to_absolute_path(),
-        BTreeSet::from(["out".to_string()]),
-    );
+    foo_drv
+        .input_derivations
+        .insert(bar_drv_path, BTreeSet::from(["out".to_string()]));
 
     // calculate foo output paths
     let foo_calc_result = foo_drv.calculate_output_paths(
diff --git a/tvix/nix-compat/src/derivation/validate.rs b/tvix/nix-compat/src/derivation/validate.rs
index d711f5ce1de2..757326816c63 100644
--- a/tvix/nix-compat/src/derivation/validate.rs
+++ b/tvix/nix-compat/src/derivation/validate.rs
@@ -53,14 +53,7 @@ impl Derivation {
         // Validate all input_derivations
         for (input_derivation_path, output_names) in &self.input_derivations {
             // Validate input_derivation_path
-            if let Err(e) = StorePathRef::from_absolute_path(input_derivation_path.as_bytes()) {
-                return Err(DerivationError::InvalidInputDerivationPath(
-                    input_derivation_path.to_string(),
-                    e,
-                ));
-            }
-
-            if !input_derivation_path.ends_with(".drv") {
+            if !input_derivation_path.name().ends_with(".drv") {
                 return Err(DerivationError::InvalidInputDerivationPrefix(
                     input_derivation_path.to_string(),
                 ));
diff --git a/tvix/nix-compat/src/derivation/write.rs b/tvix/nix-compat/src/derivation/write.rs
index f3b16d9cf9a4..83106cd9e60b 100644
--- a/tvix/nix-compat/src/derivation/write.rs
+++ b/tvix/nix-compat/src/derivation/write.rs
@@ -149,7 +149,7 @@ pub(crate) fn write_outputs(
 
 pub(crate) fn write_input_derivations(
     writer: &mut impl Write,
-    input_derivations: &BTreeMap<BString, &BTreeSet<String>>,
+    input_derivations: &BTreeMap<impl AtermWriteable, BTreeSet<String>>,
 ) -> Result<(), io::Error> {
     write_char(writer, BRACKET_OPEN)?;
 
@@ -159,7 +159,7 @@ pub(crate) fn write_input_derivations(
         }
 
         write_char(writer, PAREN_OPEN)?;
-        writer.write_all(input_derivation_aterm)?;
+        input_derivation_aterm.aterm_write(writer)?;
         write_char(writer, COMMA)?;
 
         write_char(writer, BRACKET_OPEN)?;