Add option to read script from command line argument

src/driver.cpp

@@ -2,21 +2,34 @@
 
 #include "driver.h"
 
+extern FILE *yyin;
+extern void *yy_scan_string(const char *yystr);
+
 namespace ebpf {
 namespace bpftrace {
 
-int Driver::parse()
+int Driver::parse_stdin()
 {
   return parser_.parse();
 }
 
-int Driver::parse(const std::string &f)
+int Driver::parse_str(const std::string &script)
+{
+  void *buffer = yy_scan_string(script.c_str());
+  int result = parser_.parse();
+  free(buffer);
+  return result;
+}
+
+int Driver::parse_file(const std::string &f)
 {
   if (!(yyin = fopen(f.c_str(), "r"))) {
-    std::cerr << "Could not open file" << std::endl;
+    std::cerr << "Error: Could not open file '" << f << "'" << std::endl;
     return -1;
   }
-  return parser_.parse();
+  int result = parser_.parse();
+  fclose(yyin);
+  return result;
 }
 
 } // namespace bpftrace

src/driver.h

@@ -6,8 +6,6 @@
 #define YY_DECL ebpf::bpftrace::Parser::symbol_type yylex(ebpf::bpftrace::Driver &driver)
 YY_DECL;
 
-extern FILE *yyin;
-
 namespace ebpf {
 namespace bpftrace {
 
@@ -15,8 +13,9 @@ class Driver {
 public:
   Driver() : parser_(*this) { }
 
-  int parse();
-  int parse(const std::string &f);
+  int parse_stdin();
+  int parse_str(const std::string &script);
+  int parse_file(const std::string &f);
 
   void error(const location &l, const std::string &m)
   {

src/main.cpp

 #include <iostream>
+#include <unistd.h>
 #include "bpftrace.h"
 #include "codegen_llvm.h"
 #include "driver.h"
@@ -7,14 +8,54 @@
 
 using namespace ebpf::bpftrace;
 
+void usage()
+{
+  std::cerr << "Usage:" << std::endl;
+  std::cerr << "  bpftrace filename" << std::endl;
+  std::cerr << "  bpftrace -e 'script'" << std::endl;
+}
+
 int main(int argc, char *argv[])
 {
   int result;
   Driver driver;
-  if (argc == 1)
-    result = driver.parse();
+
+  std::string script;
+  int c;
+  while ((c = getopt(argc, argv, "e:")) != -1)
+  {
+    switch (c)
+    {
+      case 'e':
+        script = optarg;
+        break;
+      default:
+        usage();
+        return 1;
+    }
+  }
+
+  if (script.empty())
+  {
+    // There should only be 1 non-option argument (the script file)
+    if (optind != argc-1)
+    {
+      usage();
+      return 1;
+    }
+    char *file_name = argv[optind];
+    result = driver.parse_file(file_name);
+  }
   else
-    result = driver.parse(argv[1]);
+  {
+    // Script is provided as a command line argument
+    if (optind != argc)
+    {
+      usage();
+      return 1;
+    }
+    result = driver.parse_str(script);
+  }
 
   if (result)
     return result;