From 9d5bbe8a309c5760166febf7edfc3b790d77ab7e Mon Sep 17 00:00:00 2001
From: Cristian Cadar <c.cadar@imperial.ac.uk>
Date: Fri, 20 Oct 2023 22:03:38 +0100
Subject: On a symbolic allocation, retrieve size from a seed, if available

---
 lib/Core/Executor.cpp               | 47 +++++++++++++++++++++----------------
 test/Feature/SeedConcretizeMalloc.c | 32 +++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 20 deletions(-)
 create mode 100644 test/Feature/SeedConcretizeMalloc.c

diff --git a/lib/Core/Executor.cpp b/lib/Core/Executor.cpp
index ebbcfaea..bc9e31f7 100644
--- a/lib/Core/Executor.cpp
+++ b/lib/Core/Executor.cpp
@@ -4205,24 +4205,30 @@ void Executor::executeAlloc(ExecutionState &state,
     size = optimizer.optimizeExpr(size, true);
 
     ref<ConstantExpr> example;
-    bool success =
-        solver->getValue(state.constraints, size, example, state.queryMetaData);
-    assert(success && "FIXME: Unhandled solver failure");
-    (void) success;
-    
-    // Try and start with a small example.
-    Expr::Width W = example->getWidth();
-    while (example->Ugt(ConstantExpr::alloc(128, W))->isTrue()) {
-      ref<ConstantExpr> tmp = example->LShr(ConstantExpr::alloc(1, W));
-      bool res;
-      bool success =
-          solver->mayBeTrue(state.constraints, EqExpr::create(tmp, size), res,
-                            state.queryMetaData);
-      assert(success && "FIXME: Unhandled solver failure");      
-      (void) success;
-      if (!res)
-        break;
-      example = tmp;
+    // Check if in seed mode, then try to replicate size from a seed
+    ref<Expr> value = nullptr;
+    if (auto found = seedMap.find(&state); found != seedMap.end())
+      value = getValueFromSeeds(found->second, size);
+    /* If no seed evaluation results in a constant, call the solver */
+    if (!value || !(example = dyn_cast<ConstantExpr>(value))) {
+      bool success = solver->getValue(state.constraints, size, example,
+                                      state.queryMetaData);
+      assert(success && "FIXME: Unhandled solver failure");
+      (void)success;
+
+      // Try and start with a small example.
+      Expr::Width W = example->getWidth();
+      while (example->Ugt(ConstantExpr::alloc(128, W))->isTrue()) {
+        ref<ConstantExpr> tmp = example->LShr(ConstantExpr::alloc(1, W));
+        bool res;
+        [[maybe_unused]] bool success =
+            solver->mayBeTrue(state.constraints, EqExpr::create(tmp, size), res,
+                              state.queryMetaData);
+        assert(success && "FIXME: Unhandled solver failure");
+        if (!res)
+          break;
+        example = tmp;
+      }
     }
 
     StatePair fixedSize =
@@ -4249,8 +4255,9 @@ void Executor::executeAlloc(ExecutionState &state,
         // malloc will fail for it, so lets fork and return 0.
         StatePair hugeSize =
             fork(*fixedSize.second,
-                 UltExpr::create(ConstantExpr::alloc(1U << 31, W), size), true,
-                 BranchType::Alloc);
+                 UltExpr::create(
+                     ConstantExpr::alloc(1U << 31, example->getWidth()), size),
+                 true, BranchType::Alloc);
         if (hugeSize.first) {
           klee_message("NOTE: found huge malloc, returning 0");
           bindLocal(target, *hugeSize.first, 
diff --git a/test/Feature/SeedConcretizeMalloc.c b/test/Feature/SeedConcretizeMalloc.c
new file mode 100644
index 00000000..1f916723
--- /dev/null
+++ b/test/Feature/SeedConcretizeMalloc.c
@@ -0,0 +1,32 @@
+// RUN: %clang -emit-llvm -c -g %s -o %t.bc
+// RUN: rm -rf %t.klee-out
+// RUN: %klee --output-dir=%t.klee-out --entry-point=SeedGen %t.bc
+// RUN: test -f %t.klee-out/test000001.ktest
+// RUN: not test -f %t.klee-out/test000002.ktest
+
+// RUN: rm -rf %t.klee-out-2
+// RUN: %klee --output-dir=%t.klee-out-2 --seed-file %t.klee-out/test000001.ktest %t.bc | FileCheck --allow-empty %s
+
+#include "klee/klee.h"
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+void SeedGen() {
+  unsigned x;
+  klee_make_symbolic(&x, sizeof(x), "x");
+  klee_assume(x == 100);
+}
+
+int main(int argc, char **argv) {
+  unsigned s;
+  klee_make_symbolic(&s, sizeof(s), "size");
+  char *p = (char *)malloc(s);
+  if (!p)
+    return 0;
+
+  if (s != 100)
+    printf("Error\n");
+    // CHECK-NOT: Error
+}
-- 
cgit 1.4.1