4 files changed, 65 insertions, 2 deletions
diff --git a/include/klee/Internal/ADT/DiscretePDF.inc b/include/klee/Internal/ADT/DiscretePDF.inc
index 32328c95..743a69b5 100644
--- a/include/klee/Internal/ADT/DiscretePDF.inc
+++ b/include/klee/Internal/ADT/DiscretePDF.inc
@@ -7,6 +7,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include <cassert>
 namespace klee {
 
 template <class T>
@@ -162,8 +163,7 @@ void DiscretePDF<T>::update(T item, weight_type weight) {
 
 template <class T>
 T DiscretePDF<T>::choose(double p) {
-  if ((p < 0.0) || (p >= 1.0))
-    assert(0 && "choose: argument(p) outside valid range");
+  assert (!((p < 0.0) || (p >= 1.0)) && "choose: argument(p) outside valid range");
 
   if (!m_root)
     assert(0 && "choose: choose() called on empty tree");
diff --git a/unittests/CMakeLists.txt b/unittests/CMakeLists.txt
index 63934336..00d4cbfa 100644
--- a/unittests/CMakeLists.txt
+++ b/unittests/CMakeLists.txt
@@ -85,6 +85,7 @@ add_subdirectory(Expr)
 add_subdirectory(Ref)
 add_subdirectory(Solver)
 add_subdirectory(TreeStream)
+add_subdirectory(DiscretePDF)
 
 # Set up lit configuration
 set (UNIT_TEST_EXE_SUFFIX "Test")
diff --git a/unittests/DiscretePDF/CMakeLists.txt b/unittests/DiscretePDF/CMakeLists.txt
new file mode 100644
index 00000000..07ec1f89
--- /dev/null
+++ b/unittests/DiscretePDF/CMakeLists.txt
@@ -0,0 +1,4 @@
+add_klee_unit_test(DiscretePDFTest
+  DiscretePDFTest.cpp)
+# FIXME add the following line to link against libgtest.a
+target_link_libraries(DiscretePDFTest PRIVATE kleaverSolver)
\ No newline at end of file
diff --git a/unittests/DiscretePDF/DiscretePDFTest.cpp b/unittests/DiscretePDF/DiscretePDFTest.cpp
new file mode 100644
index 00000000..2180360c
--- /dev/null
+++ b/unittests/DiscretePDF/DiscretePDFTest.cpp
@@ -0,0 +1,58 @@
+#include "klee/Internal/ADT/DiscretePDF.h"
+#include "gtest/gtest.h"
+#include <iostream>
+#include <vector>
+
+int finished = 0;
+
+using namespace klee;
+
+TEST(DiscretePDFTest, Rotatation) {
+  DiscretePDF<int> testTree;
+
+  ASSERT_TRUE(testTree.empty());
+
+  for (auto i = 0; i < 9; ++i)
+    testTree.insert(i + 10, 0.1 * i);
+
+  for (auto i = 9; i > 0; --i)
+    testTree.insert(i + 100, 0.01 * i);
+
+  ASSERT_FALSE(testTree.empty());
+
+  ASSERT_TRUE(testTree.inTree(101));
+  testTree.remove(101);
+  ASSERT_FALSE(testTree.inTree(101));
+  testTree.insert(101, 0.01);
+  ASSERT_EQ(0.01, testTree.getWeight(101));
+
+  ASSERT_EQ(11, testTree.choose(0.01));
+  testTree.choose(0);
+  testTree.choose(0.9999999);
+
+  testTree.update(101, 0.9);
+
+  for (auto i = 0; i < 9; ++i)
+    testTree.remove(i + 10);
+  for (auto i = 0; i < 50; ++i) {
+    testTree.insert(50 + i, 2 * i);
+  }
+
+  for (auto i = 9; i > 0; --i)
+    testTree.remove(i + 100);
+
+  testTree.insert(1, 0);
+#ifndef NDEBUG
+  ASSERT_DEATH({ testTree.insert(1, 0); }, "already in tree");
+#endif
+
+  while (!testTree.empty()) {
+    testTree.remove(testTree.choose(0));
+  }
+
+  testTree.insert(1, 1);
+  testTree.insert(2, 2);
+
+  ASSERT_EQ(1, testTree.getWeight(1));
+  ASSERT_EQ(2, testTree.getWeight(2));
+}