Merge pull request #704 from ccadar/TreeStream

Fixes bug in TreeStreamWriter::write reported by @gladtbx in #562.  A…

6 files changed, 65 insertions, 8 deletions

diff --git a/lib/Support/TreeStream.cpp b/lib/Support/TreeStream.cpp
index ef59b2a9..a0e1596f 100644
--- a/lib/Support/TreeStream.cpp
+++ b/lib/Support/TreeStream.cpp
@@ -63,7 +63,6 @@ TreeOStream TreeStreamWriter::open(const TreeOStream &os) {
 }
 
 void TreeStreamWriter::write(TreeOStream &os, const char *s, unsigned size) {
-#if 1
   if (bufferCount && 
       (os.id!=lastID || size+bufferCount>bufferSize))
     flushBuffer();
@@ -77,13 +76,8 @@ void TreeStreamWriter::write(TreeOStream &os, const char *s, unsigned size) {
   } else {
     output->write(reinterpret_cast<const char*>(&os.id), 4);
     output->write(reinterpret_cast<const char*>(&size), 4);
-    output->write(buffer, size);
+    output->write(s, size);
   }
-#else
-  output->write(reinterpret_cast<const char*>(&os.id), 4);
-  output->write(reinterpret_cast<const char*>(&size), 4);
-  output->write(s, size);
-#endif
 }
 
 void TreeStreamWriter::flushBuffer() {
diff --git a/unittests/CMakeLists.txt b/unittests/CMakeLists.txt
index a85c0a74..85b7b061 100644
--- a/unittests/CMakeLists.txt
+++ b/unittests/CMakeLists.txt
@@ -79,6 +79,7 @@ add_subdirectory(Assignment)
 add_subdirectory(Expr)
 add_subdirectory(Ref)
 add_subdirectory(Solver)
+add_subdirectory(TreeStream)
 
 # Set up lit configuration
 set (UNIT_TEST_EXE_SUFFIX "Test")
diff --git a/unittests/Makefile b/unittests/Makefile
index 2a73b809..582b8d8d 100644
--- a/unittests/Makefile
+++ b/unittests/Makefile
@@ -17,7 +17,7 @@ CPP.Flags += -I$(LLVM_SRC_ROOT)/utils/unittest/googletest/include/
 CPP.Flags += -Wno-variadic-macros
 
 # FIXME: Parallel dirs is broken?
-DIRS = Expr Solver Ref Assignment
+DIRS = Expr Solver Ref Assignment TreeStream
 
 include $(LEVEL)/Makefile.common
 
diff --git a/unittests/TreeStream/CMakeLists.txt b/unittests/TreeStream/CMakeLists.txt
new file mode 100644
index 00000000..4b298724
--- /dev/null
+++ b/unittests/TreeStream/CMakeLists.txt
@@ -0,0 +1,3 @@
+add_klee_unit_test(TreeStreamTest
+  TreeStreamTest.cpp)
+target_link_libraries(TreeStreamTest PRIVATE kleeBasic)
diff --git a/unittests/TreeStream/Makefile b/unittests/TreeStream/Makefile
new file mode 100644
index 00000000..a56659b9
--- /dev/null
+++ b/unittests/TreeStream/Makefile
@@ -0,0 +1,11 @@
+##===- unittests/Expr/Makefile -----------------------------*- Makefile -*-===##
+
+LEVEL := ../..
+include $(LEVEL)/Makefile.config
+
+TESTNAME := TreeStreamTest
+USEDLIBS := kleeBasic.a kleeSupport.a
+LINK_COMPONENTS := support
+
+include $(LLVM_SRC_ROOT)/unittests/Makefile.unittest
+
diff --git a/unittests/TreeStream/TreeStreamTest.cpp b/unittests/TreeStream/TreeStreamTest.cpp
new file mode 100644
index 00000000..09d62c6e
--- /dev/null
+++ b/unittests/TreeStream/TreeStreamTest.cpp
@@ -0,0 +1,48 @@
+#include "klee/Internal/ADT/TreeStream.h"
+#include <vector>
+#include <cstring>
+
+#include "gtest/gtest.h"
+
+using namespace klee;
+
+/* Basic test, checking that after writing "abc" and then "defg", we
+   get a {'a', 'b', 'c', 'c', 'd', 'e', 'f', 'g' } back.  */
+TEST(TreeStreamTest, Basic) {
+  TreeStreamWriter tsw("tsw1.out");
+  ASSERT_TRUE(tsw.good());
+  
+  TreeOStream tos = tsw.open();
+  tos.write("abc", 3);
+  tos.write("defg", 4);
+  tos.flush();
+
+  std::vector<unsigned char> out;
+  tsw.readStream(tos.getID(), out);
+  ASSERT_EQ(out.size(), 7);
+  
+  for (unsigned char c = 'a'; c <= 'g'; c++)
+    ASSERT_EQ(out[c - 'a'], c);
+}
+
+
+/* This tests the case when we perform a write with a size larger than
+   the buffer size, which is a constant set to 4*4096.  This test fails 
+   without #704 */
+TEST(TreeStreamTest, WriteLargerThanBufferSize) {
+  TreeStreamWriter tsw("tsw2.out");
+  ASSERT_TRUE(tsw.good());
+  
+  TreeOStream tos = tsw.open();
+#define NBYTES 5*4096
+  char buf[NBYTES];
+  memset(buf, 'A', sizeof(buf));
+  tos.write(buf, NBYTES);
+  tos.flush();
+
+  std::vector<unsigned char> out;
+  tsw.readStream(tos.getID(), out);
+  ASSERT_EQ(out.size(), NBYTES);
+  for (unsigned i=0; i<out.size(); i++)
+    ASSERT_EQ('A', out[i]);
+}