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(in)), + std::istreambuf_iterator()); +} + 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; }