Merge pull request #391 from MartinNowack/feat_zipstream_compress

Support gzip-based compression of raw_outstreams

4 files changed, 176 insertions, 10 deletions

diff --git a/lib/Core/Executor.cpp b/lib/Core/Executor.cpp
index f23fb3c0..46e163ea 100644
--- a/lib/Core/Executor.cpp
+++ b/lib/Core/Executor.cpp
@@ -89,6 +89,10 @@
 #include "llvm/IR/CallSite.h"
 #endif
 
+#ifdef HAVE_ZLIB_H
+#include "klee/Internal/Support/CompressionStream.h"
+#endif
+
 #include <cassert>
 #include <algorithm>
 #include <iomanip>
@@ -158,6 +162,11 @@ namespace {
                      "[inst_id]"),
           clEnumValEnd),
       llvm::cl::CommaSeparated);
+#ifdef HAVE_ZLIB_H
+  cl::opt<bool> DebugCompressInstructions(
+      "debug-compress-instructions", cl::init(false),
+      cl::desc("Compress the logged instructions in gzip format."));
+#endif
 
   cl::opt<bool>
   DebugCheckForImpliedValues("debug-check-for-implied-values");
@@ -332,6 +341,10 @@ Executor::Executor(const InterpreterOptions &opts, InterpreterHandler *ih)
     std::string debug_file_name =
         interpreterHandler->getOutputFilename("instructions.txt");
     std::string ErrorInfo;
+#ifdef HAVE_ZLIB_H
+    if (!DebugCompressInstructions) {
+#endif
+
 #if LLVM_VERSION_CODE >= LLVM_VERSION(3, 5)
     debugInstFile = new llvm::raw_fd_ostream(debug_file_name.c_str(), ErrorInfo,
                                              llvm::sys::fs::OpenFlags::F_Text),
@@ -339,6 +352,16 @@ Executor::Executor(const InterpreterOptions &opts, InterpreterHandler *ih)
     debugInstFile =
         new llvm::raw_fd_ostream(debug_file_name.c_str(), ErrorInfo);
 #endif
+#ifdef HAVE_ZLIB_H
+    } else {
+      debugInstFile = new compressed_fd_ostream(
+          (debug_file_name + ".gz").c_str(), ErrorInfo);
+    }
+#endif
+    if (ErrorInfo != "") {
+      klee_error("Could not open file %s : %s", debug_file_name.c_str(),
+                 ErrorInfo.c_str());
+    }
   }
 }
 
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
diff --git a/lib/Support/CompressionStream.cpp b/lib/Support/CompressionStream.cpp
new file mode 100644
index 00000000..eb208edf
--- /dev/null
+++ b/lib/Support/CompressionStream.cpp
@@ -0,0 +1,119 @@
+//===-- CompressionStream.cpp --------------------------------------------===//
+//
+//                     The KLEE Symbolic Virtual Machine
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+#include "klee/Config/config.h"
+#include "klee/Config/Version.h"
+#ifdef HAVE_ZLIB_H
+#include "klee/Internal/Support/CompressionStream.h"
+#if LLVM_VERSION_CODE >= LLVM_VERSION(3, 3)
+#include "llvm/Support/system_error.h"
+#else
+#include <fcntl.h>
+#include <errno.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+#endif
+
+namespace klee {
+
+compressed_fd_ostream::compressed_fd_ostream(const char *Filename,
+                                             std::string &ErrorInfo)
+    : llvm::raw_ostream(), pos(0) {
+  ErrorInfo = "";
+#if LLVM_VERSION_CODE >= LLVM_VERSION(3, 3)
+  // Open file in binary mode
+  llvm::error_code EC =
+      llvm::sys::fs::openFileForWrite(Filename, FD, llvm::sys::fs::F_Binary);
+
+  if (EC) {
+    ErrorInfo = EC.message();
+    FD = -1;
+  }
+#else
+  FD = ::open(Filename, O_WRONLY | O_CREAT | O_TRUNC, 0666);
+  if (FD < 0) {
+    ErrorInfo = "Could not open file.";
+    FD = -1;
+  }
+#endif
+  // Initialize the compression library
+  strm.zalloc = 0;
+  strm.zfree = 0;
+  strm.next_out = buffer;
+  strm.avail_out = BUFSIZE;
+
+  deflateInit2(&strm, Z_BEST_COMPRESSION, Z_DEFLATED, 15 | 16,
+               8 /* memory usage default, 0 smalles, 9 highest*/,
+               Z_DEFAULT_STRATEGY);
+}
+
+void compressed_fd_ostream::writeFullCompressedData() {
+  // Check if no space available and write the buffer
+  if (strm.avail_out == 0) {
+    write_file(reinterpret_cast<const char *>(buffer), BUFSIZE);
+    strm.next_out = buffer;
+    strm.avail_out = BUFSIZE;
+  }
+}
+
+void compressed_fd_ostream::flush_compressed_data() {
+  // flush data from the raw buffer
+  flush();
+
+  // write the remaining data
+  int deflate_res = Z_OK;
+  while (deflate_res == Z_OK) {
+    // Check if no space available and write the buffer
+    writeFullCompressedData();
+    deflate_res = deflate(&strm, Z_FINISH);
+  }
+  assert(deflate_res == Z_STREAM_END);
+  write_file(reinterpret_cast<const char *>(buffer), BUFSIZE - strm.avail_out);
+}
+
+compressed_fd_ostream::~compressed_fd_ostream() {
+  if (FD >= 0) {
+    // write the remaining data
+    flush_compressed_data();
+    close(FD);
+  }
+  deflateEnd(&strm);
+}
+
+void compressed_fd_ostream::write_impl(const char *Ptr, size_t Size) {
+  strm.next_in =
+      const_cast<unsigned char *>(reinterpret_cast<const unsigned char *>(Ptr));
+  strm.avail_in = Size;
+
+  // Check if there is still data to compress
+  while (strm.avail_in != 0) {
+    // compress data
+    int res = deflate(&strm, Z_NO_FLUSH);
+    assert(res == Z_OK);
+    writeFullCompressedData();
+  }
+}
+
+void compressed_fd_ostream::write_file(const char *Ptr, size_t Size) {
+  pos += Size;
+  assert(FD >= 0 && "File already closed");
+  do {
+    ssize_t ret = ::write(FD, Ptr, Size);
+    if (ret < 0) {
+      if (errno == EINTR || errno == EAGAIN)
+        continue;
+      assert(0 && "Could not write to file");
+      break;
+    }
+
+    Ptr += ret;
+    Size -= ret;
+  } while (Size > 0);
+}
+}
+#endif