From 9e37409b4b40e4ca5b026ea9de7ea9784da80ff2 Mon Sep 17 00:00:00 2001
From: Jesse Beder <jbeder@gmail.com>
Date: Sat, 30 Jan 2016 18:28:27 -0600
Subject: [PATCH] Add features to read binary:

Flag -n N repeats parsing N times.
Flag -c (--cache) caches the input in a string and uses that to parse.
---
 util/read.cpp | 72 ++++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 65 insertions(+), 7 deletions(-)

diff --git a/util/read.cpp b/util/read.cpp
index 4734e88..b1f787e 100644
--- a/util/read.cpp
+++ b/util/read.cpp
@@ -26,19 +26,77 @@ class NullEventHandler : public YAML::EventHandler {
   virtual void OnMapEnd() {}
 };
 
-void run(YAML::Parser& parser) {
+void run(std::istream& in) {
+  YAML::Parser parser(in);
   NullEventHandler handler;
   parser.HandleNextDocument(handler);
 }
 
+void usage() { std::cerr << "Usage: read [-n N] [-c, --cache] [filename]\n"; }
+
+std::string read_stream(std::istream& in) {
+  return std::string((std::istreambuf_iterator<char>(in)),
+                     std::istreambuf_iterator<char>());
+}
+
 int main(int argc, char** argv) {
-  if (argc > 1) {
-    std::ifstream in(argv[1]);
-    YAML::Parser parser(in);
-    run(parser);
+  int N = 1;
+  bool cache = false;
+  std::string filename;
+  for (int i = 1; i < argc; i++) {
+    std::string arg = argv[i];
+    if (arg == "-n") {
+      i++;
+      if (i >= argc) {
+        usage();
+        return -1;
+      }
+      N = std::atoi(argv[i]);
+      if (N <= 0) {
+        usage();
+        return -1;
+      }
+    } else if (arg == "-c" || arg == "--cache") {
+      cache = true;
+    } else {
+      filename = argv[i];
+      if (i + 1 != argc) {
+        usage();
+        return -1;
+      }
+    }
+  }
+
+  if (N > 1 && !cache && filename == "") {
+    usage();
+    return -1;
+  }
+
+  if (cache) {
+    std::string input;
+    if (filename != "") {
+      std::ifstream in(filename);
+      input = read_stream(in);
+    } else {
+      input = read_stream(std::cin);
+    }
+    std::istringstream in(input);
+    for (int i = 0; i < N; i++) {
+      in.seekg(std::ios_base::beg);
+      run(in);
+    }
   } else {
-    YAML::Parser parser(std::cin);
-    run(parser);
+    if (filename != "") {
+      std::ifstream in(filename);
+      for (int i = 0; i < N; i++) {
+        in.seekg(std::ios_base::beg);
+        run(in);
+      }
+    } else {
+      for (int i = 0; i < N; i++) {
+        run(std::cin);
+      }
+    }
   }
   return 0;
 }