work list algorithm for exploration

Author: butterunderflow

headers/wasm/concolic_driver.hpp

@@ -42,6 +42,7 @@ class ConcolicDriver {
   Solver solver;
   std::function<void()> entrypoint;
   std::optional<std::string> tree_file;
+  std::vector<NodeBox *> work_list;
 };
 
 class ManagedConcolicCleanup {
@@ -64,19 +65,39 @@ class ManagedConcolicCleanup {
 static std::monostate reset_stacks();
 
 inline void ConcolicDriver::main_exploration_loop() {
-  while (true) {
+
+  // Register a collector to ExploreTree to add new nodes to work_list
+  ExploreTree.register_new_node_collector(
+      [&](NodeBox *new_node) { work_list.push_back(new_node); });
+
+  std::set<NodeBox *> visited;
+
+  assert(ExploreTree.get_root()->isUnexplored() &&
+         "Before main loop, root should be unexplored!");
+  work_list.push_back(ExploreTree.get_root());
+
+  while (!work_list.empty()) {
     ManagedConcolicCleanup cleanup{*this};
+    // Pick an unexplored node from the work list
+    auto node = work_list.back();
+    work_list.pop_back();
 
-    auto unexplored = ExploreTree.pick_unexplored();
-    if (!unexplored) {
-      GENSYM_INFO("No unexplored nodes found, exiting...");
-      return;
+    if (visited.find(node) != visited.end()) {
+      continue;
+    } else {
+      visited.insert(node);
     }
-    auto cond = unexplored->collect_path_conds();
+
+    if (!node->isUnexplored()) {
+      // if it's not unexplored anymore, skip it
+      continue;
+    }
+
+    auto cond = node->collect_path_conds();
     auto result = solver.solve(cond);
     if (!result.has_value()) {
       GENSYM_INFO("Found an unreachable path, marking it as unreachable...");
-      unexplored->fillUnreachableNode();
+      node->fillUnreachableNode();
       continue;
     }
     auto new_env = result.value();
@@ -87,8 +108,8 @@ inline void ConcolicDriver::main_exploration_loop() {
       GENSYM_INFO("Now execute the program with symbolic environment: ");
       GENSYM_INFO(SymEnv.to_string());
       if (auto snapshot_node =
-              dynamic_cast<SnapshotNode *>(unexplored->node.get())) {
-        snapshot_node->get_snapshot().resume_execution(SymEnv, unexplored);
+              dynamic_cast<SnapshotNode *>(node->node.get())) {
+        snapshot_node->get_snapshot().resume_execution(SymEnv, node);
       } else {
         auto timer = ManagedTimer();
         reset_stacks();

headers/wasm/symbolic_rt.hpp

@@ -420,7 +420,7 @@ struct NodeBox {
   std::unique_ptr<Node> node;
   NodeBox *parent;
 
-  std::monostate fillIfElseNode(SymVal cond, int id);
+  bool fillIfElseNode(SymVal cond, int id);
   std::monostate fillFinishedNode();
   std::monostate fillFailedNode();
   std::monostate fillUnreachableNode();
@@ -610,15 +610,16 @@ inline NodeBox::NodeBox(NodeBox *parent)
       /* TODO: avoid allocation of unexplored node */
       parent(parent) {}
 
-inline std::monostate NodeBox::fillIfElseNode(SymVal cond, int id) {
+inline bool NodeBox::fillIfElseNode(SymVal cond, int id) {
   // fill the current NodeBox with an ifelse branch node when it's unexplored
   if (this->isUnexplored()) {
     node = std::make_unique<IfElseNode>(cond, this, id);
+    return true;
   }
   assert(
       dynamic_cast<IfElseNode *>(node.get()) != nullptr &&
       "Current node is not an Unexplored nor an IfElseNode, cannot fill it!");
-  return std::monostate();
+  return false;
 }
 
 inline std::monostate NodeBox::fillSnapshotNode(Snapshot_t snapshot) {
@@ -722,7 +723,12 @@ class ExploreTree_t {
   std::monostate fillFailedNode() { return cursor->fillFailedNode(); }
 
   std::monostate fillIfElseNode(SymVal cond, int id) {
-    return cursor->fillIfElseNode(cond, id);
+    if (cursor->fillIfElseNode(cond, id)) {
+      auto if_else_node = dynamic_cast<IfElseNode *>(cursor->node.get());
+      register_new_node(if_else_node->true_branch.get());
+      register_new_node(if_else_node->false_branch.get());
+    }
+    return std::monostate();
   }
 
   std::monostate moveCursor(bool branch, Snapshot_t snapshot) {
@@ -764,16 +770,6 @@ class ExploreTree_t {
     return std::monostate();
   }
 
-  std::optional<std::vector<SymVal>> get_unexplored_conditions() {
-    // Get all unexplored conditions in the tree
-    std::vector<SymVal> result;
-    auto box = pick_unexplored();
-    if (!box) {
-      return std::nullopt;
-    }
-    return box->collect_path_conds();
-  }
-
   NodeBox *pick_unexplored() {
     // Pick an unexplored node from the tree
     // For now, we just iterate through the tree and return the first unexplored
@@ -793,6 +789,12 @@ class ExploreTree_t {
     return true;
   }
 
+  NodeBox *get_root() const { return root.get(); }
+
+  void register_new_node_collector(std::function<void(NodeBox *)> func) {
+    new_node_collectors.push_back(func);
+  }
+
 private:
   NodeBox *pick_unexplored_of(NodeBox *node) {
     if (node->isUnexplored()) {
@@ -808,8 +810,14 @@ class ExploreTree_t {
     }
     return nullptr; // No unexplored node found
   }
+  void register_new_node(NodeBox *node) {
+    for (auto &func : new_node_collectors) {
+      func(node);
+    }
+  }
   std::unique_ptr<NodeBox> root;
   NodeBox *cursor;
+  std::vector<std::function<void(NodeBox *)>> new_node_collectors;
 };
 
 static ExploreTree_t ExploreTree;