feat(tvix/glue): wire up nix refscanning r/8743

After this, attempting to build the nixpkgs still fails in the same way,
because the references are not yet properly used by the code at
`tvix/glue/src/tvix_store_io.rs`.

Change-Id: I8a59ef8ef3c9a6f6aa7b05106dd9eef2e9ac0d0f
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12532
Reviewed-by: Brian Olsen <me@griff.name>
Autosubmit: yuka <yuka@yuka.dev>
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>

2 files changed, 73 insertions, 12 deletions

diff --git a/tvix/glue/src/tvix_build.rs b/tvix/glue/src/tvix_build.rs
index 5187a76d48e5..7a901a593047 100644
--- a/tvix/glue/src/tvix_build.rs
+++ b/tvix/glue/src/tvix_build.rs
@@ -4,7 +4,7 @@
 use std::collections::BTreeMap;
 
 use bytes::Bytes;
-use nix_compat::{derivation::Derivation, nixbase32};
+use nix_compat::{derivation::Derivation, nixbase32, store_path::StorePath};
 use sha2::{Digest, Sha256};
 use tvix_build::proto::{
     build_request::{AdditionalFile, BuildConstraints, EnvVar},
@@ -29,6 +29,20 @@ const NIX_ENVIRONMENT_VARS: [(&str, &str); 12] = [
     ("TMPDIR", "/build"),
 ];
 
+/// 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
+/// 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,
+) -> impl Iterator<Item = &StorePath<String>> {
+    derivation
+        .outputs
+        .values()
+        .filter_map(|output| output.path.as_ref())
+        .chain(derivation.input_sources.iter())
+        .chain(derivation.input_derivations.keys())
+}
+
 /// Takes a [Derivation] and turns it into a [BuildRequest].
 /// It assumes the Derivation has been validated.
 /// It needs two lookup functions:
@@ -100,7 +114,11 @@ pub(crate) fn derivation_to_build_request(
     });
 
     let build_request = BuildRequest {
-        refscan_needles: vec![], // TODO refscan
+        // 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,
 
@@ -277,7 +295,10 @@ mod test {
                 additional_files: vec![],
                 working_dir: "build".into(),
                 scratch_paths: vec!["build".into(), "nix/store".into()],
-                refscan_needles: vec![],
+                refscan_needles: vec![
+                    "fhaj6gmwns62s6ypkcldbaj2ybvkhx3p".into(),
+                    "ss2p4wmxijn652haqyd7dckxwl4c7hxx".into()
+                ],
             },
             build_request
         );
@@ -347,7 +368,7 @@ mod test {
                 additional_files: vec![],
                 working_dir: "build".into(),
                 scratch_paths: vec!["build".into(), "nix/store".into()],
-                refscan_needles: vec![],
+                refscan_needles: vec!["4q0pg5zpfmznxscq3avycvf9xdvx50n3".into()],
             },
             build_request
         );
@@ -434,7 +455,7 @@ mod test {
                 ],
                 working_dir: "build".into(),
                 scratch_paths: vec!["build".into(), "nix/store".into()],
-                refscan_needles: vec![],
+                refscan_needles: vec!["pp17lwra2jkx8rha15qabg2q3wij72lj".into()],
             },
             build_request
         );
diff --git a/tvix/glue/src/tvix_store_io.rs b/tvix/glue/src/tvix_store_io.rs
index 93675b438f25..8a11e293f0ac 100644
--- a/tvix/glue/src/tvix_store_io.rs
+++ b/tvix/glue/src/tvix_store_io.rs
@@ -296,12 +296,37 @@ impl TvixStoreIO {
                             .await
                             .map_err(|e| std::io::Error::new(io::ErrorKind::Other, e))?;
 
-                        // TODO: refscan
+                        // Maps from the index in refscan_needles to the full store path
+                        // Used to map back to the actual store path from the found needles
+                        // Importantly, this must match the order of the needles generated in derivation_to_build_request
+                        let refscan_needles =
+                            crate::tvix_build::get_refscan_needles(&drv).collect::<Vec<_>>();
 
                         // For each output, insert a PathInfo.
-                        for output in &build_result.outputs {
-                            let (output_name, output_node) =
-                                output.clone().into_name_and_node().expect("invalid node");
+                        for ((output, output_needles), drv_output) in build_result
+                            .outputs
+                            .iter()
+                            .zip(build_result.outputs_needles.iter())
+                            .zip(drv.outputs.iter())
+                        {
+                            let (_, output_node) = output
+                                .clone()
+                                .into_name_bytes_and_node()
+                                .expect("invalid node");
+
+                            let output_needles: Vec<_> = output_needles
+                                .needles
+                                .iter()
+                                // Map each output needle index back to the refscan_needle
+                                .map(|idx| {
+                                    refscan_needles
+                                        .get(*idx as usize)
+                                        .ok_or(std::io::Error::new(
+                                            std::io::ErrorKind::Other,
+                                            "invalid build response",
+                                        ))
+                                })
+                                .collect::<Result<_, std::io::Error>>()?;
 
                             // calculate the nar representation
                             let (nar_size, nar_sha256) = self
@@ -312,15 +337,30 @@ impl TvixStoreIO {
                             // assemble the PathInfo to persist
                             let path_info = PathInfo {
                                 node: Some(tvix_castore::proto::Node::from_name_and_node(
-                                    output_name.into(),
+                                    drv_output
+                                        .1
+                                        .path
+                                        .as_ref()
+                                        .ok_or(std::io::Error::new(
+                                            std::io::ErrorKind::Other,
+                                            "missing output store path",
+                                        ))?
+                                        .to_string()
+                                        .into(),
                                     output_node,
                                 )),
-                                references: vec![], // TODO: refscan
+                                references: output_needles
+                                    .iter()
+                                    .map(|path| Bytes::from(path.digest().as_slice().to_vec()))
+                                    .collect(),
                                 narinfo: Some(tvix_store::proto::NarInfo {
                                     nar_size,
                                     nar_sha256: Bytes::from(nar_sha256.to_vec()),
                                     signatures: vec![],
-                                    reference_names: vec![], // TODO: refscan
+                                    reference_names: output_needles
+                                        .iter()
+                                        .map(|path| path.to_string())
+                                        .collect(),
                                     deriver: Some(tvix_store::proto::StorePath {
                                         name: drv_path
                                             .name()