refactor(tvix/derivation): refactor the derivation serialization r/5564

This refactors the code to serialize a derivation. The original code
has beed moved to seperate crates for better code structure.

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

6 files changed, 248 insertions, 215 deletions

diff --git a/tvix/derivation/src/derivation.rs b/tvix/derivation/src/derivation.rs
new file mode 100644
index 000000000000..142fa3c9f309
--- /dev/null
+++ b/tvix/derivation/src/derivation.rs
@@ -0,0 +1,34 @@
+use crate::output::Output;
+use crate::write;
+use serde::{Deserialize, Serialize};
+use std::{collections::BTreeMap, fmt, fmt::Write};
+
+#[derive(Serialize, Deserialize)]
+pub struct Derivation {
+    outputs: BTreeMap<String, Output>,
+    input_sources: Vec<String>,
+    input_derivations: BTreeMap<String, Vec<String>>,
+    platform: String,
+    builder: String,
+    arguments: Vec<String>,
+    environment: BTreeMap<String, String>,
+}
+
+impl Derivation {
+    pub fn serialize(self: Self, writer: &mut impl Write) -> Result<(), fmt::Error> {
+        writer.write_str(write::DERIVATION_PREFIX)?;
+        writer.write_char(write::PAREN_OPEN)?;
+
+        write::write_outputs(writer, self.outputs)?;
+        write::write_input_derivations(writer, self.input_derivations)?;
+        write::write_input_sources(writer, self.input_sources)?;
+        write::write_platfrom(writer, &self.platform)?;
+        write::write_builder(writer, &self.builder)?;
+        write::write_arguments(writer, self.arguments)?;
+        write::write_enviroment(writer, self.environment)?;
+
+        writer.write_char(write::PAREN_CLOSE)?;
+
+        Ok(())
+    }
+}
diff --git a/tvix/derivation/src/lib.rs b/tvix/derivation/src/lib.rs
index a0ddee410b80..a902943493d5 100644
--- a/tvix/derivation/src/lib.rs
+++ b/tvix/derivation/src/lib.rs
@@ -1,216 +1,8 @@
-use serde::{Deserialize, Serialize};
-use std::{collections::BTreeMap, fmt, fmt::Write};
+mod output;
+mod string_escape;
+mod write;
+
+mod derivation;
 
 #[cfg(test)]
 mod tests;
-
-const DERIVATION_PREFIX: &str = "Derive";
-const PAREN_OPEN: char = '(';
-const PAREN_CLOSE: char = ')';
-const BRACKET_OPEN: char = '[';
-const BRACKET_CLOSE: char = ']';
-const COMMA: char = ',';
-const QUOTE: char = '"';
-
-const STRING_ESCAPER: [(char, &str); 5] = [
-    ('\\', "\\\\"),
-    ('\n', "\\n"),
-    ('\r', "\\r"),
-    ('\t', "\\t"),
-    ('\"', "\\\""),
-];
-
-fn default_resource() -> String {
-    "".to_string()
-}
-
-#[derive(Serialize, Deserialize)]
-pub struct Output {
-    path: String,
-    #[serde(default = "default_resource")]
-    hash_algorithm: String,
-    #[serde(default = "default_resource")]
-    hash: String,
-}
-
-#[derive(Serialize, Deserialize)]
-pub struct Derivation {
-    outputs: BTreeMap<String, Output>,
-    input_sources: Vec<String>,
-    input_derivations: BTreeMap<String, Vec<String>>,
-    platform: String,
-    builder: String,
-    arguments: Vec<String>,
-    environment: BTreeMap<String, String>,
-}
-
-fn escape_string(s: &String) -> String {
-    let mut s_replaced = s.clone();
-
-    for escape_sequence in STRING_ESCAPER {
-        s_replaced = s_replaced.replace(escape_sequence.0, escape_sequence.1);
-    }
-
-    return format!("\"{}\"", s_replaced);
-}
-
-fn write_array_elements(
-    writer: &mut impl Write,
-    quote: bool,
-    open: &str,
-    closing: &str,
-    elements: Vec<&String>,
-) -> Result<(), fmt::Error> {
-    writer.write_str(open)?;
-
-    for (index, element) in elements.iter().enumerate() {
-        if index > 0 {
-            writer.write_char(COMMA)?;
-        }
-
-        if quote {
-            writer.write_char(QUOTE)?;
-        }
-
-        writer.write_str(element)?;
-
-        if quote {
-            writer.write_char(QUOTE)?;
-        }
-    }
-
-    writer.write_str(closing)?;
-
-    return Ok(());
-}
-
-pub fn serialize_derivation(
-    derivation: Derivation,
-    writer: &mut impl Write,
-) -> Result<(), fmt::Error> {
-    writer.write_str(DERIVATION_PREFIX)?;
-    writer.write_char(PAREN_OPEN)?;
-
-    // Step 1: Write outputs
-    {
-        writer.write_char(BRACKET_OPEN)?;
-        for (ii, (output_name, output)) in derivation.outputs.iter().enumerate() {
-            if ii > 0 {
-                writer.write_char(COMMA)?;
-            }
-
-            // TODO(jrhahn) option to strip output
-            let elements = vec![
-                output_name,
-                &output.path,
-                &output.hash_algorithm,
-                &output.hash,
-            ];
-
-            write_array_elements(
-                writer,
-                true,
-                &PAREN_OPEN.to_string(),
-                &PAREN_CLOSE.to_string(),
-                elements,
-            )?
-        }
-        writer.write_char(BRACKET_CLOSE)?;
-    }
-
-    // Step 2: Write input_derivations
-    {
-        writer.write_char(COMMA)?;
-        writer.write_char(BRACKET_OPEN)?;
-
-        for (ii, (input_derivation_path, input_derivation)) in
-            derivation.input_derivations.iter().enumerate()
-        {
-            if ii > 0 {
-                writer.write_char(COMMA)?;
-            }
-
-            writer.write_char(PAREN_OPEN)?;
-            writer.write_char(QUOTE)?;
-            writer.write_str(input_derivation_path.as_str())?;
-            writer.write_char(QUOTE)?;
-            writer.write_char(COMMA)?;
-
-            write_array_elements(
-                writer,
-                true,
-                &BRACKET_OPEN.to_string(),
-                &BRACKET_CLOSE.to_string(),
-                input_derivation.iter().map(|s| s).collect(),
-            )?;
-
-            writer.write_char(PAREN_CLOSE)?;
-        }
-
-        writer.write_char(BRACKET_CLOSE)?;
-    }
-
-    // Step 3: Write input_sources
-    {
-        writer.write_char(COMMA)?;
-        write_array_elements(
-            writer,
-            true,
-            &BRACKET_OPEN.to_string(),
-            &BRACKET_CLOSE.to_string(),
-            derivation.input_sources.iter().map(|s| s).collect(),
-        )?;
-    }
-
-    // Step 4: Write platform
-    {
-        writer.write_char(COMMA)?;
-        writer.write_str(&escape_string(&derivation.platform).as_str())?;
-    }
-
-    // Step 5: Write builder
-    {
-        writer.write_char(COMMA)?;
-        writer.write_str(&escape_string(&derivation.builder).as_str())?;
-    }
-
-    // Step 6: Write arguments
-    {
-        writer.write_char(COMMA)?;
-        write_array_elements(
-            writer,
-            true,
-            &BRACKET_OPEN.to_string(),
-            &BRACKET_CLOSE.to_string(),
-            derivation.arguments.iter().map(|s| s).collect(),
-        )?;
-    }
-
-    // Step 7: Write env
-    {
-        writer.write_char(COMMA)?;
-        writer.write_char(BRACKET_OPEN)?;
-
-        for (ii, (key, environment)) in derivation.environment.iter().enumerate() {
-            if ii > 0 {
-                writer.write_char(COMMA)?;
-            }
-
-            // TODO(jrhahn) add strip option
-            write_array_elements(
-                writer,
-                false,
-                &PAREN_OPEN.to_string(),
-                &PAREN_CLOSE.to_string(),
-                vec![&escape_string(key), &escape_string(&environment)],
-            )?;
-        }
-
-        writer.write_char(BRACKET_CLOSE)?;
-    }
-
-    // Step 8: Close Derive call
-    writer.write_char(PAREN_CLOSE)?;
-
-    return Ok(());
-}
diff --git a/tvix/derivation/src/output.rs b/tvix/derivation/src/output.rs
new file mode 100644
index 000000000000..0d764011fc1c
--- /dev/null
+++ b/tvix/derivation/src/output.rs
@@ -0,0 +1,16 @@
+use serde::{Deserialize, Serialize};
+
+// This function is required by serde to deserialize files
+// with missing keys.
+fn default_resource() -> String {
+    "".to_string()
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct Output {
+    pub path: String,
+    #[serde(default = "default_resource")]
+    pub hash_algorithm: String,
+    #[serde(default = "default_resource")]
+    pub hash: String,
+}
diff --git a/tvix/derivation/src/string_escape.rs b/tvix/derivation/src/string_escape.rs
new file mode 100644
index 000000000000..0e1dbe516f73
--- /dev/null
+++ b/tvix/derivation/src/string_escape.rs
@@ -0,0 +1,17 @@
+const STRING_ESCAPER: [(char, &str); 5] = [
+    ('\\', "\\\\"),
+    ('\n', "\\n"),
+    ('\r', "\\r"),
+    ('\t', "\\t"),
+    ('\"', "\\\""),
+];
+
+pub fn escape_string(s: &str) -> String {
+    let mut s_replaced = s.to_string();
+
+    for escape_sequence in STRING_ESCAPER {
+        s_replaced = s_replaced.replace(escape_sequence.0, escape_sequence.1);
+    }
+
+    format!("\"{}\"", s_replaced)
+}
diff --git a/tvix/derivation/src/tests/mod.rs b/tvix/derivation/src/tests/mod.rs
index 9c8ffd061c11..800806926976 100644
--- a/tvix/derivation/src/tests/mod.rs
+++ b/tvix/derivation/src/tests/mod.rs
@@ -1,4 +1,4 @@
-use super::{serialize_derivation, Derivation};
+use crate::derivation::Derivation;
 use std::fs::File;
 use std::io::Read;
 use std::path::Path;
@@ -19,7 +19,7 @@ fn assert_derivation_ok(path_to_drv_file: &str) {
     let derivation: Derivation = serde_json::from_str(&data).expect("JSON was not well-formatted");
 
     let mut serialized_derivation = String::new();
-    serialize_derivation(derivation, &mut serialized_derivation).unwrap();
+    derivation.serialize(&mut serialized_derivation).unwrap();
 
     let expected = read_file(path_to_drv_file);
 
diff --git a/tvix/derivation/src/write.rs b/tvix/derivation/src/write.rs
new file mode 100644
index 000000000000..8e5899b9ec32
--- /dev/null
+++ b/tvix/derivation/src/write.rs
@@ -0,0 +1,174 @@
+use crate::output::Output;
+use crate::string_escape::escape_string;
+use std::{collections::BTreeMap, fmt, fmt::Write};
+
+pub const DERIVATION_PREFIX: &str = "Derive";
+pub const PAREN_OPEN: char = '(';
+pub const PAREN_CLOSE: char = ')';
+pub const BRACKET_OPEN: char = '[';
+pub const BRACKET_CLOSE: char = ']';
+pub const COMMA: char = ',';
+pub const QUOTE: char = '"';
+
+fn write_array_elements(
+    writer: &mut impl Write,
+    quote: bool,
+    open: &str,
+    closing: &str,
+    elements: Vec<&String>,
+) -> Result<(), fmt::Error> {
+    writer.write_str(open)?;
+
+    for (index, element) in elements.iter().enumerate() {
+        if index > 0 {
+            writer.write_char(COMMA)?;
+        }
+
+        if quote {
+            writer.write_char(QUOTE)?;
+        }
+
+        writer.write_str(element)?;
+
+        if quote {
+            writer.write_char(QUOTE)?;
+        }
+    }
+
+    writer.write_str(closing)?;
+
+    Ok(())
+}
+
+pub fn write_outputs(
+    writer: &mut impl Write,
+    outputs: BTreeMap<String, Output>,
+) -> Result<(), fmt::Error> {
+    writer.write_char(BRACKET_OPEN)?;
+    for (ii, (output_name, output)) in outputs.iter().enumerate() {
+        if ii > 0 {
+            writer.write_char(COMMA)?;
+        }
+
+        // TODO(jrhahn) option to strip output
+        let elements = vec![
+            output_name,
+            &output.path,
+            &output.hash_algorithm,
+            &output.hash,
+        ];
+
+        write_array_elements(
+            writer,
+            true,
+            &PAREN_OPEN.to_string(),
+            &PAREN_CLOSE.to_string(),
+            elements,
+        )?
+    }
+    writer.write_char(BRACKET_CLOSE)?;
+
+    Ok(())
+}
+
+pub fn write_input_derivations(
+    writer: &mut impl Write,
+    input_derivations: BTreeMap<String, Vec<String>>,
+) -> Result<(), fmt::Error> {
+    writer.write_char(COMMA)?;
+    writer.write_char(BRACKET_OPEN)?;
+
+    for (ii, (input_derivation_path, input_derivation)) in input_derivations.iter().enumerate() {
+        if ii > 0 {
+            writer.write_char(COMMA)?;
+        }
+
+        writer.write_char(PAREN_OPEN)?;
+        writer.write_char(QUOTE)?;
+        writer.write_str(input_derivation_path.as_str())?;
+        writer.write_char(QUOTE)?;
+        writer.write_char(COMMA)?;
+
+        write_array_elements(
+            writer,
+            true,
+            &BRACKET_OPEN.to_string(),
+            &BRACKET_CLOSE.to_string(),
+            input_derivation.iter().collect(),
+        )?;
+
+        writer.write_char(PAREN_CLOSE)?;
+    }
+
+    writer.write_char(BRACKET_CLOSE)?;
+
+    Ok(())
+}
+
+pub fn write_input_sources(
+    writer: &mut impl Write,
+    input_sources: Vec<String>,
+) -> Result<(), fmt::Error> {
+    writer.write_char(COMMA)?;
+    write_array_elements(
+        writer,
+        true,
+        &BRACKET_OPEN.to_string(),
+        &BRACKET_CLOSE.to_string(),
+        input_sources.iter().collect(),
+    )?;
+
+    Ok(())
+}
+
+pub fn write_platfrom(writer: &mut impl Write, platform: &str) -> Result<(), fmt::Error> {
+    writer.write_char(COMMA)?;
+    writer.write_str(escape_string(platform).as_str())?;
+    Ok(())
+}
+
+pub fn write_builder(writer: &mut impl Write, builder: &str) -> Result<(), fmt::Error> {
+    writer.write_char(COMMA)?;
+    writer.write_str(escape_string(builder).as_str())?;
+    Ok(())
+}
+
+pub fn write_arguments(writer: &mut impl Write, arguments: Vec<String>) -> Result<(), fmt::Error> {
+    writer.write_char(COMMA)?;
+    write_array_elements(
+        writer,
+        true,
+        &BRACKET_OPEN.to_string(),
+        &BRACKET_CLOSE.to_string(),
+        arguments.iter().collect(),
+    )?;
+
+    Ok(())
+}
+
+pub fn write_enviroment(
+    writer: &mut impl Write,
+    environment: BTreeMap<String, String>,
+) -> Result<(), fmt::Error> {
+    writer.write_char(COMMA)?;
+    writer.write_char(BRACKET_OPEN)?;
+
+    for (ii, (key, environment)) in environment.iter().enumerate() {
+        if ii > 0 {
+            writer.write_char(COMMA)?;
+        }
+
+        // TODO(jrhahn) add strip option
+        write_array_elements(
+            writer,
+            false,
+            &PAREN_OPEN.to_string(),
+            &PAREN_CLOSE.to_string(),
+            vec![&escape_string(key), &escape_string(environment)],
+        )?;
+    }
+
+    writer.write_char(BRACKET_CLOSE)?;
+
+    Ok(())
+}