refactor(tvix/build): use stricter BuildRequest type r/8855

Change-Id: Ifadd190e10ec22570ab3ccb4df54f64ae5ef0a44
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12674
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

2 files changed, 80 insertions, 85 deletions

diff --git a/tvix/glue/src/tvix_build.rs b/tvix/glue/src/tvix_build.rs
index 6a3c6ab40ccc..fa73224992e6 100644
--- a/tvix/glue/src/tvix_build.rs
+++ b/tvix/glue/src/tvix_build.rs
@@ -1,16 +1,13 @@
 //! This module contains glue code translating from
-//! [nix_compat::derivation::Derivation] to [tvix_build::proto::BuildRequest].
+//! [nix_compat::derivation::Derivation] to [tvix_build::buildservice::BuildRequest].
 
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashSet};
+use std::path::PathBuf;
 
 use bytes::Bytes;
 use nix_compat::{derivation::Derivation, nixbase32, store_path::StorePath};
 use sha2::{Digest, Sha256};
-use tvix_build::buildservice::BuildRequest;
-use tvix_build::proto::{
-    self,
-    build_request::{AdditionalFile, BuildConstraints, EnvVar},
-};
+use tvix_build::buildservice::{AdditionalFile, BuildConstraints, BuildRequest, EnvVar};
 use tvix_castore::Node;
 
 /// These are the environment variables that Nix sets in its sandbox for every
@@ -31,7 +28,7 @@ const NIX_ENVIRONMENT_VARS: [(&str, &str); 12] = [
 ];
 
 /// Get an iterator of store paths whose nixbase32 hashes will be the needles for refscanning
-/// Importantly, the returned order will match the one used by derivation_to_build_request
+/// Importantly, the returned order will match the one used by [derivation_to_build_request]
 /// so users may use this function to map back from the found needles to a store path
 pub(crate) fn get_refscan_needles(
     derivation: &Derivation,
@@ -44,7 +41,7 @@ pub(crate) fn get_refscan_needles(
         .chain(derivation.input_derivations.keys())
 }
 
-/// Takes a [Derivation] and turns it into a [proto::BuildRequest].
+/// Takes a [Derivation] and turns it into a [buildservice::BuildRequest].
 /// It assumes the Derivation has been validated.
 /// It needs two lookup functions:
 /// - one translating input sources to a castore node
@@ -53,8 +50,8 @@ pub(crate) fn get_refscan_needles(
 ///   castore nodes of the selected outpus (`fn_input_drvs_to_output_nodes`).
 pub(crate) fn derivation_to_build_request(
     derivation: &Derivation,
-    inputs: BTreeMap<bytes::Bytes, Node>,
-) -> std::io::Result<proto::BuildRequest> {
+    inputs: BTreeMap<StorePath<String>, Node>,
+) -> std::io::Result<BuildRequest> {
     debug_assert!(derivation.validate(true).is_ok(), "drv must validate");
 
     // produce command_args, which is builder and arguments in a Vec.
@@ -62,16 +59,6 @@ pub(crate) fn derivation_to_build_request(
     command_args.push(derivation.builder.clone());
     command_args.extend_from_slice(&derivation.arguments);
 
-    // produce output_paths, which is the absolute path of each output (sorted)
-    let mut output_paths: Vec<String> = derivation
-        .outputs
-        .values()
-        .map(|e| e.path_str()[1..].to_owned())
-        .collect();
-
-    // Sort the outputs. We can use sort_unstable, as these are unique strings.
-    output_paths.sort_unstable();
-
     // Produce environment_vars and additional files.
     // We use a BTreeMap while producing, and only realize the resulting Vec
     // while populating BuildRequest, so we don't need to worry about ordering.
@@ -100,28 +87,41 @@ pub(crate) fn derivation_to_build_request(
     // TODO: handle __json (structured attrs, provide JSON file and source-able bash script)
 
     // Produce constraints.
-    let constraints = Some(BuildConstraints {
-        system: derivation.system.clone(),
-        min_memory: 0,
-        available_ro_paths: vec![],
-        // in case this is a fixed-output derivation, allow network access.
-        network_access: derivation.outputs.len() == 1
-            && derivation
-                .outputs
-                .get("out")
-                .expect("invalid derivation")
-                .is_fixed(),
-        provide_bin_sh: true,
-    });
+    let mut constraints = HashSet::from([
+        BuildConstraints::System(derivation.system.clone()),
+        BuildConstraints::ProvideBinSh,
+    ]);
+
+    if derivation.outputs.len() == 1
+        && derivation
+            .outputs
+            .get("out")
+            .expect("Tvix bug: Derivation has no out output")
+            .is_fixed()
+    {
+        constraints.insert(BuildConstraints::NetworkAccess);
+    }
 
-    let build_request = proto::BuildRequest {
+    Ok(BuildRequest {
         // Importantly, this must match the order of get_refscan_needles, since users may use that
         // function to map back from the found needles to a store path
         refscan_needles: get_refscan_needles(derivation)
             .map(|path| nixbase32::encode(path.digest()))
             .collect(),
         command_args,
-        outputs: output_paths,
+
+        outputs: {
+            // produce output_paths, which is the absolute path of each output (sorted)
+            let mut output_paths: Vec<PathBuf> = derivation
+                .outputs
+                .values()
+                .map(|e| PathBuf::from(e.path_str()[1..].to_owned()))
+                .collect();
+
+            // Sort the outputs. We can use sort_unstable, as these are unique strings.
+            output_paths.sort_unstable();
+            output_paths
+        },
 
         // Turn this into a sorted-by-key Vec<EnvVar>.
         environment_vars: environment_vars
@@ -130,7 +130,15 @@ pub(crate) fn derivation_to_build_request(
             .collect(),
         inputs: inputs
             .into_iter()
-            .map(|(name, node)| tvix_castore::proto::Node::from_name_and_node(name, node))
+            .map(|(path, node)| {
+                (
+                    path.to_string()
+                        .as_str()
+                        .try_into()
+                        .expect("Tvix bug: unable to convert store path basename to PathComponent"),
+                    node,
+                )
+            })
             .collect(),
         inputs_dir: nix_compat::store_path::STORE_DIR[1..].into(),
         constraints,
@@ -138,17 +146,12 @@ pub(crate) fn derivation_to_build_request(
         scratch_paths: vec!["build".into(), "nix/store".into()],
         additional_files: additional_files
             .into_iter()
-            .map(|(path, contents)| AdditionalFile { path, contents })
+            .map(|(path, contents)| AdditionalFile {
+                path: PathBuf::from(path),
+                contents,
+            })
             .collect(),
-    };
-
-    // FUTUREWORK: switch this function to construct the stricter BuildRequest directly.
-    debug_assert!(
-        BuildRequest::try_from(build_request.clone()).is_ok(),
-        "Tvix bug: BuildRequest would not be valid"
-    );
-
-    Ok(build_request)
+    })
 }
 
 /// handle passAsFile, if set.
@@ -212,15 +215,13 @@ fn calculate_pass_as_file_env(k: &str) -> (String, String) {
 #[cfg(test)]
 mod test {
     use bytes::Bytes;
-    use nix_compat::derivation::Derivation;
-    use std::collections::BTreeMap;
+    use nix_compat::{derivation::Derivation, store_path::StorePath};
+    use std::collections::{BTreeMap, HashSet};
     use std::sync::LazyLock;
-    use tvix_build::proto::{
-        build_request::{AdditionalFile, BuildConstraints, EnvVar},
-        BuildRequest,
-    };
     use tvix_castore::fixtures::DUMMY_DIGEST;
-    use tvix_castore::Node;
+    use tvix_castore::{Node, PathComponent};
+
+    use tvix_build::buildservice::{AdditionalFile, BuildConstraints, BuildRequest, EnvVar};
 
     use crate::tvix_build::NIX_ENVIRONMENT_VARS;
 
@@ -242,7 +243,10 @@ mod test {
 
         let build_request = derivation_to_build_request(
             &derivation,
-            BTreeMap::from([(INPUT_NODE_FOO_NAME.clone(), INPUT_NODE_FOO.clone())]),
+            BTreeMap::from([(
+                StorePath::<String>::from_bytes(&INPUT_NODE_FOO_NAME.clone()).unwrap(),
+                INPUT_NODE_FOO.clone(),
+            )]),
         )
         .expect("must succeed");
 
@@ -281,18 +285,15 @@ mod test {
                 command_args: vec![":".into()],
                 outputs: vec!["nix/store/fhaj6gmwns62s6ypkcldbaj2ybvkhx3p-foo".into()],
                 environment_vars: expected_environment_vars,
-                inputs: vec![tvix_castore::proto::Node::from_name_and_node(
-                    INPUT_NODE_FOO_NAME.clone(),
+                inputs: BTreeMap::from([(
+                    PathComponent::try_from(INPUT_NODE_FOO_NAME.clone()).unwrap(),
                     INPUT_NODE_FOO.clone()
-                )],
+                )]),
                 inputs_dir: "nix/store".into(),
-                constraints: Some(BuildConstraints {
-                    system: derivation.system.clone(),
-                    min_memory: 0,
-                    network_access: false,
-                    available_ro_paths: vec![],
-                    provide_bin_sh: true,
-                }),
+                constraints: HashSet::from([
+                    BuildConstraints::System(derivation.system.clone()),
+                    BuildConstraints::ProvideBinSh
+                ]),
                 additional_files: vec![],
                 working_dir: "build".into(),
                 scratch_paths: vec!["build".into(), "nix/store".into()],
@@ -357,15 +358,13 @@ mod test {
                 command_args: vec![":".to_string()],
                 outputs: vec!["nix/store/4q0pg5zpfmznxscq3avycvf9xdvx50n3-bar".into()],
                 environment_vars: expected_environment_vars,
-                inputs: vec![],
+                inputs: BTreeMap::new(),
                 inputs_dir: "nix/store".into(),
-                constraints: Some(BuildConstraints {
-                    system: derivation.system.clone(),
-                    min_memory: 0,
-                    network_access: true,
-                    available_ro_paths: vec![],
-                    provide_bin_sh: true,
-                }),
+                constraints: HashSet::from([
+                    BuildConstraints::System(derivation.system.clone()),
+                    BuildConstraints::NetworkAccess,
+                    BuildConstraints::ProvideBinSh
+                ]),
                 additional_files: vec![],
                 working_dir: "build".into(),
                 scratch_paths: vec!["build".into(), "nix/store".into()],
@@ -431,15 +430,12 @@ mod test {
                 command_args: vec![":".to_string()],
                 outputs: vec!["nix/store/pp17lwra2jkx8rha15qabg2q3wij72lj-foo".into()],
                 environment_vars: expected_environment_vars,
-                inputs: vec![],
+                inputs: BTreeMap::new(),
                 inputs_dir: "nix/store".into(),
-                constraints: Some(BuildConstraints {
-                    system: derivation.system.clone(),
-                    min_memory: 0,
-                    network_access: false,
-                    available_ro_paths: vec![],
-                    provide_bin_sh: true,
-                }),
+                constraints: HashSet::from([
+                    BuildConstraints::System(derivation.system.clone()),
+                    BuildConstraints::ProvideBinSh,
+                ]),
                 additional_files: vec![
                     // baz env
                     AdditionalFile {
diff --git a/tvix/glue/src/tvix_store_io.rs b/tvix/glue/src/tvix_store_io.rs
index c1fe53aa873f..67a88e13c54b 100644
--- a/tvix/glue/src/tvix_store_io.rs
+++ b/tvix/glue/src/tvix_store_io.rs
@@ -1,5 +1,4 @@
 //! This module provides an implementation of EvalIO talking to tvix-store.
-use bytes::Bytes;
 use futures::{StreamExt, TryStreamExt};
 use nix_compat::{nixhash::CAHash, store_path::StorePath};
 use std::collections::BTreeMap;
@@ -181,7 +180,7 @@ impl TvixStoreIO {
                         // derivation_to_build_request needs castore nodes for all inputs.
                         // Provide them, which means, here is where we recursively build
                         // all dependencies.
-                        let mut inputs: BTreeMap<Bytes, Node> =
+                        let mut inputs: BTreeMap<StorePath<String>, Node> =
                             futures::stream::iter(drv.input_derivations.iter())
                                 .map(|(input_drv_path, output_names)| {
                                     // look up the derivation object
@@ -224,7 +223,7 @@ impl TvixStoreIO {
                                                 .await?;
 
                                             if let Some(node) = node {
-                                                Ok((output_path.to_string().into(), node))
+                                                Ok((output_path, node))
                                             } else {
                                                 Err(io::Error::other("no node produced"))
                                             }
@@ -250,7 +249,7 @@ impl TvixStoreIO {
                                                 .store_path_to_node(&input_source, Path::new(""))
                                                 .await?;
                                             if let Some(node) = node {
-                                                Ok((input_source.to_string().into(), node))
+                                                Ok((input_source, node))
                                             } else {
                                                 Err(io::Error::other("no node produced"))
                                             }