Merge pull request #391 from MartinNowack/feat_zipstream_compress

Support gzip-based compression of raw_outstreams

2 files changed, 34 insertions, 10 deletions

diff --git a/lib/Solver/QueryLoggingSolver.cpp b/lib/Solver/QueryLoggingSolver.cpp
index f93bec3c..fb957b65 100644
--- a/lib/Solver/QueryLoggingSolver.cpp
+++ b/lib/Solver/QueryLoggingSolver.cpp
@@ -7,8 +7,13 @@
 //
 //===----------------------------------------------------------------------===//
 #include "QueryLoggingSolver.h"
+#include "klee/Config/config.h"
 #include "klee/Internal/System/Time.h"
 #include "klee/Statistics.h"
+#ifdef HAVE_ZLIB_H
+#include "klee/Internal/Support/CompressionStream.h"
+#include "klee/Internal/Support/ErrorHandling.h"
+#endif
 
 #if LLVM_VERSION_CODE >= LLVM_VERSION(3, 5)
 #include "llvm/Support/FileSystem.h"
@@ -20,30 +25,49 @@ namespace {
 llvm::cl::opt<bool> DumpPartialQueryiesEarly(
     "log-partial-queries-early", llvm::cl::init(false),
     llvm::cl::desc("Log queries before calling the solver (default=off)"));
+
+#ifdef HAVE_ZLIB_H
+llvm::cl::opt<bool> CreateCompressedQueryLog(
+    "compress-query-log", llvm::cl::init(false),
+    llvm::cl::desc("Compress query log files (default=off)"));
+#endif
 }
 
 QueryLoggingSolver::QueryLoggingSolver(Solver *_solver, std::string path,
                                        const std::string &commentSign,
                                        int queryTimeToLog)
-    : solver(_solver),
+    : solver(_solver), os(0), BufferString(""), logBuffer(BufferString),
+      queryCount(0), minQueryTimeToLog(queryTimeToLog), startTime(0.0f),
+      lastQueryTime(0.0f), queryCommentSign(commentSign) {
+  if (!CreateCompressedQueryLog) {
 #if LLVM_VERSION_CODE >= LLVM_VERSION(3, 5)
-      os(path.c_str(), ErrorInfo, llvm::sys::fs::OpenFlags::F_Text),
+    os = new llvm::raw_fd_ostream(path.c_str(), ErrorInfo,
+                                  llvm::sys::fs::OpenFlags::F_Text);
 #else
-      os(path.c_str(), ErrorInfo),
+    os = new llvm::raw_fd_ostream(path.c_str(), ErrorInfo);
+#endif
+  }
+#ifdef HAVE_ZLIB_H
+  else {
+    os = new compressed_fd_ostream((path + ".gz").c_str(), ErrorInfo);
+  }
+  if (ErrorInfo != "") {
+    klee_error("Could not open file %s : %s", path.c_str(), ErrorInfo.c_str());
+  }
 #endif
-      BufferString(""), logBuffer(BufferString), queryCount(0),
-      minQueryTimeToLog(queryTimeToLog), startTime(0.0f), lastQueryTime(0.0f),
-      queryCommentSign(commentSign) {
   assert(0 != solver);
 }
 
-QueryLoggingSolver::~QueryLoggingSolver() { delete solver; }
+QueryLoggingSolver::~QueryLoggingSolver() {
+  delete solver;
+  delete os;
+}
 
 void QueryLoggingSolver::flushBufferConditionally(bool writeToFile) {
   logBuffer.flush();
   if (writeToFile) {
-    os << logBuffer.str();
-    os.flush();
+    *os << logBuffer.str();
+    os->flush();
   }
   // prepare the buffer for reuse
   BufferString = "";
diff --git a/lib/Solver/QueryLoggingSolver.h b/lib/Solver/QueryLoggingSolver.h
index bb266c67..7ac783d1 100644
--- a/lib/Solver/QueryLoggingSolver.h
+++ b/lib/Solver/QueryLoggingSolver.h
@@ -26,7 +26,7 @@ class QueryLoggingSolver : public SolverImpl {
 protected:
   Solver *solver;
   std::string ErrorInfo;
-  llvm::raw_fd_ostream os;
+  llvm::raw_ostream *os;
   // @brief Buffer used by logBuffer
   std::string BufferString;
   // @brief buffer to store logs before flushing to file