fill snapshot into SnapshotNode

Author: butterunderflow

headers/wasm/symbolic_rt.hpp

@@ -228,14 +228,17 @@ class SymFrames_t {
 // A snapshot of the symbolic state and execution context (control)
 class Snapshot_t {
 public:
-  explicit Snapshot_t();
+  explicit Snapshot_t(Cont_t cont, MCont_t mcont);
 
   SymStack_t get_stack() const { return stack; }
   SymFrames_t get_frames() const { return frames; }
 
 private:
   SymStack_t stack;
   SymFrames_t frames;
+  // The continuation at the snapshot point
+  Cont_t cont;
+  MCont_t mcont;
 };
 
 inline void SymStack_t::reuse(Snapshot_t snapshot) {
@@ -273,7 +276,7 @@ struct NodeBox {
   std::monostate fillFinishedNode();
   std::monostate fillFailedNode();
   std::monostate fillUnreachableNode();
-  std::monostate fillSnapshotNode();
+  std::monostate fillSnapshotNode(Snapshot_t snapshot);
   bool isUnexplored() const;
   std::vector<SymVal> collect_path_conds();
 };
@@ -386,7 +389,7 @@ struct UnExploredNode : Node {
 };
 
 struct SnapshotNode : Node {
-  SnapshotNode() {}
+  SnapshotNode(Snapshot_t snapshot) : snapshot(snapshot) {}
   std::string to_string() override { return "SnapshotNode"; }
 
 protected:
@@ -399,6 +402,9 @@ struct SnapshotNode : Node {
       graphviz_edge(os, parent_dot_id, current_node_dot_id, edge_label);
     }
   }
+
+private:
+  Snapshot_t snapshot;
 };
 
 struct Finished : Node {
@@ -465,9 +471,9 @@ inline std::monostate NodeBox::fillIfElseNode(SymVal cond) {
   return std::monostate();
 }
 
-inline std::monostate NodeBox::fillSnapshotNode() {
+inline std::monostate NodeBox::fillSnapshotNode(Snapshot_t snapshot) {
   if (this->isUnexplored()) {
-    node = std::make_unique<SnapshotNode>();
+    node = std::make_unique<SnapshotNode>(snapshot);
   }
   return std::monostate();
 }
@@ -545,12 +551,11 @@ class Reuse_t {
 
 static Reuse_t Reuse;
 
-inline Snapshot_t::Snapshot_t() : stack(SymStack), frames(SymFrames) {
+inline Snapshot_t::Snapshot_t(Cont_t cont, MCont_t mcont)
+    : stack(SymStack), frames(SymFrames), cont(cont), mcont(mcont) {
 #ifdef DEBUG
   std::cout << "Creating snapshot of size " << stack.size() << std::endl;
 #endif
-  assert(!Reuse.is_reusing() &&
-         "Creating snapshot while reusing the symbolic stack");
 }
 
 class ExploreTree_t {
@@ -576,17 +581,17 @@ class ExploreTree_t {
     return cursor->fillIfElseNode(cond);
   }
 
-  std::monostate moveCursor(bool branch) {
+  std::monostate moveCursor(bool branch, Snapshot_t snapshot) {
     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!");
     if (branch) {
-      if_else_node->false_branch->fillSnapshotNode();
+      if_else_node->false_branch->fillSnapshotNode(snapshot);
       cursor = if_else_node->true_branch.get();
     } else {
-      if_else_node->true_branch->fillSnapshotNode();
+      if_else_node->true_branch->fillSnapshotNode(snapshot);
       cursor = if_else_node->false_branch.get();
     }
 

src/main/scala/wasm/StagedConcolicMiniWasm.scala

@@ -165,9 +165,9 @@ trait StagedWasmEvaluator extends SAIOps {
   trait Snapshot
 
   // Create a snapshot of the symbolic execution, we should ensure that current symstack is in use
-  // We don't need to store the control information, since the control is totally decided by concrete states
-  def makeSnapshot(): Rep[Snapshot] = {
-    "snapshot-make".reflectCtrlWith[Snapshot]()
+  // We need to store the control information, so we can resume the execution later
+  def makeSnapshot(kont: Rep[Cont[Unit]], mkont: Rep[MCont[Unit]]): Rep[Snapshot] = {
+    "snapshot-make".reflectCtrlWith[Snapshot](kont, mkont)
   }
 
   def eval(insts: List[Instr],
@@ -365,14 +365,23 @@ trait StagedWasmEvaluator extends SAIOps {
           val newRestCtx = restCtx.shift(offset, funcTy.out.size)
           eval(rest, kont, mk, trail)(newRestCtx)
         })
-        // TODO: put the cond.s to path condition
         ExploreTree.fillWithIfElse(symCond.s)
+        def thnK: Rep[Cont[Unit]] = topFun((mk: Rep[MCont[Unit]]) => {
+          info("Entering the true branch of the if")
+          eval(thn, restK _, mk, restK _ :: trail)(newCtx)
+        })
+        def elsK: Rep[Cont[Unit]] = topFun((mk: Rep[MCont[Unit]]) => {
+          info("Entering the false branch of the if")
+          eval(els, restK _, mk, restK _ :: trail)(newCtx)
+        })
         if (cond.toInt != 0) {
-          ExploreTree.moveCursor(true)
-          eval(thn, restK _, mkont, restK _ :: trail)(newCtx)
+          val snapshot = makeSnapshot(elsK, mkont)
+          ExploreTree.moveCursor(true, snapshot)
+          thnK(mkont)
         } else {
-          ExploreTree.moveCursor(false)
-          eval(els, restK _, mkont, restK _ :: trail)(newCtx)
+          val snapshot = makeSnapshot(thnK, mkont)
+          ExploreTree.moveCursor(false, snapshot)
+          elsK(mkont)
         }
         ()
       case Br(label) =>
@@ -384,40 +393,63 @@ trait StagedWasmEvaluator extends SAIOps {
         val symCond = Stack.popS(ty)
         info(s"The br_if(${label})'s condition is ", cond.toInt)
         ExploreTree.fillWithIfElse(symCond.s)
+        def thnK: Rep[Cont[Unit]] = topFun((mk: Rep[MCont[Unit]]) => {
+          trail(label)(newCtx)(mk)
+        })
+        def elsK: Rep[Cont[Unit]] = topFun((mk: Rep[MCont[Unit]]) => {
+          eval(rest, kont, mk, trail)(newCtx)
+        })
         if (cond.toInt != 0) {
           info(s"Jump to $label")
-          ExploreTree.moveCursor(true)
-          trail(label)(newCtx)(mkont)
+          val snapshot = makeSnapshot(elsK, mkont)
+          ExploreTree.moveCursor(true, snapshot)
+          thnK(mkont)
         } else {
           info(s"Continue")
-          ExploreTree.moveCursor(false)
-          eval(rest, kont, mkont, trail)(newCtx)
+          val snapshot = makeSnapshot(thnK, mkont)
+          ExploreTree.moveCursor(false, snapshot)
+          elsK(mkont)
         }
         ()
       case BrTable(labels, default) =>
         val (ty, newCtx) = ctx.pop()
-        val label = Stack.popC(ty)
-        val labelSym = Stack.popS(ty)
-        def aux(choices: List[Int], idx: Int): Rep[Unit] = {
-          if (choices.isEmpty) trail(default)(newCtx)(mkont)
-          else {
+        def aux(choices: List[Int], idx: Int, mkont: Rep[MCont[Unit]]): Rep[Unit] = {
+          if (choices.isEmpty) {
+            Stack.popC(ty)
+            Stack.popS(ty)
+            trail(default)(newCtx)(mkont)
+          } else {
+            val label = Stack.peekC(ty)
+            val labelSym = Stack.peekS(ty)
             val cond = (label - toStagedNum(I32V(idx))).isZero()
             val condSym = (labelSym - toStagedSymbolicNum(I32V(idx))).isZero()
             ExploreTree.fillWithIfElse(condSym.s)
             // When moving the cursor to a branch, we mark another branch as
             // snapshotNode (this is done by moveCursor's runtime implementation)
             // TODO: store snapshot into this snapshot node
+            def thnK: Rep[Cont[Unit]] = topFun((mk: Rep[MCont[Unit]]) => {
+              info("Entering the true branch of the br_table")
+              Stack.popC(ty)
+              Stack.popS(ty)
+              trail(choices.head)(newCtx)(mk)
+            })
+            def elsK: Rep[Cont[Unit]] = topFun((mk: Rep[MCont[Unit]]) => {
+              info("Entering the false branch of the br_table")
+              aux(choices.tail, idx + 1, mk)
+            })
             if (cond.toInt != 0) {
-              ExploreTree.moveCursor(true)
-              trail(choices.head)(newCtx)(mkont)
+              val snapshot = makeSnapshot(elsK, mkont)
+              ExploreTree.moveCursor(true, snapshot)
+              thnK(mkont)
             }
             else {
-              ExploreTree.moveCursor(false)
-              aux(choices.tail, idx + 1)
+              val snapshot = makeSnapshot(thnK, mkont)
+              ExploreTree.moveCursor(false, snapshot)
+              elsK(mkont)
             }
           }
         }
-        aux(labels, 0)
+        aux(labels, 0, mkont)
       case Return        => trail.last(ctx)(mkont)
       case Call(f)       => evalCall(rest, kont, mkont, trail, f, false)
       case ReturnCall(f) => evalCall(rest, kont, mkont, trail, f, true)
@@ -888,9 +920,9 @@ trait StagedWasmEvaluator extends SAIOps {
       "tree-fill-finished".reflectCtrlWith[Unit]()
     }
 
-    def moveCursor(branch: Boolean): Rep[Unit] = {
+    def moveCursor(branch: Boolean, snapshot: Rep[Snapshot]): Rep[Unit] = {
       // when moving cursor from to an unexplored node, we need to change the reuse state
-      "tree-move-cursor".reflectCtrlWith[Unit](branch)
+      "tree-move-cursor".reflectCtrlWith[Unit](branch, snapshot)
     }
 
     def print(): Rep[Unit] = {
@@ -1323,8 +1355,8 @@ trait StagedWasmCppGen extends CGenBase with CppSAICodeGenBase {
       emit("Stack.pop()")
     case Node(_, "sym-stack-pop", _, _) =>
       emit("SymStack.pop()")
-    case Node(_, "snapshot-make", _, _) => 
-      emit("Snapshot_t()")
+    case Node(_, "snapshot-make", List(k, mk), _) =>
+      emit("Snapshot_t("); shallow(k); emit(", "); shallow(mk); emit(")")
     case Node(_, "frame-pop", List(i), _) =>
       emit("Frames.popFrame("); shallow(i); emit(")")
     case Node(_, "sym-frame-pop", List(i), _) =>
@@ -1420,8 +1452,8 @@ trait StagedWasmCppGen extends CGenBase with CppSAICodeGenBase {
       emit("ExploreTree.fillIfElseNode("); shallow(sym); emit(")")
     case Node(_, "tree-fill-finished", List(), _) =>
       emit("ExploreTree.fillFinishedNode()")
-    case Node(_, "tree-move-cursor", List(b), _) =>
-      emit("ExploreTree.moveCursor("); shallow(b); emit(")")
+    case Node(_, "tree-move-cursor", List(b, snapshot), _) =>
+      emit("ExploreTree.moveCursor("); shallow(b); emit(", "); shallow(snapshot); emit(")")
     case Node(_, "tree-print", List(), _) =>
       emit("ExploreTree.print()")
     case Node(_, "tree-dump-graphviz", List(f), _) =>