Fix bug in `AssignmentEvaluator` where NotOptimizedExpr would not (#466)

* Add unittest to check that the `Assignment` class can evaluate
expressions containing a `NotOptimizedExpr`.

* Fix the `AssignmentTest.FoldNotOptimized` unit test by
teaching the `ExprEvaluator` to fold `NotOptimizedExpr` nodes.

5 files changed, 58 insertions, 1 deletions

diff --git a/include/klee/util/ExprEvaluator.h b/include/klee/util/ExprEvaluator.h
index 6b67a1cf..313ac0a5 100644
--- a/include/klee/util/ExprEvaluator.h
+++ b/include/klee/util/ExprEvaluator.h
@@ -25,6 +25,7 @@ namespace klee {
     Action visitSDiv(const SDivExpr &e);
     Action visitURem(const URemExpr &e);
     Action visitSRem(const SRemExpr &e);
+    Action visitExprPost(const Expr& e);
       
   public:
     ExprEvaluator() {}
diff --git a/lib/Expr/ExprEvaluator.cpp b/lib/Expr/ExprEvaluator.cpp
index 1a146641..6b84cd6f 100644
--- a/lib/Expr/ExprEvaluator.cpp
+++ b/lib/Expr/ExprEvaluator.cpp
@@ -97,3 +97,12 @@ ExprVisitor::Action ExprEvaluator::visitURem(const URemExpr &e) {
 ExprVisitor::Action ExprEvaluator::visitSRem(const SRemExpr &e) { 
   return protectedDivOperation(e); 
 }
+
+ExprVisitor::Action ExprEvaluator::visitExprPost(const Expr& e) {
+  // When evaluating an assignment we should fold NotOptimizedExpr
+  // nodes so we can fully evaluate.
+  if (e.getKind() == Expr::NotOptimized) {
+    return Action::changeTo(static_cast<const NotOptimizedExpr&>(e).src);
+  }
+  return Action::skipChildren();
+}
diff --git a/unittests/Assignment/AssignmentTest.cpp b/unittests/Assignment/AssignmentTest.cpp
new file mode 100644
index 00000000..0eaa28f1
--- /dev/null
+++ b/unittests/Assignment/AssignmentTest.cpp
@@ -0,0 +1,35 @@
+#include "klee/util/ArrayCache.h"
+#include "klee/util/Assignment.h"
+#include "gtest/gtest.h"
+#include <iostream>
+#include <vector>
+
+int finished = 0;
+
+using namespace klee;
+
+TEST(AssignmentTest, FoldNotOptimized)
+{
+  ArrayCache ac;
+  const Array* array = ac.CreateArray("simple_array", /*size=*/ 1);
+  // Create a simple assignment
+  std::vector<const Array*> objects;
+  std::vector<unsigned char> value;
+  std::vector< std::vector<unsigned char> > values;
+  objects.push_back(array);
+  value.push_back(128);
+  values.push_back(value);
+  // We want to simplify to a constant so allow free values so
+  // if the assignment is incomplete we don't get back a constant.
+  Assignment assignment(objects, values, /*_allowFreeValues=*/true);
+
+  // Now make an expression that reads from the array at position
+  // zero.
+  ref<Expr> read = NotOptimizedExpr::alloc(Expr::createTempRead(array, Expr::Int8));
+
+  // Now evaluate. The OptimizedExpr should be folded
+  ref<Expr> evaluated = assignment.evaluate(read);
+  const ConstantExpr* asConstant = dyn_cast<ConstantExpr>(evaluated);
+  ASSERT_TRUE(asConstant != NULL);
+  ASSERT_EQ(asConstant->getZExtValue(), (unsigned) 128);
+}
diff --git a/unittests/Assignment/Makefile b/unittests/Assignment/Makefile
new file mode 100644
index 00000000..d4110c80
--- /dev/null
+++ b/unittests/Assignment/Makefile
@@ -0,0 +1,12 @@
+##===- unittests/Assignment/Makefile -----------------------*- Makefile -*-===##
+
+LEVEL := ../..
+include $(LEVEL)/Makefile.config
+
+TESTNAME := Assignment
+USEDLIBS := kleaverExpr.a
+LINK_COMPONENTS := support
+
+include $(LLVM_SRC_ROOT)/unittests/Makefile.unittest
+
+CXXFLAGS += -DLLVM_29_UNITTEST
diff --git a/unittests/Makefile b/unittests/Makefile
index b7e51371..2a73b809 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
+DIRS = Expr Solver Ref Assignment
 
 include $(LEVEL)/Makefile.common