From a35dadf9f0c7c98d2ae8761428936a0b9d30489e Mon Sep 17 00:00:00 2001
From: Vincent Ambo <mail@tazj.in>
Date: Fri, 13 Jan 2023 20:19:48 +0300
Subject: refactor(tvix/derivation): use BTreeSet for derivation outputs

When constructing derivations inside builtins.derivationStrict, we'd
have to very frequently check whether certain outputs have already
been inserted into the derivation inputs.

Using a set type is much easier, especially as this has to be ordered
and the source data that is being inserted also comes from a set,
which might let us pass this more efficiently in the future.

Note that the validate function no longer checks the order of the
entries, as that is now guaranteed by the type.

Change-Id: I2fbb984facba3e668075f6f8df8992092368c63d
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7826
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
---
 tvix/derivation/src/derivation.rs |  8 +++++---
 tvix/derivation/src/tests/mod.rs  |  8 +++++---
 tvix/derivation/src/validate.rs   | 10 +---------
 tvix/derivation/src/write.rs      |  3 ++-
 4 files changed, 13 insertions(+), 16 deletions(-)

(limited to 'tvix/derivation/src')

diff --git a/tvix/derivation/src/derivation.rs b/tvix/derivation/src/derivation.rs
index 54e0bbf82ec2..d1f46de8c5ee 100644
--- a/tvix/derivation/src/derivation.rs
+++ b/tvix/derivation/src/derivation.rs
@@ -3,6 +3,7 @@ use crate::output::{Hash, Output};
 use crate::write;
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
+use std::collections::BTreeSet;
 use std::{collections::BTreeMap, fmt, fmt::Write};
 use tvix_store::nixbase32::NIXBASE32;
 use tvix_store::store_path::{ParseStorePathError, StorePath, STORE_DIR};
@@ -18,7 +19,7 @@ pub struct Derivation {
     pub environment: BTreeMap<String, String>,
 
     #[serde(rename = "inputDrvs")]
-    pub input_derivations: BTreeMap<String, Vec<String>>,
+    pub input_derivations: BTreeMap<String, BTreeSet<String>>,
 
     #[serde(rename = "inputSrcs")]
     pub input_sources: Vec<String>,
@@ -173,13 +174,14 @@ impl Derivation {
                 hasher.finalize()
             }
             None => {
-                let mut replaced_input_derivations: BTreeMap<String, Vec<String>> = BTreeMap::new();
+                let mut replaced_input_derivations: BTreeMap<String, BTreeSet<String>> =
+                    BTreeMap::new();
 
                 // For each input_derivation, look up the replacement.
                 for (drv_path, input_derivation) in &self.input_derivations {
                     replaced_input_derivations.insert(
                         fn_get_drv_replacement(drv_path).to_string(),
-                        input_derivation.to_vec(),
+                        input_derivation.clone(),
                     );
                 }
 
diff --git a/tvix/derivation/src/tests/mod.rs b/tvix/derivation/src/tests/mod.rs
index cb98a27cdfe4..bd5d9f34cb12 100644
--- a/tvix/derivation/src/tests/mod.rs
+++ b/tvix/derivation/src/tests/mod.rs
@@ -1,5 +1,6 @@
 use crate::derivation::Derivation;
 use crate::output::{Hash, Output};
+use std::collections::BTreeSet;
 use std::fs::File;
 use std::io::Read;
 use std::path::Path;
@@ -272,9 +273,10 @@ fn output_path_construction() {
     );
 
     // assemble foo input_derivations
-    foo_drv
-        .input_derivations
-        .insert(bar_drv_path.to_absolute_path(), vec!["out".to_string()]);
+    foo_drv.input_derivations.insert(
+        bar_drv_path.to_absolute_path(),
+        BTreeSet::from(["out".to_string()]),
+    );
 
     // calculate foo output paths
     let foo_calc_result = foo_drv.calculate_output_paths(
diff --git a/tvix/derivation/src/validate.rs b/tvix/derivation/src/validate.rs
index eda0457ae564..5335f43636d3 100644
--- a/tvix/derivation/src/validate.rs
+++ b/tvix/derivation/src/validate.rs
@@ -50,21 +50,13 @@ impl Derivation {
                 );
             }
 
-            for (i, output_name) in output_names.iter().enumerate() {
+            for output_name in output_names.iter() {
                 if output_name.is_empty() {
                     bail!(
                         "output name entry for {} may not be empty",
                         input_derivation_path
                     )
                 }
-                // if i is at least 1, peek at the previous element to ensure output_names are sorted.
-                if i > 0 && (output_names[i - 1] >= *output_name) {
-                    bail!(
-                        "invalid input derivation output order: {} < {}",
-                        output_name,
-                        output_names[i - 1],
-                    );
-                }
             }
         }
 
diff --git a/tvix/derivation/src/write.rs b/tvix/derivation/src/write.rs
index 842b2c4ab6ad..8e9cdfcccaf6 100644
--- a/tvix/derivation/src/write.rs
+++ b/tvix/derivation/src/write.rs
@@ -1,5 +1,6 @@
 use crate::output::Output;
 use crate::string_escape::escape_string;
+use std::collections::BTreeSet;
 use std::{collections::BTreeMap, fmt, fmt::Write};
 
 pub const DERIVATION_PREFIX: &str = "Derive";
@@ -84,7 +85,7 @@ pub fn write_outputs(
 
 pub fn write_input_derivations(
     writer: &mut impl Write,
-    input_derivations: &BTreeMap<String, Vec<String>>,
+    input_derivations: &BTreeMap<String, BTreeSet<String>>,
 ) -> Result<(), fmt::Error> {
     writer.write_char(COMMA)?;
     writer.write_char(BRACKET_OPEN)?;
-- 
cgit 1.4.1