Merge remote-tracking branch 'upstream/master'

1 files changed, 60 insertions, 0 deletions

diff --git a/klm/util/usage.cc b/klm/util/usage.cc
index e5cf76f0..16a004bb 100644
--- a/klm/util/usage.cc
+++ b/klm/util/usage.cc
@@ -1,13 +1,17 @@
 #include "util/usage.hh"
 
+#include "util/exception.hh"
+
 #include <fstream>
 #include <ostream>
+#include <sstream>
 
 #include <string.h>
 #include <ctype.h>
 #if !defined(_WIN32) && !defined(_WIN64)
 #include <sys/resource.h>
 #include <sys/time.h>
+#include <unistd.h>
 #endif
 
 namespace util {
@@ -43,4 +47,60 @@ void PrintUsage(std::ostream &out) {
 #endif
 }
 
+uint64_t GuessPhysicalMemory() {
+#if defined(_WIN32) || defined(_WIN64)
+  return 0;
+#elif defined(_SC_PHYS_PAGES) && defined(_SC_PAGESIZE)
+  long pages = sysconf(_SC_PHYS_PAGES);
+  if (pages == -1) return 0;
+  long page_size = sysconf(_SC_PAGESIZE);
+  if (page_size == -1) return 0;
+  return static_cast<uint64_t>(pages) * static_cast<uint64_t>(page_size);
+#else
+  return 0;
+#endif
+}
+
+namespace {
+class SizeParseError : public Exception {
+  public:
+    explicit SizeParseError(const std::string &str) throw() {
+      *this << "Failed to parse " << str << " into a memory size ";
+    }
+};
+
+template <class Num> uint64_t ParseNum(const std::string &arg) {
+  std::stringstream stream(arg);
+  Num value;
+  stream >> value;
+  UTIL_THROW_IF_ARG(!stream, SizeParseError, (arg), "for the leading number.");
+  std::string after;
+  stream >> after;
+  UTIL_THROW_IF_ARG(after.size() > 1, SizeParseError, (arg), "because there are more than two characters after the number.");
+  std::string throwaway;
+  UTIL_THROW_IF_ARG(stream >> throwaway, SizeParseError, (arg), "because there was more cruft " << throwaway << " after the number.");
+
+  // Silly sort, using kilobytes as your default unit.  
+  if (after.empty()) after == "K";
+  if (after == "%") {
+    uint64_t mem = GuessPhysicalMemory();
+    UTIL_THROW_IF_ARG(!mem, SizeParseError, (arg), "because % was specified but the physical memory size could not be determined.");
+    return static_cast<double>(value) * static_cast<double>(mem) / 100.0;
+  }
+  
+  std::string units("bKMGTPEZY");
+  std::string::size_type index = units.find(after[0]);
+  UTIL_THROW_IF_ARG(index == std::string::npos, SizeParseError, (arg), "the allowed suffixes are " << units << "%.");
+  for (std::string::size_type i = 0; i < index; ++i) {
+    value *= 1024;
+  }
+  return value;
+}
+
+} // namespace
+
+uint64_t ParseSize(const std::string &arg) {
+  return arg.find('.') == std::string::npos ? ParseNum<double>(arg) : ParseNum<uint64_t>(arg);
+}
+
 } // namespace util