Support binary unit prefixes in command line arguments

4 files changed, 29 insertions, 27 deletions

diff --git a/src/libmain/shared.cc b/src/libmain/shared.cc
index f1feeddd6..e797c2fb9 100644
--- a/src/libmain/shared.cc
+++ b/src/libmain/shared.cc
@@ -219,10 +219,8 @@ LegacyArgs::LegacyArgs(const std::string & programName,
             .description = description,
             .labels = {"n"},
             .handler = {[=](std::string s) {
-                if (auto n = string2Int<unsigned int>(s))
-                    settings.set(dest, std::to_string(*n));
-                else
-                    throw UsageError("'%s' is not an integer", s);
+                auto n = string2IntWithUnitPrefix<uint64_t>(s);
+                settings.set(dest, std::to_string(n));
             }}
         });
     };
diff --git a/src/libmain/shared.hh b/src/libmain/shared.hh
index 38f627b44..edc7b5efa 100644
--- a/src/libmain/shared.hh
+++ b/src/libmain/shared.hh
@@ -57,22 +57,7 @@ template<class N> N getIntArg(const string & opt,
 {
     ++i;
     if (i == end) throw UsageError("'%1%' requires an argument", opt);
-    string s = *i;
-    N multiplier = 1;
-    if (allowUnit && !s.empty()) {
-        char u = std::toupper(*s.rbegin());
-        if (std::isalpha(u)) {
-            if (u == 'K') multiplier = 1ULL << 10;
-            else if (u == 'M') multiplier = 1ULL << 20;
-            else if (u == 'G') multiplier = 1ULL << 30;
-            else if (u == 'T') multiplier = 1ULL << 40;
-            else throw UsageError("invalid unit specifier '%1%'", u);
-            s.resize(s.size() - 1);
-        }
-    }
-    if (auto n = string2Int<N>(s))
-        return *n * multiplier;
-    throw UsageError("'%1%' requires an integer argument", opt);
+    return string2IntWithUnitPrefix<N>(*i);
 }
 
 
diff --git a/src/libutil/args.hh b/src/libutil/args.hh
index 823d843aa..3783bc84f 100644
--- a/src/libutil/args.hh
+++ b/src/libutil/args.hh
@@ -87,10 +87,7 @@ protected:
         template<class I>
         Handler(I * dest)
             : fun([=](std::vector<std::string> ss) {
-                if (auto n = string2Int<I>(ss[0]))
-                    *dest = *n;
-                else
-                    throw UsageError("'%s' is not an integer", ss[0]);
+                *dest = string2IntWithUnitPrefix<I>(ss[0]);
               })
             , arity(1)
         { }
diff --git a/src/libutil/util.hh b/src/libutil/util.hh
index 7a4d5fe92..ab0bd865a 100644
--- a/src/libutil/util.hh
+++ b/src/libutil/util.hh
@@ -401,12 +401,34 @@ template<class N>
 std::optional<N> string2Int(const std::string & s)
 {
     if (s.substr(0, 1) == "-" && !std::numeric_limits<N>::is_signed)
-        return {};
+        return std::nullopt;
     std::istringstream str(s);
     N n;
     str >> n;
     if (str && str.get() == EOF) return n;
-    return {};
+    return std::nullopt;
+}
+
+/* Like string2Int(), but support an optional suffix 'K', 'M', 'G' or
+   'T' denoting a binary unit prefix. */
+template<class N>
+N string2IntWithUnitPrefix(std::string s)
+{
+    N multiplier = 1;
+    if (!s.empty()) {
+        char u = std::toupper(*s.rbegin());
+        if (std::isalpha(u)) {
+            if (u == 'K') multiplier = 1ULL << 10;
+            else if (u == 'M') multiplier = 1ULL << 20;
+            else if (u == 'G') multiplier = 1ULL << 30;
+            else if (u == 'T') multiplier = 1ULL << 40;
+            else throw UsageError("invalid unit specifier '%1%'", u);
+            s.resize(s.size() - 1);
+        }
+    }
+    if (auto n = string2Int<N>(s))
+        return *n * multiplier;
+    throw UsageError("'%s' is not an integer", s);
 }
 
 /* Parse a string into a float. */
@@ -417,7 +439,7 @@ std::optional<N> string2Float(const string & s)
     N n;
     str >> n;
     if (str && str.get() == EOF) return n;
-    return {};
+    return std::nullopt;
 }