Make yes command behaviour match gnu

Author: morr0ne

src/bin/yes.rs

@@ -1,6 +1,5 @@
 use std::{
     borrow::Cow,
-    ffi::OsString,
     io::{BufWriter, Write, stdout},
     os::unix::ffi::OsStringExt,
 };
@@ -20,9 +19,9 @@ fn main() -> Result {
     // No point in locking stdout since we only use it once in this program
     let mut stdout = stdout();
 
-    let mut first_value: Option<Vec<u8>> = None;
+    let mut buffer: Option<Vec<u8>> = None;
 
-    if let Some(arg) = arg_parser.forward()? {
+    while let Some(arg) = arg_parser.forward()? {
         match arg {
             Long("version") => {
                 stdout.write_all(VERSION)?;
@@ -36,29 +35,22 @@ fn main() -> Result {
 
                 return Ok(());
             }
+            Value(value) => {
+                extend_buffer(&mut buffer, value.as_bytes().to_vec());
+            }
             Long(_) | Short(_) => return Err(Exit::ArgError(arg.into_error(None))),
-            Value(value) => first_value = Some(value.as_bytes().to_vec()),
         }
     }
 
-    let mut args = arg_parser.into_inner();
-
-    if first_value.is_none() {
-        first_value = args.next().map(OsString::into_vec);
-    }
-
-    // We prepare the output so it doesn't need to go through the formatting each time
-    let output: Cow<'_, [u8]> = if let Some(mut first) = first_value {
-        for arg in args {
-            first.push(b' '); // Manually put the space
-            first.append(&mut arg.into_vec()); // Append will move the data efficiently from the other vector
-        }
-
-        first.push(b'\n');
+    arg_parser
+        .into_inner()
+        .for_each(|arg| extend_buffer(&mut buffer, arg.into_vec()));
 
-        Cow::Owned(first)
+    let output: Cow<'_, [u8]> = if let Some(mut buffer) = buffer {
+        buffer.push(b'\n');
+        Cow::Owned(buffer)
     } else {
-        Cow::Borrowed(b"y\n") // If there are no args we can just hardcode it and avoid allocation
+        Cow::Borrowed(b"y\n")
     };
 
     // Write everything to stdout, BufWriter will handle the buffering
@@ -68,3 +60,13 @@ fn main() -> Result {
         stdout.write_all(&output)?;
     }
 }
+
+#[inline]
+fn extend_buffer(buffer: &mut Option<Vec<u8>>, mut arg: Vec<u8>) {
+    if let Some(buffer) = buffer {
+        buffer.push(b' '); // Manually put the space
+        buffer.append(&mut arg);
+    } else {
+        *buffer = Some(arg)
+    }
+}