Avoid internalization of non-standard entry point (i.e. not the main function) (#455)

5 files changed, 19 insertions, 10 deletions

diff --git a/include/klee/Interpreter.h b/include/klee/Interpreter.h
index b40ad0d5..3a4d40b4 100644
--- a/include/klee/Interpreter.h
+++ b/include/klee/Interpreter.h
@@ -51,15 +51,16 @@ public:
   /// registering a module with the interpreter.
   struct ModuleOptions {
     std::string LibraryDir;
+    std::string EntryPoint;
     bool Optimize;
     bool CheckDivZero;
     bool CheckOvershift;
 
-    ModuleOptions(const std::string& _LibraryDir, 
-                  bool _Optimize, bool _CheckDivZero,
-                  bool _CheckOvershift)
-      : LibraryDir(_LibraryDir), Optimize(_Optimize), 
-        CheckDivZero(_CheckDivZero), CheckOvershift(_CheckOvershift) {}
+    ModuleOptions(const std::string &_LibraryDir,
+                  const std::string &_EntryPoint, bool _Optimize,
+                  bool _CheckDivZero, bool _CheckOvershift)
+        : LibraryDir(_LibraryDir), EntryPoint(_EntryPoint), Optimize(_Optimize),
+          CheckDivZero(_CheckDivZero), CheckOvershift(_CheckOvershift) {}
   };
 
   enum LogType
diff --git a/lib/Module/KModule.cpp b/lib/Module/KModule.cpp
index 01165e94..57346a31 100644
--- a/lib/Module/KModule.cpp
+++ b/lib/Module/KModule.cpp
@@ -132,7 +132,7 @@ KModule::~KModule() {
 /***/
 
 namespace llvm {
-extern void Optimize(Module*);
+extern void Optimize(Module *, const std::string &EntryPoint);
 }
 
 // what a hack
@@ -308,7 +308,7 @@ void KModule::prepare(const Interpreter::ModuleOptions &opts,
   pm.run(*module);
 
   if (opts.Optimize)
-    Optimize(module);
+    Optimize(module, opts.EntryPoint);
 #if LLVM_VERSION_CODE < LLVM_VERSION(3, 3)
   // Force importing functions required by intrinsic lowering. Kind of
   // unfortunate clutter when we don't need them but we won't know
diff --git a/lib/Module/Optimize.cpp b/lib/Module/Optimize.cpp
index ce43cd96..3d9c8cc1 100644
--- a/lib/Module/Optimize.cpp
+++ b/lib/Module/Optimize.cpp
@@ -163,7 +163,7 @@ static void AddStandardCompilePasses(PassManager &PM) {
 /// Optimize - Perform link time optimizations. This will run the scalar
 /// optimizations, any loaded plugin-optimization modules, and then the
 /// inter-procedural optimizations if applicable.
-void Optimize(Module* M) {
+void Optimize(Module *M, const std::string &EntryPoint) {
 
   // Instantiate the pass manager to organize the passes.
   PassManager Passes;
@@ -192,7 +192,8 @@ void Optimize(Module* M) {
     // internal.
     if (!DisableInternalize) {
 #if LLVM_VERSION_CODE >= LLVM_VERSION(3, 2)
-      ModulePass *pass = createInternalizePass(std::vector<const char *>(1, "main"));
+      ModulePass *pass = createInternalizePass(
+          std::vector<const char *>(1, EntryPoint.c_str()));
 #else
       ModulePass *pass = createInternalizePass(true);
 #endif
diff --git a/test/regression/2016-08-11-entry-point-internalize-pass.c b/test/regression/2016-08-11-entry-point-internalize-pass.c
new file mode 100644
index 00000000..4cd8ff8d
--- /dev/null
+++ b/test/regression/2016-08-11-entry-point-internalize-pass.c
@@ -0,0 +1,7 @@
+// RUN: %llvmgcc %s -emit-llvm -g -O0 -c -o %t.bc
+// RUN: rm -rf %t.klee-out
+// RUN: %klee --output-dir=%t.klee-out --entry-point=entry %t.bc
+
+int entry() {
+  return 0;
+}
diff --git a/tools/klee/main.cpp b/tools/klee/main.cpp
index 2143c95b..1eab170f 100644
--- a/tools/klee/main.cpp
+++ b/tools/klee/main.cpp
@@ -1286,7 +1286,7 @@ int main(int argc, char **argv, char **envp) {
   }
 
   std::string LibraryDir = KleeHandler::getRunTimeLibraryPath(argv[0]);
-  Interpreter::ModuleOptions Opts(LibraryDir.c_str(),
+  Interpreter::ModuleOptions Opts(LibraryDir.c_str(), EntryPoint,
                                   /*Optimize=*/OptimizeModule,
                                   /*CheckDivZero=*/CheckDivZero,
                                   /*CheckOvershift=*/CheckOvershift);