Fix memory bug (#293)

* Fix memory bug

* Add SmallString

Author: wsmoses

tools/cgeist/ArgumentList.h

@@ -10,16 +10,23 @@
 #define MLIR_TOOLS_CGEIST_ARGUMENTLIST_H
 
 #include <llvm/ADT/ArrayRef.h>
+#include <llvm/ADT/SmallString.h>
 #include <llvm/Support/raw_ostream.h>
 
 namespace mlirclang {
 /// Class to pass options to a compilation tool.
 class ArgumentList {
+private:
+  /// Helper storage.
+  llvm::SmallVector<llvm::SmallString<0>> Storage;
+  /// List of arguments
+  llvm::SmallVector<const char *> Args;
+
 public:
   /// Add argument.
   ///
   /// The element stored will not be owned by this.
-  void push_back(llvm::StringRef Arg) { Args.push_back(Arg.data()); }
+  void push_back(const char *Arg) { Args.push_back(Arg); }
 
   /// Add argument and ensure it will be valid before this passer's destruction.
   ///
@@ -29,21 +36,18 @@ class ArgumentList {
     std::string Buffer;
     llvm::raw_string_ostream Stream(Buffer);
     (Stream << ... << Args);
-    Storage.push_back(Stream.str());
-    push_back(Storage.back());
+    emplace_back(llvm::StringRef(Stream.str()));
+  }
+
+  void emplace_back(llvm::StringRef &&Arg) {
+    push_back(Storage.emplace_back(Arg).c_str());
   }
 
   /// Return the underling argument list.
   ///
   /// The return value of this operation could be invalidated by subsequent
   /// calls to push_back() or emplace_back().
   llvm::ArrayRef<const char *> getArguments() const { return Args; }
-
-private:
-  /// Helper storage.
-  llvm::SmallVector<std::string> Storage;
-  /// List of arguments
-  llvm::SmallVector<const char *> Args;
 };
 } // end namespace mlirclang
 

tools/cgeist/Lib/clang-mlir.cc

@@ -5514,13 +5514,13 @@ static bool parseMLIR(const char *Argv0, std::vector<std::string> filenames,
   mlirclang::ArgumentList Argv;
   Argv.push_back(binary);
   for (const auto &filename : filenames) {
-    Argv.push_back(filename);
+    Argv.emplace_back(filename);
   }
   if (FOpenMP)
     Argv.push_back("-fopenmp");
   if (TargetTripleOpt != "") {
     Argv.push_back("-target");
-    Argv.push_back(TargetTripleOpt);
+    Argv.emplace_back(TargetTripleOpt);
   }
   if (McpuOpt != "") {
     Argv.emplace_back("-mcpu=", McpuOpt);
@@ -5530,11 +5530,11 @@ static bool parseMLIR(const char *Argv0, std::vector<std::string> filenames,
   }
   if (ResourceDir != "") {
     Argv.push_back("-resource-dir");
-    Argv.push_back(ResourceDir);
+    Argv.emplace_back(ResourceDir);
   }
   if (SysRoot != "") {
     Argv.push_back("--sysroot");
-    Argv.push_back(SysRoot);
+    Argv.emplace_back(SysRoot);
   }
   if (Verbose) {
     Argv.push_back("-v");
@@ -5556,14 +5556,14 @@ static bool parseMLIR(const char *Argv0, std::vector<std::string> filenames,
   }
   for (const auto &dir : includeDirs) {
     Argv.push_back("-I");
-    Argv.push_back(dir);
+    Argv.emplace_back(dir);
   }
   for (const auto &define : defines) {
     Argv.emplace_back("-D", define);
   }
   for (const auto &Include : Includes) {
     Argv.push_back("-include");
-    Argv.push_back(Include);
+    Argv.emplace_back(Include);
   }
 
   Argv.push_back("-emit-ast");

tools/cgeist/Test/Verification/defines.c

@@ -0,0 +1,15 @@
+// RUN: cgeist -DOUTPUT -DTEST -c %s -S | FileCheck %s
+
+int main(int argc, char **argv)
+{
+#ifdef OUTPUT
+        return 1;
+#else
+        return 2;
+#endif
+}
+
+// CHECK:   func.func @main(%arg0: i32, %arg1: memref<?xmemref<?xi8>>) -> i32 
+// CHECK-NEXT:     %c1_i32 = arith.constant 1 : i32
+// CHECK-NEXT:     return %c1_i32 : i32
+// CHECK-NEXT:   }

tools/cgeist/driver.cc

@@ -274,7 +274,7 @@ int emitBinary(char *Argv0, const char *filename,
     Argv.push_back("-fopenmp");
   if (ResourceDir != "") {
     Argv.push_back("-resource-dir");
-    Argv.push_back(ResourceDir);
+    Argv.emplace_back(ResourceDir);
   }
   if (Verbose) {
     Argv.push_back("-v");
@@ -299,7 +299,7 @@ int emitBinary(char *Argv0, const char *filename,
   }
   if (Output != "") {
     Argv.push_back("-o");
-    Argv.push_back(Output);
+    Argv.emplace_back(Output);
   }
   for (const auto *arg : LinkArgs)
     Argv.push_back(arg);