feat(users/Profpatsch/netencode.rs): parse multiple stdin values r/4966

Adds support for parsing multiple netencode values from stdin.

This is overly complicated for my tastes, but I don’t see a better way
of writing this logic that does not read all of stdin before starting
to parse the first value.

A kingdom for a conduit.

Change-Id: Ia4f849d4096c43e887756b756d2a85d7f9cd380a
Reviewed-on: https://cl.tvl.fyi/c/depot/+/6631
Autosubmit: Profpatsch <mail@profpatsch.de>
Reviewed-by: Profpatsch <mail@profpatsch.de>
Tested-by: BuildkiteCI

4 files changed, 104 insertions, 31 deletions

diff --git a/users/Profpatsch/netencode/Netencode.hs b/users/Profpatsch/netencode/Netencode.hs
index 8398d8246fd8..dfc57ce8dc27 100644
--- a/users/Profpatsch/netencode/Netencode.hs
+++ b/users/Profpatsch/netencode/Netencode.hs
@@ -133,11 +133,11 @@ record = T . Fix . Record . coerce @(NEMap Text T) @(NEMap Text (Fix TF))
 list :: [T] -> T
 list = T . Fix . List . coerce @[T] @([Fix TF])
 
--- Stable encoding of a netencode value. Record keys will be sorted lexicographically ascending.
+-- | Stable encoding of a netencode value. Record keys will be sorted lexicographically ascending.
 netencodeEncodeStable :: T -> Builder
 netencodeEncodeStable (T fix) = Fix.foldFix (netencodeEncodeStableF id) fix
 
--- Stable encoding of a netencode functor value. Record keys will be sorted lexicographically ascending.
+-- | Stable encoding of a netencode functor value. Record keys will be sorted lexicographically ascending.
 --
 -- The given function is used for encoding the recursive values.
 netencodeEncodeStableF :: (rec -> Builder) -> TF rec -> Builder
diff --git a/users/Profpatsch/netencode/default.nix b/users/Profpatsch/netencode/default.nix
index 964cb4b5d639..00fadf695357 100644
--- a/users/Profpatsch/netencode/default.nix
+++ b/users/Profpatsch/netencode/default.nix
@@ -62,9 +62,8 @@ let
 
     fn main() {
       let (_, prog) = exec_helpers::args_for_exec("netencode-pretty", 0);
-      let mut buf = vec![];
-      let u = netencode::u_from_stdin_or_die_user_error("netencode-pretty", &mut buf);
-      match netencode_pretty::Pretty::from_u(u).print_multiline(&mut std::io::stdout()) {
+      let t = netencode::t_from_stdin_or_die_user_error("netencode-pretty");
+      match netencode_pretty::Pretty::from_u(t.to_u()).print_multiline(&mut std::io::stdout()) {
         Ok(()) => {},
         Err(err) => exec_helpers::die_temporary("netencode-pretty", format!("could not write to stdout: {}", err))
       }
@@ -89,24 +88,21 @@ let
       dependencies = [
         netencode-rs
         depot.users.Profpatsch.execline.exec-helpers
-        depot.users.Profpatsch.arglib.netencode.rust
       ];
     } ''
     extern crate netencode;
-    extern crate arglib_netencode;
     extern crate exec_helpers;
     use netencode::{encode, dec};
     use netencode::dec::{Decoder, DecodeError};
 
     fn main() {
-        let mut buf = vec![];
         let args = exec_helpers::args("record-get", 1);
         let field = match std::str::from_utf8(&args[0]) {
             Ok(f) => f,
             Err(_e) => exec_helpers::die_user_error("record-get", format!("The field name needs to be valid unicode"))
         };
-        let u = netencode::u_from_stdin_or_die_user_error("record-get", &mut buf);
-        match (dec::RecordDot {field, inner: dec::AnyU }).dec(u) {
+        let t = netencode::t_from_stdin_or_die_user_error("record-get");
+        match (dec::RecordDot {field, inner: dec::AnyU }).dec(t.to_u()) {
             Ok(u) => encode(&mut std::io::stdout(), &u).expect("encoding to stdout failed"),
             Err(DecodeError(err)) => exec_helpers::die_user_error("record-get", err)
         }
@@ -126,10 +122,9 @@ let
     use netencode::dec::{Record, Try, ScalarAsBytes, Decoder, DecodeError};
 
     fn main() {
-        let mut buf = vec![];
-        let u = netencode::u_from_stdin_or_die_user_error("record-splice-env", &mut buf);
+        let t = netencode::t_from_stdin_or_die_user_error("record-splice-env");
         let (_, prog) = exec_helpers::args_for_exec("record-splice-env", 0);
-        match Record(Try(ScalarAsBytes)).dec(u) {
+        match Record(Try(ScalarAsBytes)).dec(t.to_u()) {
             Ok(map) => {
                 exec_helpers::exec_into_args(
                     "record-splice-env",
diff --git a/users/Profpatsch/netencode/netencode.rs b/users/Profpatsch/netencode/netencode.rs
index bb08dca4aa57..34a8fcef0990 100644
--- a/users/Profpatsch/netencode/netencode.rs
+++ b/users/Profpatsch/netencode/netencode.rs
@@ -198,25 +198,103 @@ pub fn text(s: String) -> T {
     T::Text(s)
 }
 
-pub fn u_from_stdin_or_die_user_error<'a>(prog_name: &'_ str, stdin_buf: &'a mut Vec<u8>) -> U<'a> {
-    std::io::stdin().lock().read_to_end(stdin_buf);
-    let u = match parse::u_u(stdin_buf) {
-        Ok((rest, u)) => match rest {
-            b"" => u,
-            _ => exec_helpers::die_user_error(
+pub fn t_from_stdin_or_die_user_error<'a>(prog_name: &'_ str) -> T {
+    match t_from_stdin_or_die_user_error_with_rest(prog_name, &vec![]) {
+        None => exec_helpers::die_user_error(prog_name, "stdin was empty"),
+        Some((rest, t)) => {
+            if rest.is_empty() {
+                t
+            } else {
+                exec_helpers::die_user_error(
+                    prog_name,
+                    format!(
+                        "stdin contained some soup after netencode value: {:?}",
+                        String::from_utf8_lossy(&rest)
+                    ),
+                )
+            }
+        }
+    }
+}
+
+/// Read a netencode value from stdin incrementally, return bytes that could not be read.
+/// Nothing if there was nothing to read from stdin & no initial_bytes were provided.
+/// These can be passed back as `initial_bytes` if more values should be read.
+pub fn t_from_stdin_or_die_user_error_with_rest<'a>(
+    prog_name: &'_ str,
+    initial_bytes: &[u8],
+) -> Option<(Vec<u8>, T)> {
+    let mut chonker = Chunkyboi::new(std::io::stdin().lock(), 4096);
+    // The vec to pass to the parser on each step
+    let mut parser_vec: Vec<u8> = initial_bytes.to_vec();
+    // whether stdin was already empty
+    let mut was_empty: bool = false;
+    loop {
+        match chonker.next() {
+            None => {
+                if parser_vec.is_empty() {
+                    return None;
+                } else {
+                    was_empty = true
+                }
+            }
+            Some(Err(err)) => exec_helpers::die_temporary(
                 prog_name,
-                format!(
-                    "stdin contained some soup after netencode value: {:?}",
-                    String::from_utf8_lossy(rest)
-                ),
+                &format!("could not read from stdin: {:?}", err),
+            ),
+            Some(Ok(mut new_bytes)) => parser_vec.append(&mut new_bytes),
+        }
+
+        match parse::t_t(&parser_vec) {
+            Ok((rest, t)) => return Some((rest.to_owned(), t)),
+            Err(nom::Err::Incomplete(Needed)) => {
+                if was_empty {
+                    exec_helpers::die_user_error(
+                        prog_name,
+                        &format!(
+                            "unable to parse netencode from stdin, input incomplete: {:?}",
+                            parser_vec
+                        ),
+                    );
+                }
+                // read more from stdin and try parsing again
+                continue;
+            }
+            Err(err) => exec_helpers::die_user_error(
+                prog_name,
+                &format!("unable to parse netencode from stdin: {:?}", err),
             ),
-        },
-        Err(err) => exec_helpers::die_user_error(
-            prog_name,
-            format!("unable to parse netencode from stdin: {:?}", err),
-        ),
-    };
-    u
+        }
+    }
+}
+
+// iter helper
+// TODO: put into its own module
+struct Chunkyboi<T> {
+    inner: T,
+    buf: Vec<u8>,
+}
+
+impl<R: Read> Chunkyboi<R> {
+    fn new(inner: R, chunksize: usize) -> Self {
+        let buf = vec![0; chunksize];
+        Chunkyboi { inner, buf }
+    }
+}
+
+impl<R: Read> Iterator for Chunkyboi<R> {
+    type Item = std::io::Result<Vec<u8>>;
+
+    fn next(&mut self) -> Option<std::io::Result<Vec<u8>>> {
+        match self.inner.read(&mut self.buf) {
+            Ok(0) => None,
+            Ok(read) => {
+                // clone a new buffer so we can reuse the internal one
+                Some(Ok(self.buf[..read].to_owned()))
+            }
+            Err(err) => Some(Err(err)),
+        }
+    }
 }
 
 pub mod parse {
diff --git a/users/Profpatsch/read-http.rs b/users/Profpatsch/read-http.rs
index efaded87e6cd..2b24e6beb10a 100644
--- a/users/Profpatsch/read-http.rs
+++ b/users/Profpatsch/read-http.rs
@@ -220,7 +220,7 @@ fn write_dict<'buf, 'a>(
 }
 
 // iter helper
-
+// TODO: put into its own module
 struct Chunkyboi<T> {
     inner: T,
     buf: Vec<u8>,