support-stdin-input-llvm-ir2vec

Author: svkeerthy

llvm/tools/llvm-ir2vec/llvm-ir2vec.cpp

@@ -38,6 +38,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
 
@@ -48,10 +49,10 @@ namespace ir2vec {
 
 static cl::OptionCategory IR2VecToolCategory("IR2Vec Tool Options");
 
-static cl::opt<std::string> InputFilename(cl::Positional,
-                                          cl::desc("<input bitcode file>"),
-                                          cl::Required,
-                                          cl::cat(IR2VecToolCategory));
+static cl::opt<std::string>
+    InputFilename(cl::Positional,
+                  cl::desc("<input bitcode file or '-' for stdin>"),
+                  cl::init("-"), cl::cat(IR2VecToolCategory));
 
 static cl::opt<std::string> OutputFilename("o", cl::desc("Output filename"),
                                            cl::value_desc("filename"),
@@ -283,10 +284,24 @@ int main(int argc, char **argv) {
   if (Mode == TripletMode && Level.getNumOccurrences() > 0)
     errs() << "Warning: --level option is ignored in triplet mode\n";
 
-  // Parse the input LLVM IR file
+  // Parse the input LLVM IR file or stdin
   SMDiagnostic Err;
   LLVMContext Context;
-  std::unique_ptr<Module> M = parseIRFile(InputFilename, Err, Context);
+  std::unique_ptr<Module> M;
+
+  if (InputFilename == "-") {
+    // Read from stdin
+    auto StdinBuffer = MemoryBuffer::getSTDIN();
+    if (std::error_code EC = StdinBuffer.getError()) {
+      errs() << "Error reading from stdin: " << EC.message() << "\n";
+      return 1;
+    }
+    M = parseIR(StdinBuffer.get()->getMemBufferRef(), Err, Context);
+  } else {
+    // Read from file
+    M = parseIRFile(InputFilename, Err, Context);
+  }
+
   if (!M) {
     Err.print(argv[0], errs());
     return 1;