1 files changed, 32 insertions, 0 deletions
diff --git a/lib/Core/Executor.cpp b/lib/Core/Executor.cpp
index 73dc21fa..4763fdb5 100644
--- a/lib/Core/Executor.cpp
+++ b/lib/Core/Executor.cpp
@@ -62,6 +62,7 @@
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/InlineAsm.h"
 #include "llvm/IR/Instructions.h"
@@ -135,6 +136,10 @@ cl::opt<bool> SingleObjectResolution(
     cl::desc("Try to resolve memory reads/writes to single objects "
              "when offsets are symbolic (default=false)"),
     cl::init(false), cl::cat(MiscCat));
+
+cl::opt<std::string> PatchedFile(
+    "patched-file",
+    cl::desc("Relative path to patched file"));
 } // namespace klee
 
 namespace {
@@ -2120,6 +2125,30 @@ Function *Executor::getTargetFunction(Value *calledVal) {
   }
 }
 
+void Executor::registerOutput(ExecutionState &state, Function *fn,
+    std::string var, Expr::Width bits, ref<Expr> val) {
+  if (state.patchLocs == 0 || fn->getName() == "__choose"
+      || (!PatchedFile.empty()
+          && fn->getSubprogram()->getFilename() != PatchedFile))
+    return;
+  size_t bytes = bits == Expr::Bool ? 1u : bits >> 3;
+  auto const inst = state.prevPC->inst;
+  auto const mo = this->memory->allocate(bytes, true, false,
+                                         &state, inst, bytes);
+  std::string baseName = "out!" + fn->getName().str() + "!" + var + "!";
+  std::string name;
+  unsigned id = 0;
+  do
+    name = baseName + llvm::utostr(id++);
+  while (!state.arrayNames.insert(name).second);
+  this->executeMakeSymbolic(state, mo, name);
+  // Cached symbolic array
+  auto const outArray = this->arrayCache.CreateArray(name, bytes);
+  auto const outVal = Expr::createTempRead(outArray, bits);
+  auto const eq = EqExpr::create(outVal, val);
+  this->addConstraint(state, eq);
+}
+
 void Executor::executeInstruction(ExecutionState &state, KInstruction *ki) {
   Instruction *i = ki->inst;
   switch (i->getOpcode()) {
@@ -2206,6 +2235,9 @@ void Executor::executeInstruction(ExecutionState &state, KInstruction *ki) {
             }
           }
 
+          auto const fp = cast<CallBase>(*caller).getCalledOperand();
+          if (auto const fn = this->getTargetFunction(fp))
+            registerOutput(state, fn, "ret", to, result);
           bindLocal(kcaller, state, result);
         }
       } else {