1 files changed, 89 insertions, 0 deletions
diff --git a/tvix/cli/src/derivation.rs b/tvix/cli/src/derivation.rs
index ad8fec226bfd..2d951f7ae65c 100644
--- a/tvix/cli/src/derivation.rs
+++ b/tvix/cli/src/derivation.rs
@@ -1,9 +1,11 @@
 //! Implements `builtins.derivation`, the core of what makes Nix build packages.
 
+use std::collections::{btree_map, BTreeSet};
 use tvix_derivation::Derivation;
 use tvix_eval::{AddContext, ErrorKind, NixList, VM};
 
 use crate::errors::Error;
+use crate::known_paths::{KnownPaths, PathType};
 
 /// Helper function for populating the `drv.outputs` field from a
 /// manually specified set of outputs, instead of the default
@@ -30,6 +32,46 @@ fn populate_outputs(vm: &mut VM, drv: &mut Derivation, outputs: NixList) -> Resu
     Ok(())
 }
 
+/// Populate the inputs of a derivation from the build references
+/// found when scanning the derivation's parameters.
+fn populate_inputs<I: IntoIterator<Item = String>>(
+    drv: &mut Derivation,
+    known_paths: &KnownPaths,
+    references: I,
+) {
+    for reference in references.into_iter() {
+        match &known_paths[&reference] {
+            PathType::Plain => {
+                drv.input_sources.insert(reference.to_string());
+            }
+
+            PathType::Output { name, derivation } => {
+                match drv.input_derivations.entry(derivation.clone()) {
+                    btree_map::Entry::Vacant(entry) => {
+                        entry.insert(BTreeSet::from([name.clone()]));
+                    }
+
+                    btree_map::Entry::Occupied(mut entry) => {
+                        entry.get_mut().insert(name.clone());
+                    }
+                }
+            }
+
+            PathType::Derivation { output_names } => {
+                match drv.input_derivations.entry(reference.to_string()) {
+                    btree_map::Entry::Vacant(entry) => {
+                        entry.insert(output_names.clone());
+                    }
+
+                    btree_map::Entry::Occupied(mut entry) => {
+                        entry.get_mut().extend(output_names.clone().into_iter());
+                    }
+                }
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -85,4 +127,51 @@ mod tests {
         populate_outputs(&mut vm, &mut drv, outputs)
             .expect_err("supplying duplicate outputs should fail");
     }
+
+    #[test]
+    fn populate_inputs_empty() {
+        let mut drv = Derivation::default();
+        let paths = KnownPaths::default();
+        let inputs = vec![];
+
+        populate_inputs(&mut drv, &paths, inputs);
+
+        assert!(drv.input_sources.is_empty());
+        assert!(drv.input_derivations.is_empty());
+    }
+
+    #[test]
+    fn populate_inputs_all() {
+        let mut drv = Derivation::default();
+
+        let mut paths = KnownPaths::default();
+        paths.plain("/nix/store/fn7zvafq26f0c8b17brs7s95s10ibfzs-foo");
+        paths.drv(
+            "/nix/store/aqffiyqx602lbam7n1zsaz3yrh6v08pc-bar.drv",
+            &["out"],
+        );
+        paths.output(
+            "/nix/store/zvpskvjwi72fjxg0vzq822sfvq20mq4l-bar",
+            "out",
+            "/nix/store/aqffiyqx602lbam7n1zsaz3yrh6v08pc-bar.drv",
+        );
+
+        let inputs: Vec<String> = vec![
+            "/nix/store/fn7zvafq26f0c8b17brs7s95s10ibfzs-foo".into(),
+            "/nix/store/aqffiyqx602lbam7n1zsaz3yrh6v08pc-bar.drv".into(),
+            "/nix/store/zvpskvjwi72fjxg0vzq822sfvq20mq4l-bar".into(),
+        ];
+
+        populate_inputs(&mut drv, &paths, inputs);
+
+        assert_eq!(drv.input_sources.len(), 1);
+        assert!(drv
+            .input_sources
+            .contains("/nix/store/fn7zvafq26f0c8b17brs7s95s10ibfzs-foo"));
+
+        assert_eq!(drv.input_derivations.len(), 1);
+        assert!(drv
+            .input_derivations
+            .contains_key("/nix/store/aqffiyqx602lbam7n1zsaz3yrh6v08pc-bar.drv"));
+    }
 }