From 4e6e897781a47e0c2cc180f670cdb816f2024275 Mon Sep 17 00:00:00 2001
From: Himadri Bhattacharjee <lavafroth@protonmail.com>
Date: Sun, 11 Feb 2024 11:19:02 +0000
Subject: [PATCH] string repeat: allow omission of `-n` (#10282)

---
 src/builtins/string/repeat.rs | 60 +++++++++++++++++++++++++++--------
 tests/checks/string.fish      | 15 +++++++--
 2 files changed, 59 insertions(+), 16 deletions(-)

diff --git a/src/builtins/string/repeat.rs b/src/builtins/string/repeat.rs
index eab03eee4..735834601 100644
--- a/src/builtins/string/repeat.rs
+++ b/src/builtins/string/repeat.rs
@@ -3,8 +3,8 @@ use crate::wutil::fish_wcstol;
 
 #[derive(Default)]
 pub struct Repeat {
-    count: usize,
-    max: usize,
+    count: Option<usize>,
+    max: Option<usize>,
     quiet: bool,
     no_newline: bool,
 }
@@ -21,14 +21,17 @@ impl StringSubCommand<'_> for Repeat {
     fn parse_opt(&mut self, name: &wstr, c: char, arg: Option<&wstr>) -> Result<(), StringError> {
         match c {
             'n' => {
-                self.count = fish_wcstol(arg.unwrap())?
-                    .try_into()
-                    .map_err(|_| invalid_args!("%ls: Invalid count value '%ls'\n", name, arg))?
+                self.count =
+                    Some(fish_wcstol(arg.unwrap())?.try_into().map_err(|_| {
+                        invalid_args!("%ls: Invalid count value '%ls'\n", name, arg)
+                    })?)
             }
             'm' => {
-                self.max = fish_wcstol(arg.unwrap())?
-                    .try_into()
-                    .map_err(|_| invalid_args!("%ls: Invalid max value '%ls'\n", name, arg))?
+                self.max = Some(
+                    fish_wcstol(arg.unwrap())?
+                        .try_into()
+                        .map_err(|_| invalid_args!("%ls: Invalid max value '%ls'\n", name, arg))?,
+                )
             }
             'q' => self.quiet = true,
             'N' => self.no_newline = true,
@@ -37,6 +40,33 @@ impl StringSubCommand<'_> for Repeat {
         return Ok(());
     }
 
+    fn take_args(
+        &mut self,
+        optind: &mut usize,
+        args: &[&'_ wstr],
+        streams: &mut IoStreams,
+    ) -> Option<c_int> {
+        if self.count.is_some() || self.max.is_some() {
+            return STATUS_CMD_OK;
+        }
+
+        let name = args[0];
+
+        let Some(arg) = args.get(*optind) else {
+            string_error!(streams, BUILTIN_ERR_ARG_COUNT0, name);
+            return STATUS_INVALID_ARGS;
+        };
+        *optind += 1;
+
+        let Ok(Ok(count)) = fish_wcstol(arg).map(|count| count.try_into()) else {
+            string_error!(streams, "%ls: Invalid count value '%ls'\n", name, arg);
+            return STATUS_INVALID_ARGS;
+        };
+
+        self.count = Some(count);
+        return STATUS_CMD_OK;
+    }
+
     fn handle(
         &mut self,
         _parser: &Parser,
@@ -44,7 +74,10 @@ impl StringSubCommand<'_> for Repeat {
         optind: &mut usize,
         args: &[&wstr],
     ) -> Option<libc::c_int> {
-        if self.max == 0 && self.count == 0 {
+        let max = self.max.unwrap_or_default();
+        let count = self.count.unwrap_or_default();
+
+        if max == 0 && count == 0 {
             // XXX: This used to be allowed, but returned 1.
             // Keep it that way for now instead of adding an error.
             // streams.err.append(L"Count or max must be greater than zero");
@@ -75,13 +108,12 @@ impl StringSubCommand<'_> for Repeat {
 
             // The maximum size of the string is either the "max" characters,
             // or it's the "count" repetitions, whichever ends up lower.
-            let max = if self.max == 0
-                || (self.count > 0 && w.len().wrapping_mul(self.count) < self.max)
-            {
+            let max_repeat_len = w.len().wrapping_mul(count);
+            let max = if max == 0 || (count > 0 && max_repeat_len < max) {
                 // TODO: we should disallow overflowing unless max <= w.len().checked_mul(self.count).unwrap_or(usize::MAX)
-                w.len().wrapping_mul(self.count)
+                max_repeat_len
             } else {
-                self.max
+                max
             };
 
             // Reserve a string to avoid writing constantly.
diff --git a/tests/checks/string.fish b/tests/checks/string.fish
index 871052247..b24e5c706 100644
--- a/tests/checks/string.fish
+++ b/tests/checks/string.fish
@@ -468,9 +468,21 @@ string repeat -n 2 foo
 string repeat --count 2 foo
 # CHECK: foofoo
 
+string repeat 2 foo
+# CHECK: foofoo
+
+string repeat 2 -n 3
+# CHECK: 222
+
 echo foo | string repeat -n 2
 # CHECK: foofoo
 
+echo foo | string repeat 2
+# CHECK: foofoo
+
+string repeat foo
+# CHECKERR: string repeat: Invalid count value 'foo'
+
 string repeat -n2 -q foo; and echo "exit 0"
 # CHECK: exit 0
 
@@ -566,8 +578,7 @@ string repeat -n; and echo "exit 0"
 # DONTCHECKERR: string repeat: Unknown option '-l'
 
 string repeat ""
-or echo string repeat empty string failed
-# CHECK: string repeat empty string failed
+# CHECKERR: string repeat: Invalid count value ''
 
 string repeat -n3 ""
 or echo string repeat empty string failed