From 2b9158ddab17135a36aa6ddd1b7123e3cdecde7d Mon Sep 17 00:00:00 2001
From: Andrew Prokhorenkov <andrew.prokhorenkov@gmail.com>
Date: Fri, 19 Jun 2020 19:52:10 -0500
Subject: [PATCH] builtin_string: add "--max" for "string pad"

---
 doc_src/cmds/string-pad.rst |  4 +--
 src/builtin_string.cpp      | 54 +++++++++++++++++++++++++++++++------
 2 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/doc_src/cmds/string-pad.rst b/doc_src/cmds/string-pad.rst
index aacd26df2..c798e9ecb 100644
--- a/doc_src/cmds/string-pad.rst
+++ b/doc_src/cmds/string-pad.rst
@@ -8,7 +8,7 @@ Synopsis
 
 ::
 
-    string pad [(-l | --left)] [(-r | --right)] [(-c | --char) CHAR] [(-n | --count) INTEGER] [(-q | --quiet)] [STRING...]
+    string pad [(-l | --left)] [(-r | --right)] [(-c | --char) CHAR] [(-n | --count) INTEGER] [(-m | --max) INTEGER] [(-q | --quiet)] [STRING...]
 
 .. END SYNOPSIS
 
@@ -17,7 +17,7 @@ Description
 
 .. BEGIN DESCRIPTION
 
-``string pad`` pads before and after the string specified character for each STRING. If ``-l`` or ``--left`` is given, only padded before string. Left only is the default padding. If ``-r`` or ``--right`` is given, only padded after string. The ``-c`` or ``--char`` switch causes the characters in CHAR to be padded instead of whitespace. The ``-n`` or ``--count`` integer specifies the amount of characters to be padded. The default padding count is 0. Exit status: 0 if string was padded, or 1 otherwise.
+``string pad`` pads before and after the string specified character for each STRING. If ``-l`` or ``--left`` is given, only padded before string. Left only is the default padding. If ``-r`` or ``--right`` is given, only padded after string. The ``-c`` or ``--char`` switch causes the characters in CHAR to be padded instead of whitespace. The ``-n`` or ``--count`` integer specifies the amount of characters to be padded. The default padding count is 0. The ``-m` or ``--max`` integer specifies the resulting length of the after after adding pad characters. Note, that the ``-n`` and ``-m`` flags are mutually exclusive, you can only use one of them. Exit status: 0 if string was padded, or 1 otherwise.
 
 .. END DESCRIPTION
 
diff --git a/src/builtin_string.cpp b/src/builtin_string.cpp
index 2c777aba3..3742bf221 100644
--- a/src/builtin_string.cpp
+++ b/src/builtin_string.cpp
@@ -1257,9 +1257,12 @@ static int string_collect(parser_t &parser, io_streams_t &streams, int argc, wch
 }
 
 static int string_pad(parser_t &parser, io_streams_t &streams, int argc, wchar_t **argv) {
+    wchar_t *cmd = argv[0];
+
     options_t opts;
     opts.char_to_pad_valid = true;
     opts.count_valid = true;
+    opts.max_valid = true;
     opts.left_valid = true;
     opts.right_valid = true;
     opts.quiet_valid = true;
@@ -1273,21 +1276,56 @@ static int string_pad(parser_t &parser, io_streams_t &streams, int argc, wchar_t
         opts.right = false;
     }
 
+    // If both count and max is specified, we raise an error.
+    if (opts.count && opts.max) {
+        streams.err.append_format(BUILTIN_ERR_COMBO2, cmd,
+                                  _(L"--count and --max are mutually exclusive"));
+        return STATUS_INVALID_ARGS;
+    }
+
     size_t npad = 0;
 
     arg_iterator_t aiter(argv, optind, streams);
     while (const wcstring *arg = aiter.nextstr()) {
-        wcstring padded_arg = wcstring(*arg, 0, arg->size());
-        if (opts.right) {
-            padded_arg.append(opts.count, opts.char_to_pad);
+        wcstring padded = wcstring(*arg);
+        size_t pad_left = 0;
+        size_t pad_right = 0;
+
+        if (opts.count && opts.right) {
+            pad_right = opts.count;
         }
-        if (opts.left) {
-            padded_arg.insert(0, opts.count, opts.char_to_pad);
+
+        if (opts.count && opts.left) {
+            pad_left = opts.count;
         }
-        // assert(begin <= end && end <= arg->size());
-        npad += arg->size() + 2 * opts.count;
+
+        if (opts.max >= padded.size()) {
+            if (opts.left && opts.right) {
+                // pad_left will get one more character
+                pad_left = (opts.max - padded.size() + 1) / 2;
+                pad_right = (opts.max - padded.size()) / 2;
+            }
+
+            if (!opts.left && opts.right) {
+                pad_right = opts.max - padded.size();
+            }
+
+            if (opts.left && !opts.right) {
+                pad_left = opts.max - padded.size();
+            }
+        }
+
+        if (pad_left) {
+            padded.insert(0, pad_left, opts.char_to_pad);
+        }
+
+        if (pad_right) {
+            padded.append(pad_right, opts.char_to_pad);
+        }
+
+        npad += padded.size();
         if (!opts.quiet) {
-            streams.out.append(padded_arg);
+            streams.out.append(padded);
             streams.out.append(L'\n');
         }
     }