1 files changed, 55 insertions, 40 deletions

diff --git a/tools/llvm-stress/llvm-stress.cpp b/tools/llvm-stress/llvm-stress.cpp
index 727d03f9d6ea..6a1a248a0572 100644
--- a/tools/llvm-stress/llvm-stress.cpp
+++ b/tools/llvm-stress/llvm-stress.cpp
@@ -31,7 +31,8 @@
 #include <set>
 #include <sstream>
 #include <vector>
-using namespace llvm;
+
+namespace llvm {
 
 static cl::opt<unsigned> SeedCL("seed",
   cl::desc("Seed used for randomness"), cl::init(0));
@@ -42,16 +43,39 @@ static cl::opt<std::string>
 OutputFilename("o", cl::desc("Override output filename"),
                cl::value_desc("filename"));
 
-static cl::opt<bool> GenHalfFloat("generate-half-float",
-  cl::desc("Generate half-length floating-point values"), cl::init(false));
-static cl::opt<bool> GenX86FP80("generate-x86-fp80",
-  cl::desc("Generate 80-bit X86 floating-point values"), cl::init(false));
-static cl::opt<bool> GenFP128("generate-fp128",
-  cl::desc("Generate 128-bit floating-point values"), cl::init(false));
-static cl::opt<bool> GenPPCFP128("generate-ppc-fp128",
-  cl::desc("Generate 128-bit PPC floating-point values"), cl::init(false));
-static cl::opt<bool> GenX86MMX("generate-x86-mmx",
-  cl::desc("Generate X86 MMX floating-point values"), cl::init(false));
+namespace cl {
+template <> class parser<Type*> final : public basic_parser<Type*> {
+public:
+  parser(Option &O) : basic_parser(O) {}
+
+  // Parse options as IR types. Return true on error.
+  bool parse(Option &O, StringRef, StringRef Arg, Type *&Value) {
+    auto &Context = getGlobalContext();
+    if      (Arg == "half")      Value = Type::getHalfTy(Context);
+    else if (Arg == "fp128")     Value = Type::getFP128Ty(Context);
+    else if (Arg == "x86_fp80")  Value = Type::getX86_FP80Ty(Context);
+    else if (Arg == "ppc_fp128") Value = Type::getPPC_FP128Ty(Context);
+    else if (Arg == "x86_mmx")   Value = Type::getX86_MMXTy(Context);
+    else if (Arg.startswith("i")) {
+      unsigned N = 0;
+      Arg.drop_front().getAsInteger(10, N);
+      if (N > 0)
+        Value = Type::getIntNTy(Context, N);
+    }
+
+    if (!Value)
+      return O.error("Invalid IR scalar type: '" + Arg + "'!");
+    return false;
+  }
+
+  const char *getValueName() const override { return "IR scalar type"; }
+};
+}
+
+
+static cl::list<Type*> AdditionalScalarTypes("types", cl::CommaSeparated,
+  cl::desc("Additional IR scalar types "
+           "(always includes i1, i8, i16, i32, i64, float and double)"));
 
 namespace {
 /// A utility class to provide a pseudo-random number generator which is
@@ -243,35 +267,22 @@ protected:
 
   /// Pick a random scalar type.
   Type *pickScalarType() {
-    Type *t = nullptr;
-    do {
-      switch (Ran->Rand() % 30) {
-      case 0: t = Type::getInt1Ty(Context); break;
-      case 1: t = Type::getInt8Ty(Context); break;
-      case 2: t = Type::getInt16Ty(Context); break;
-      case 3: case 4:
-      case 5: t = Type::getFloatTy(Context); break;
-      case 6: case 7:
-      case 8: t = Type::getDoubleTy(Context); break;
-      case 9: case 10:
-      case 11: t = Type::getInt32Ty(Context); break;
-      case 12: case 13:
-      case 14: t = Type::getInt64Ty(Context); break;
-      case 15: case 16:
-      case 17: if (GenHalfFloat) t = Type::getHalfTy(Context); break;
-      case 18: case 19:
-      case 20: if (GenX86FP80) t = Type::getX86_FP80Ty(Context); break;
-      case 21: case 22:
-      case 23: if (GenFP128) t = Type::getFP128Ty(Context); break;
-      case 24: case 25:
-      case 26: if (GenPPCFP128) t = Type::getPPC_FP128Ty(Context); break;
-      case 27: case 28:
-      case 29: if (GenX86MMX) t = Type::getX86_MMXTy(Context); break;
-      default: llvm_unreachable("Invalid scalar value");
-      }
-    } while (t == nullptr);
+    static std::vector<Type*> ScalarTypes;
+    if (ScalarTypes.empty()) {
+      ScalarTypes.assign({
+        Type::getInt1Ty(Context),
+        Type::getInt8Ty(Context),
+        Type::getInt16Ty(Context),
+        Type::getInt32Ty(Context),
+        Type::getInt64Ty(Context),
+        Type::getFloatTy(Context),
+        Type::getDoubleTy(Context)
+      });
+      ScalarTypes.insert(ScalarTypes.end(),
+        AdditionalScalarTypes.begin(), AdditionalScalarTypes.end());
+    }
 
-    return t;
+    return ScalarTypes[Ran->Rand() % ScalarTypes.size()];
   }
 
   /// Basic block to populate
@@ -665,9 +676,13 @@ static void IntroduceControlFlow(Function *F, Random &R) {
   }
 }
 
+}
+
 int main(int argc, char **argv) {
+  using namespace llvm;
+
   // Init LLVM, call llvm_shutdown() on exit, parse args, etc.
-  llvm::PrettyStackTraceProgram X(argc, argv);
+  PrettyStackTraceProgram X(argc, argv);
   cl::ParseCommandLineOptions(argc, argv, "llvm codegen stress-tester\n");
   llvm_shutdown_obj Y;