support sym_assume to prune subtrees

Author: butterunderflow

headers/wasm/concolic_driver.hpp

@@ -48,7 +48,7 @@ class ManagedConcolicCleanup {
     if (driver.tree_file.has_value())
       ExploreTree.dump_graphviz(driver.tree_file.value());
 
-    // Clear the symbol bookkeeper
+      // Clear the symbol bookkeeper
     SymBookKeeper.clear();
   }
 };
@@ -134,6 +134,7 @@ inline void ConcolicDriver::main_exploration_loop() {
 
 inline void ConcolicDriver::run() {
   main_exploration_loop();
+  ExploreTree.print_overall_result();
   Profile.print_summary();
 }
 

headers/wasm/symbolic_rt.hpp

@@ -543,6 +543,7 @@ struct NodeBox {
   std::monostate fillFailedNode();
   std::monostate fillUnreachableNode();
   std::monostate fillSnapshotNode(Snapshot_t snapshot);
+  std::monostate fillNotToExploreNode();
   bool isUnexplored() const;
   std::vector<SymVal> collect_path_conds();
   int min_cost_of_reaching_here();
@@ -664,6 +665,22 @@ struct UnExploredNode : Node {
   }
 };
 
+struct NotToExploreNode : Node {
+  NotToExploreNode() {}
+  std::string to_string() override { return "NotToExploreNode"; }
+
+protected:
+  void generate_dot(std::ostream &os, int parent_dot_id,
+                    const std::string &edge_label) override {
+    int current_node_dot_id = current_id++;
+    graphviz_node(os, current_node_dot_id, "NotToExplore", "box", "grey");
+
+    if (parent_dot_id != -1) {
+      graphviz_edge(os, parent_dot_id, current_node_dot_id, edge_label);
+    }
+  }
+};
+
 struct SnapshotNode : Node {
   SnapshotNode(Snapshot_t snapshot) : snapshot(snapshot) {}
   std::string to_string() override { return "SnapshotNode"; }
@@ -761,6 +778,15 @@ inline std::monostate NodeBox::fillSnapshotNode(Snapshot_t snapshot) {
   return std::monostate();
 }
 
+inline std::monostate NodeBox::fillNotToExploreNode() {
+  if (this->isUnexplored()) {
+    node = std::make_unique<NotToExploreNode>();
+  } else {
+    assert(dynamic_cast<NotToExploreNode *>(node.get()) != nullptr);
+  }
+  return std::monostate();
+}
+
 inline std::monostate NodeBox::fillFinishedNode() {
   if (this->isUnexplored()) {
     node = std::make_unique<Finished>();
@@ -896,6 +922,10 @@ class ExploreTree_t {
     return std::monostate();
   }
 
+  std::monostate fillNotToExploredNode() {
+    return cursor->fillNotToExploreNode();
+  }
+
   bool worth_to_create_snapshot() {
     if (!ENABLE_COST_MODEL) {
       return REUSE_SNAPSHOT;
@@ -911,8 +941,8 @@ class ExploreTree_t {
     auto parent_cost =
         cursor->parent ? cursor->parent->min_cost_of_reaching_here() : 0;
     auto exec_from_parent_cost = reach_parent_cost + cursor->instr_cost;
-    GENSYM_INFO("The score of snapshot tendency: " + std::to_string(exec_from_parent_cost -
-                snapshot_cost));
+    GENSYM_INFO("The score of snapshot tendency: " +
+                std::to_string(exec_from_parent_cost - snapshot_cost));
     return snapshot_cost <= exec_from_parent_cost;
   }
 
@@ -924,6 +954,7 @@ class ExploreTree_t {
         if_else_node != nullptr &&
         "Can't move cursor when the branch node is not initialized correctly!");
     int cost_from_parent = CostManager.dump_instr_cost();
+    cursor->instr_cost = cost_from_parent;
     if (branch) {
       true_branch_cov_map[if_else_node->id] = true;
       if (worth_to_create_snapshot()) {
@@ -947,6 +978,26 @@ class ExploreTree_t {
     return std::monostate();
   }
 
+  std::monostate moveCursorNoControl(bool branch) {
+    Profile.step(StepProfileKind::CURSOR_MOVE);
+    assert(cursor != nullptr);
+    auto if_else_node = dynamic_cast<IfElseNode *>(cursor->node.get());
+    assert(
+        if_else_node != nullptr &&
+        "Can't move cursor when the branch node is not initialized correctly!");
+    int cost_from_parent = CostManager.dump_instr_cost();
+    cursor->instr_cost = cost_from_parent;
+    if (branch) {
+      true_branch_cov_map[if_else_node->id] = true;
+      if_else_node->false_branch->fillNotToExploreNode();
+      cursor = if_else_node->true_branch.get();
+    } else {
+      assert(false &&
+             "moveCursorNoControl should not be used for false branch");
+    }
+    return std::monostate();
+  }
+
   std::monostate print() {
     std::cout << root->node->to_string() << std::endl;
     return std::monostate();
@@ -966,6 +1017,45 @@ class ExploreTree_t {
     return std::monostate();
   }
 
+  std::monostate print_overall_result() {
+    // Print how many paths have been explored, how many paths are unreachable,
+    // how many paths are failed, how many paths are finished successfully
+    int unexplored_count = 0;
+    int finished_count = 0;
+    int failed_count = 0;
+    int not_to_explore_count = 0;
+    int unreachable_count = 0;
+    std::function<void(NodeBox *)> dfs = [&](NodeBox *node) {
+      if (auto if_else_node = dynamic_cast<IfElseNode *>(node->node.get())) {
+        dfs(if_else_node->true_branch.get());
+        dfs(if_else_node->false_branch.get());
+      } else if (dynamic_cast<UnExploredNode *>(node->node.get())) {
+        unexplored_count += 1;
+      } else if (dynamic_cast<Finished *>(node->node.get())) {
+        finished_count += 1;
+      } else if (dynamic_cast<Failed *>(node->node.get())) {
+        failed_count += 1;
+      } else if (dynamic_cast<Unreachable *>(node->node.get())) {
+        unreachable_count += 1;
+      } else if (dynamic_cast<SnapshotNode *>(node->node.get())) {
+        // Snapshot node is considered unexplored
+        unexplored_count += 1;
+      } else if (dynamic_cast<NotToExploreNode *>(node->node.get())) {
+        not_to_explore_count += 1;
+      } else {
+        throw std::runtime_error("Unknown node type in explore tree");
+      }
+    };
+    dfs(root.get());
+    std::cout << "Explore Tree Overall Result:" << std::endl;
+    std::cout << "  Unexplored paths: " << unexplored_count << std::endl;
+    std::cout << "  Finished paths: " << finished_count << std::endl;
+    std::cout << "  Failed paths: " << failed_count << std::endl;
+    std::cout << "  Unreachable paths: " << unreachable_count << std::endl;
+    std::cout << "  NotToExplore paths: " << not_to_explore_count << std::endl;
+    return std::monostate();
+  }
+
   NodeBox *pick_unexplored() {
     // Pick an unexplored node from the tree
     // For now, we just iterate through the tree and return the first unexplored

src/main/scala/wasm/StagedConcolicMiniWasm.scala

@@ -31,6 +31,12 @@ object Counter {
     dict.clear()
   }
 
+  def getId(): Int = {
+    val id = currentId
+    currentId += 1
+    id
+  }
+
   def getId(wir: WIR, nth: Int = 0): Int = {
     if (dict.contains((wir, nth))) {
       dict((wir, nth))
@@ -547,8 +553,28 @@ trait StagedWasmEvaluator extends SAIOps {
       case Import("i32", "symbolic", _) =>
         evalSymbolic(NumType(I32Type), rest, kont, mkont, trail)(ctx)
       case Import("i32", "sym_assume", _) =>
-        // TODO: implement sym_assume
-        eval(rest, kont, mkont, trail)(ctx.pop()._2)
+        // symbolic assume is just like an if else that only has one branch, while another
+        // is marked as not-to-explore
+        val (condTy, newCtx) = ctx.pop()
+        Predef.assert(condTy == NumType(I32Type), s"sym_assume only supports i32 condition, get $condTy")
+        val cond = Stack.popC(condTy)
+        val symCond = Stack.popS(condTy)
+        val id = Counter.getId()
+        ExploreTree.fillWithIfElse(symCond.s, id)
+        def thnK: Rep[Cont[Unit]] = topFun((mk: Rep[MCont[Unit]]) => {
+          info(s"Successfully assumed condition at $id")
+          eval(rest, kont, mk, trail)(newCtx)
+        })
+        if (cond.toInt != 0) {
+          ExploreTree.moveCursor(true)
+          eval(rest, kont, mkont, trail)(newCtx)
+        } else {
+          val control = makeControl(thnK, mkont)
+          ExploreTree.moveCursor(false, control)
+          // just stop the execution at here
+          ExploreTree.fillWithNotToExplore()
+        }
+        ()
       case Import("i32", "sym_assert", _) =>
         // TODO: implement sym_assert
         eval(rest, kont, mkont, trail)(ctx.pop()._2)
@@ -637,6 +663,8 @@ trait StagedWasmEvaluator extends SAIOps {
 
   def evalTop(mkont: Rep[MCont[Unit]], main: Option[String]): Rep[Unit] = {
     Counter.reset()
+    Predef.println("[DEBUG]" + module)
+    Predef.println("[DEBUG] module.defs: " + module.defs)
     val funBody: FuncBodyDef = main match {
       case Some(func_name) =>
         module.defs.flatMap({
@@ -889,6 +917,10 @@ trait StagedWasmEvaluator extends SAIOps {
       "tree-fill-if-else".reflectCtrlWith[Unit](s, id)
     }
 
+    def fillWithNotToExplore(): Rep[Unit] = {
+      "tree-fill-not-to-explore".reflectCtrlWith[Unit]()
+    }
+
     def fillWithFinished(): Rep[Unit] = {
       "tree-fill-finished".reflectCtrlWith[Unit]()
     }
@@ -898,6 +930,11 @@ trait StagedWasmEvaluator extends SAIOps {
       "tree-move-cursor".reflectCtrlWith[Unit](branch, control)
     }
 
+    def moveCursor(branch: Boolean): Rep[Unit] = {
+      // when moving cursor from to an unexplored node, we need to change the reuse state
+      "tree-move-cursor-no-control".reflectCtrlWith[Unit](branch)
+    }
+
     def print(): Rep[Unit] = {
       "tree-print".reflectCtrlWith[Unit]()
     }
@@ -1543,10 +1580,14 @@ trait StagedWasmCppGen extends CGenBase with CppSAICodeGenBase {
       emit("GENSYM_ASSERT("); shallow(cond); emit(")")
     case Node(_, "tree-fill-if-else", List(sym, id), _) =>
       emit("ExploreTree.fillIfElseNode("); shallow(sym); emit(", "); emit(id.toString); emit(")")
+    case Node(_, "tree-fill-not-to-explore", List(), _) =>
+      emit("ExploreTree.fillNotToExploredNode()")
     case Node(_, "tree-fill-finished", List(), _) =>
       emit("ExploreTree.fillFinishedNode()")
     case Node(_, "tree-move-cursor", List(b, snapshot), _) =>
       emit("ExploreTree.moveCursor("); shallow(b); emit(", "); shallow(snapshot); emit(")")
+    case Node(_, "tree-move-cursor-no-control", List(b), _) =>
+      emit("ExploreTree.moveCursorNoControl("); shallow(b); emit(")")
     case Node(_, "add-instr-cost", List(n), _) =>
     emit("CostManager.add_instr_cost("); shallow(n); emit(")")
     case Node(_, "tree-print", List(), _) =>