refactor(tvix/nix-compat/derivation): simplify r/6450

Let the escape function only take care of string escaping, not quoting.

Let write_array_elements always quote and escape strings it consumes.
Move the business of writing additional wrapping characters around it to
the caller.

Change-Id: Ib8dea69c409561b49862c531ba5a3fe6c2f061f8
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8993
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI

3 files changed, 51 insertions, 64 deletions

diff --git a/tvix/nix-compat/src/derivation/escape.rs b/tvix/nix-compat/src/derivation/escape.rs
index 03106c4420..8c2c6fce07 100644
--- a/tvix/nix-compat/src/derivation/escape.rs
+++ b/tvix/nix-compat/src/derivation/escape.rs
@@ -1,7 +1,8 @@
 use bstr::{BString, ByteSlice};
 
-pub fn escape_bstr(s: &[u8]) -> BString {
-    let mut s: Vec<u8> = s.to_owned();
+/// Escapes a byte sequence. Does not add surrounding quotes.
+pub fn escape_bstr<P: AsRef<[u8]>>(s: P) -> BString {
+    let mut s: Vec<u8> = s.as_ref().to_vec();
 
     s = s.replace(b"\\", b"\\\\");
     s = s.replace(b"\n", b"\\n");
@@ -9,12 +10,7 @@ pub fn escape_bstr(s: &[u8]) -> BString {
     s = s.replace(b"\t", b"\\t");
     s = s.replace(b"\"", b"\\\"");
 
-    let mut out: Vec<u8> = Vec::new();
-    out.push(b'\"');
-    out.append(&mut s);
-    out.push(b'\"');
-
-    out.into()
+    s.into()
 }
 
 #[cfg(test)]
@@ -22,9 +18,9 @@ mod tests {
     use super::escape_bstr;
     use test_case::test_case;
 
-    #[test_case(b"", b"\"\""; "empty")]
-    #[test_case(b"\"", b"\"\\\"\""; "doublequote")]
-    #[test_case(b":", b"\":\""; "colon")]
+    #[test_case(b"", b""; "empty")]
+    #[test_case(b"\"", b"\\\""; "doublequote")]
+    #[test_case(b":", b":"; "colon")]
     fn escape(input: &[u8], expected: &[u8]) {
         assert_eq!(expected, escape_bstr(input))
     }
diff --git a/tvix/nix-compat/src/derivation/mod.rs b/tvix/nix-compat/src/derivation/mod.rs
index 498aa1b59d..b416e745a5 100644
--- a/tvix/nix-compat/src/derivation/mod.rs
+++ b/tvix/nix-compat/src/derivation/mod.rs
@@ -5,6 +5,7 @@ use bstr::BString;
 use serde::{Deserialize, Serialize};
 use sha2::{Digest, Sha256};
 use std::collections::{BTreeMap, BTreeSet};
+use std::io;
 
 mod errors;
 mod escape;
@@ -45,8 +46,8 @@ impl Derivation {
     /// write the Derivation to the given [std::io::Write], in ATerm format.
     ///
     /// The only errors returns are these when writing to the passed writer.
-    pub fn serialize(&self, writer: &mut impl std::io::Write) -> Result<(), std::io::Error> {
-        write::write_str(writer, write::DERIVATION_PREFIX)?;
+    pub fn serialize(&self, writer: &mut impl std::io::Write) -> Result<(), io::Error> {
+        io::copy(&mut io::Cursor::new(write::DERIVATION_PREFIX), writer)?;
         write::write_char(writer, write::PAREN_OPEN)?;
 
         write::write_outputs(writer, &self.outputs)?;
diff --git a/tvix/nix-compat/src/derivation/write.rs b/tvix/nix-compat/src/derivation/write.rs
index cf62f85022..464b02fef1 100644
--- a/tvix/nix-compat/src/derivation/write.rs
+++ b/tvix/nix-compat/src/derivation/write.rs
@@ -26,39 +26,37 @@ pub(crate) fn write_char(writer: &mut impl Write, c: char) -> io::Result<()> {
     Ok(())
 }
 
-// Writes a string to the writer (as unicode)
-pub(crate) fn write_str(writer: &mut impl Write, s: &str) -> io::Result<()> {
-    io::copy(&mut Cursor::new(s.as_bytes()), writer)?;
+// Write a string `s` as a quoted field to the writer.
+// The `escape` argument controls whether escaping will be skipped.
+// This is the case if `s` is known to only contain characters that need no
+// escaping.
+pub(crate) fn write_field<S: AsRef<[u8]>>(
+    writer: &mut impl Write,
+    s: S,
+    escape: bool,
+) -> io::Result<()> {
+    write_char(writer, QUOTE)?;
+
+    if !escape {
+        io::copy(&mut Cursor::new(s), writer)?;
+    } else {
+        io::copy(&mut Cursor::new(escape_bstr(s.as_ref())), writer)?;
+    }
+
+    write_char(writer, QUOTE)?;
+
     Ok(())
 }
 
-fn write_array_elements(
-    writer: &mut impl Write,
-    quote: bool,
-    open: &str,
-    closing: &str,
-    elements: &[BString],
-) -> Result<(), io::Error> {
-    write_str(writer, open)?;
-
+fn write_array_elements(writer: &mut impl Write, elements: &[BString]) -> Result<(), io::Error> {
     for (index, element) in elements.iter().enumerate() {
         if index > 0 {
             write_char(writer, COMMA)?;
         }
 
-        if quote {
-            write_char(writer, QUOTE)?;
-        }
-
-        io::copy(&mut Cursor::new(element), writer)?;
-
-        if quote {
-            write_char(writer, QUOTE)?;
-        }
+        write_field(writer, element, true)?;
     }
 
-    write_str(writer, closing)?;
-
     Ok(())
 }
 
@@ -94,13 +92,11 @@ pub fn write_outputs(
         elements.push(e2.into());
         elements.push(e3.into());
 
-        write_array_elements(
-            writer,
-            true,
-            &PAREN_OPEN.to_string(),
-            &PAREN_CLOSE.to_string(),
-            &elements,
-        )?
+        write_char(writer, PAREN_OPEN)?;
+
+        write_array_elements(writer, &elements)?;
+
+        write_char(writer, PAREN_CLOSE)?;
     }
     write_char(writer, BRACKET_CLOSE)?;
 
@@ -121,21 +117,18 @@ pub fn write_input_derivations(
         }
 
         write_char(writer, PAREN_OPEN)?;
-        write_char(writer, QUOTE)?;
-        write_str(writer, input_derivation_path.as_str())?;
-        write_char(writer, QUOTE)?;
+        write_field(writer, input_derivation_path.as_str(), false)?;
         write_char(writer, COMMA)?;
 
+        write_char(writer, BRACKET_OPEN)?;
         write_array_elements(
             writer,
-            true,
-            &BRACKET_OPEN.to_string(),
-            &BRACKET_CLOSE.to_string(),
             &input_derivation
                 .iter()
                 .map(|s| s.as_bytes().to_vec().into())
                 .collect::<Vec<BString>>(),
         )?;
+        write_char(writer, BRACKET_CLOSE)?;
 
         write_char(writer, PAREN_CLOSE)?;
     }
@@ -151,43 +144,42 @@ pub fn write_input_sources(
 ) -> Result<(), io::Error> {
     write_char(writer, COMMA)?;
 
+    write_char(writer, BRACKET_OPEN)?;
     write_array_elements(
         writer,
-        true,
-        &BRACKET_OPEN.to_string(),
-        &BRACKET_CLOSE.to_string(),
         &input_sources
             .iter()
             .map(|s| s.as_bytes().to_vec().into())
             .collect::<Vec<BString>>(),
     )?;
+    write_char(writer, BRACKET_CLOSE)?;
 
     Ok(())
 }
 
 pub fn write_system(writer: &mut impl Write, platform: &str) -> Result<(), Error> {
     write_char(writer, COMMA)?;
-    io::copy(&mut Cursor::new(escape_bstr(platform.as_bytes())), writer)?;
+    write_field(writer, platform, true)?;
     Ok(())
 }
 
 pub fn write_builder(writer: &mut impl Write, builder: &str) -> Result<(), Error> {
     write_char(writer, COMMA)?;
-    io::copy(&mut Cursor::new(escape_bstr(builder.as_bytes())), writer)?;
+    write_field(writer, builder, true)?;
     Ok(())
 }
 pub fn write_arguments(writer: &mut impl Write, arguments: &[String]) -> Result<(), io::Error> {
     write_char(writer, COMMA)?;
+
+    write_char(writer, BRACKET_OPEN)?;
     write_array_elements(
         writer,
-        true,
-        &BRACKET_OPEN.to_string(),
-        &BRACKET_CLOSE.to_string(),
         &arguments
             .iter()
             .map(|s| s.as_bytes().to_vec().into())
             .collect::<Vec<BString>>(),
     )?;
+    write_char(writer, BRACKET_CLOSE)?;
 
     Ok(())
 }
@@ -204,13 +196,11 @@ pub fn write_enviroment(
             write_char(writer, COMMA)?;
         }
 
-        write_array_elements(
-            writer,
-            false,
-            &PAREN_OPEN.to_string(),
-            &PAREN_CLOSE.to_string(),
-            &[escape_bstr(k.as_bytes()), escape_bstr(v)],
-        )?;
+        write_char(writer, PAREN_OPEN)?;
+        write_field(writer, k, false)?;
+        write_char(writer, COMMA)?;
+        write_field(writer, v, true)?;
+        write_char(writer, PAREN_CLOSE)?;
     }
 
     write_char(writer, BRACKET_CLOSE)?;