1 files changed, 235 insertions, 0 deletions
diff --git a/lib/Core/ExecutionTree.cpp b/lib/Core/ExecutionTree.cpp
new file mode 100644
index 00000000..c5b640d3
--- /dev/null
+++ b/lib/Core/ExecutionTree.cpp
@@ -0,0 +1,235 @@
+//===-- PTree.cpp ---------------------------------------------------------===//
+//
+//                     The KLEE Symbolic Virtual Machine
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PTree.h"
+
+#include "ExecutionState.h"
+
+#include "klee/Core/Interpreter.h"
+#include "klee/Expr/ExprPPrinter.h"
+#include "klee/Module/KInstruction.h"
+#include "klee/Support/OptionCategories.h"
+
+#include <bitset>
+#include <vector>
+
+using namespace klee;
+
+namespace klee {
+llvm::cl::OptionCategory
+    PTreeCat("Process tree related options",
+             "These options affect the process tree handling.");
+}
+
+namespace {
+llvm::cl::opt<bool> CompressProcessTree(
+    "compress-process-tree",
+    llvm::cl::desc("Remove intermediate nodes in the process "
+                   "tree whenever possible (default=false)"),
+    llvm::cl::init(false), llvm::cl::cat(PTreeCat));
+
+llvm::cl::opt<bool> WritePTree(
+    "write-ptree", llvm::cl::init(false),
+    llvm::cl::desc("Write process tree into ptree.db (default=false)"),
+    llvm::cl::cat(PTreeCat));
+} // namespace
+
+// PTreeNode
+
+PTreeNode::PTreeNode(PTreeNode *parent, ExecutionState *state) noexcept
+    : parent{parent}, left{nullptr}, right{nullptr}, state{state} {
+  state->ptreeNode = this;
+}
+
+// AnnotatedPTreeNode
+
+AnnotatedPTreeNode::AnnotatedPTreeNode(PTreeNode *parent,
+                                       ExecutionState *state) noexcept
+    : PTreeNode(parent, state) {
+  id = nextID++;
+}
+
+// NoopPTree
+
+void NoopPTree::dump(llvm::raw_ostream &os) noexcept {
+  os << "digraph G {\nTreeNotAvailable [shape=box]\n}";
+}
+
+// InMemoryPTree
+
+InMemoryPTree::InMemoryPTree(ExecutionState &initialState) noexcept {
+  root = PTreeNodePtr(createNode(nullptr, &initialState));
+  initialState.ptreeNode = root.getPointer();
+}
+
+PTreeNode *InMemoryPTree::createNode(PTreeNode *parent, ExecutionState *state) {
+  return new PTreeNode(parent, state);
+}
+
+void InMemoryPTree::attach(PTreeNode *node, ExecutionState *leftState,
+                           ExecutionState *rightState,
+                           BranchType reason) noexcept {
+  assert(node && !node->left.getPointer() && !node->right.getPointer());
+  assert(node == rightState->ptreeNode &&
+         "Attach assumes the right state is the current state");
+  node->left = PTreeNodePtr(createNode(node, leftState));
+  // The current node inherits the tag
+  uint8_t currentNodeTag = root.getInt();
+  if (node->parent)
+    currentNodeTag = node->parent->left.getPointer() == node
+                         ? node->parent->left.getInt()
+                         : node->parent->right.getInt();
+  node->right = PTreeNodePtr(createNode(node, rightState), currentNodeTag);
+  updateBranchingNode(*node, reason);
+  node->state = nullptr;
+}
+
+void InMemoryPTree::remove(PTreeNode *n) noexcept {
+  assert(!n->left.getPointer() && !n->right.getPointer());
+  updateTerminatingNode(*n);
+  do {
+    PTreeNode *p = n->parent;
+    if (p) {
+      if (n == p->left.getPointer()) {
+        p->left = PTreeNodePtr(nullptr);
+      } else {
+        assert(n == p->right.getPointer());
+        p->right = PTreeNodePtr(nullptr);
+      }
+    }
+    delete n;
+    n = p;
+  } while (n && !n->left.getPointer() && !n->right.getPointer());
+
+  if (n && CompressProcessTree) {
+    // We're now at a node that has exactly one child; we've just deleted the
+    // other one. Eliminate the node and connect its child to the parent
+    // directly (if it's not the root).
+    PTreeNodePtr child = n->left.getPointer() ? n->left : n->right;
+    PTreeNode *parent = n->parent;
+
+    child.getPointer()->parent = parent;
+    if (!parent) {
+      // We're at the root.
+      root = child;
+    } else {
+      if (n == parent->left.getPointer()) {
+        parent->left = child;
+      } else {
+        assert(n == parent->right.getPointer());
+        parent->right = child;
+      }
+    }
+
+    delete n;
+  }
+}
+
+void InMemoryPTree::dump(llvm::raw_ostream &os) noexcept {
+  std::unique_ptr<ExprPPrinter> pp(ExprPPrinter::create(os));
+  pp->setNewline("\\l");
+  os << "digraph G {\n"
+     << "\tsize=\"10,7.5\";\n"
+     << "\tratio=fill;\n"
+     << "\trotate=90;\n"
+     << "\tcenter = \"true\";\n"
+     << "\tnode [style=\"filled\",width=.1,height=.1,fontname=\"Terminus\"]\n"
+     << "\tedge [arrowsize=.3]\n";
+  std::vector<const PTreeNode *> stack;
+  stack.push_back(root.getPointer());
+  while (!stack.empty()) {
+    const PTreeNode *n = stack.back();
+    stack.pop_back();
+    os << "\tn" << n << " [shape=diamond";
+    if (n->state)
+      os << ",fillcolor=green";
+    os << "];\n";
+    if (n->left.getPointer()) {
+      os << "\tn" << n << " -> n" << n->left.getPointer() << " [label=0b"
+         << std::bitset<PtrBitCount>(n->left.getInt()).to_string() << "];\n";
+      stack.push_back(n->left.getPointer());
+    }
+    if (n->right.getPointer()) {
+      os << "\tn" << n << " -> n" << n->right.getPointer() << " [label=0b"
+         << std::bitset<PtrBitCount>(n->right.getInt()).to_string() << "];\n";
+      stack.push_back(n->right.getPointer());
+    }
+  }
+  os << "}\n";
+}
+
+std::uint8_t InMemoryPTree::getNextId() noexcept {
+  static_assert(PtrBitCount <= 8);
+  std::uint8_t id = 1 << registeredIds++;
+  if (registeredIds > PtrBitCount) {
+    klee_error("PTree cannot support more than %d RandomPathSearchers",
+               PtrBitCount);
+  }
+  return id;
+}
+
+// PersistentPTree
+
+PersistentPTree::PersistentPTree(ExecutionState &initialState,
+                                 InterpreterHandler &ih) noexcept
+    : writer(ih.getOutputFilename("ptree.db")) {
+  root = PTreeNodePtr(createNode(nullptr, &initialState));
+  initialState.ptreeNode = root.getPointer();
+}
+
+void PersistentPTree::dump(llvm::raw_ostream &os) noexcept {
+  writer.batchCommit(true);
+  InMemoryPTree::dump(os);
+}
+
+PTreeNode *PersistentPTree::createNode(PTreeNode *parent,
+                                       ExecutionState *state) {
+  return new AnnotatedPTreeNode(parent, state);
+}
+
+void PersistentPTree::setTerminationType(ExecutionState &state,
+                                         StateTerminationType type) {
+  auto *annotatedNode = llvm::cast<AnnotatedPTreeNode>(state.ptreeNode);
+  annotatedNode->kind = type;
+}
+
+void PersistentPTree::updateBranchingNode(PTreeNode &node, BranchType reason) {
+  auto *annotatedNode = llvm::cast<AnnotatedPTreeNode>(&node);
+  const auto &state = *node.state;
+  const auto prevPC = state.prevPC;
+  annotatedNode->asmLine =
+      prevPC && prevPC->info ? prevPC->info->assemblyLine : 0;
+  annotatedNode->kind = reason;
+  writer.write(*annotatedNode);
+}
+
+void PersistentPTree::updateTerminatingNode(PTreeNode &node) {
+  assert(node.state);
+  auto *annotatedNode = llvm::cast<AnnotatedPTreeNode>(&node);
+  const auto &state = *node.state;
+  const auto prevPC = state.prevPC;
+  annotatedNode->asmLine =
+      prevPC && prevPC->info ? prevPC->info->assemblyLine : 0;
+  annotatedNode->stateID = state.getID();
+  writer.write(*annotatedNode);
+}
+
+// Factory
+
+std::unique_ptr<PTree> klee::createPTree(ExecutionState &initialState,
+                                         bool inMemory,
+                                         InterpreterHandler &ih) {
+  if (WritePTree)
+    return std::make_unique<PersistentPTree>(initialState, ih);
+
+  if (inMemory)
+    return std::make_unique<InMemoryPTree>(initialState);
+
+  return std::make_unique<NoopPTree>();
+};