From 12b04c63653f4e928dbd8d0c0637694f1c888b77 Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Sat, 15 Oct 2022 16:33:28 -0400
Subject: [PATCH 01/15] simplify cps

---
 .../scala/gensym/engines/ImpCPSEngine.scala   | 124 +++++++++---------
 src/test/scala/gensym/TestGS.scala            |   4 +-
 2 files changed, 61 insertions(+), 67 deletions(-)

diff --git a/src/main/scala/gensym/engines/ImpCPSEngine.scala b/src/main/scala/gensym/engines/ImpCPSEngine.scala
index 2d06efe1..31ea6b88 100644
--- a/src/main/scala/gensym/engines/ImpCPSEngine.scala
+++ b/src/main/scala/gensym/engines/ImpCPSEngine.scala
@@ -108,81 +108,81 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
   def evalFloatOp2(op: String, lhs: LLVMValue, rhs: LLVMValue, ty: LLVMType, ss: Rep[SS])(implicit ctx: Ctx): Rep[Value] =
     FloatOp2(op, eval(lhs, ty, ss), eval(rhs, ty, ss))
 
-  def execValueInst(inst: ValueInstruction, ss: Rep[SS], k: (Rep[SS], Rep[Value], Rep[Cont]) => Rep[Unit])(implicit ctx: Ctx, kk: Rep[Cont]): Rep[Unit] = {
+  def execValueInst(inst: ValueInstruction, ss: Rep[SS])(k: (Rep[SS], Rep[Value]) => Rep[Unit])(implicit ctx: Ctx): Rep[Unit] = {
     inst match {
       // Memory Access Instructions
       case AllocaInst(ty, align) =>
         val typeSize = ty.size
         val sz = ss.stackSize
         ss.allocStack(typeSize, align.n)
-        k(ss, LocV(sz, LocV.kStack, typeSize.toLong), kk)
+        k(ss, LocV(sz, LocV.kStack, typeSize.toLong))
       case LoadInst(valTy, ptrTy, value, align) =>
         val isStruct = getRealType(valTy) match {
           case Struct(types) => 1
           case _ => 0
         }
         val v = eval(value, ptrTy, ss)
-        k(ss, ss.lookup(v, valTy.size, isStruct), kk)
+        k(ss, ss.lookup(v, valTy.size, isStruct))
       case GetElemPtrInst(_, baseType, ptrType, ptrValue, typedValues) =>
         val vs = typedValues.map(tv => eval(tv.value, tv.ty, ss))
         val offset = calculateOffset(ptrType, vs)
         val v = eval(ptrValue, ptrType, ss).asRepOf[LocV] + offset
-        k(ss, v, kk)
+        k(ss, v)
       // Arith Unary Operations
-      case FNegInst(ty, op) => k(ss, evalFloatOp2("fsub", FloatConst(-0.0), op, ty, ss), kk)
+      case FNegInst(ty, op) => k(ss, evalFloatOp2("fsub", FloatConst(-0.0), op, ty, ss))
       // Arith Binary Operations
-      case AddInst(ty, lhs, rhs, _) => k(ss, evalIntOp2("add", lhs, rhs, ty, ss), kk)
-      case SubInst(ty, lhs, rhs, _) => k(ss, evalIntOp2("sub", lhs, rhs, ty, ss), kk)
-      case MulInst(ty, lhs, rhs, _) => k(ss, evalIntOp2("mul", lhs, rhs, ty, ss), kk)
-      case SDivInst(ty, lhs, rhs) => k(ss, evalIntOp2("sdiv", lhs, rhs, ty, ss), kk)
-      case UDivInst(ty, lhs, rhs) => k(ss, evalIntOp2("udiv", lhs, rhs, ty, ss), kk)
-      case FAddInst(ty, lhs, rhs) => k(ss, evalFloatOp2("fadd", lhs, rhs, ty, ss), kk)
-      case FSubInst(ty, lhs, rhs) => k(ss, evalFloatOp2("fsub", lhs, rhs, ty, ss), kk)
-      case FMulInst(ty, lhs, rhs) => k(ss, evalFloatOp2("fmul", lhs, rhs, ty, ss), kk)
-      case FDivInst(ty, lhs, rhs) => k(ss, evalFloatOp2("fdiv", lhs, rhs, ty, ss), kk)
-      case URemInst(ty, lhs, rhs) => k(ss, evalIntOp2("urem", lhs, rhs, ty, ss), kk)
-      case SRemInst(ty, lhs, rhs) => k(ss, evalIntOp2("srem", lhs, rhs, ty, ss), kk)
+      case AddInst(ty, lhs, rhs, _) => k(ss, evalIntOp2("add", lhs, rhs, ty, ss))
+      case SubInst(ty, lhs, rhs, _) => k(ss, evalIntOp2("sub", lhs, rhs, ty, ss))
+      case MulInst(ty, lhs, rhs, _) => k(ss, evalIntOp2("mul", lhs, rhs, ty, ss))
+      case SDivInst(ty, lhs, rhs) => k(ss, evalIntOp2("sdiv", lhs, rhs, ty, ss))
+      case UDivInst(ty, lhs, rhs) => k(ss, evalIntOp2("udiv", lhs, rhs, ty, ss))
+      case FAddInst(ty, lhs, rhs) => k(ss, evalFloatOp2("fadd", lhs, rhs, ty, ss))
+      case FSubInst(ty, lhs, rhs) => k(ss, evalFloatOp2("fsub", lhs, rhs, ty, ss))
+      case FMulInst(ty, lhs, rhs) => k(ss, evalFloatOp2("fmul", lhs, rhs, ty, ss))
+      case FDivInst(ty, lhs, rhs) => k(ss, evalFloatOp2("fdiv", lhs, rhs, ty, ss))
+      case URemInst(ty, lhs, rhs) => k(ss, evalIntOp2("urem", lhs, rhs, ty, ss))
+      case SRemInst(ty, lhs, rhs) => k(ss, evalIntOp2("srem", lhs, rhs, ty, ss))
       // Bitwise Operations
-      case ShlInst(ty, lhs, rhs) => k(ss, evalIntOp2("shl", lhs, rhs, ty, ss), kk)
-      case LshrInst(ty, lhs, rhs) => k(ss, evalIntOp2("lshr", lhs, rhs, ty, ss), kk)
-      case AshrInst(ty, lhs, rhs) => k(ss, evalIntOp2("ashr", lhs, rhs, ty, ss), kk)
-      case AndInst(ty, lhs, rhs) => k(ss, evalIntOp2("and", lhs, rhs, ty, ss), kk)
-      case OrInst(ty, lhs, rhs) => k(ss, evalIntOp2("or", lhs, rhs, ty, ss), kk)
-      case XorInst(ty, lhs, rhs) => k(ss, evalIntOp2("xor", lhs, rhs, ty, ss), kk)
+      case ShlInst(ty, lhs, rhs) => k(ss, evalIntOp2("shl", lhs, rhs, ty, ss))
+      case LshrInst(ty, lhs, rhs) => k(ss, evalIntOp2("lshr", lhs, rhs, ty, ss))
+      case AshrInst(ty, lhs, rhs) => k(ss, evalIntOp2("ashr", lhs, rhs, ty, ss))
+      case AndInst(ty, lhs, rhs) => k(ss, evalIntOp2("and", lhs, rhs, ty, ss))
+      case OrInst(ty, lhs, rhs) => k(ss, evalIntOp2("or", lhs, rhs, ty, ss))
+      case XorInst(ty, lhs, rhs) => k(ss, evalIntOp2("xor", lhs, rhs, ty, ss))
       // Conversion Operations
       case ZExtInst(from, value, IntType(size)) =>
-        k(ss, eval(value, from, ss).zExt(size), kk)
+        k(ss, eval(value, from, ss).zExt(size))
       case SExtInst(from, value, IntType(size)) =>
-        k(ss, eval(value, from, ss).sExt(size), kk)
+        k(ss, eval(value, from, ss).sExt(size))
       case TruncInst(from@IntType(fromSz), value, IntType(toSz)) =>
-        k(ss, eval(value, from, ss).trunc(fromSz, toSz), kk)
+        k(ss, eval(value, from, ss).trunc(fromSz, toSz))
       case FpExtInst(from, value, to) =>
-        k(ss, eval(value, from, ss), kk)
+        k(ss, eval(value, from, ss))
       case FpToUIInst(from, value, IntType(size)) =>
-        k(ss, eval(value, from, ss).fromFloatToUInt(size), kk)
+        k(ss, eval(value, from, ss).fromFloatToUInt(size))
       case FpToSIInst(from, value, IntType(size)) =>
-        k(ss, eval(value, from, ss).fromFloatToSInt(size), kk)
+        k(ss, eval(value, from, ss).fromFloatToSInt(size))
       case UiToFPInst(from, value, to) =>
-        k(ss, eval(value, from, ss).fromUIntToFloat, kk)
+        k(ss, eval(value, from, ss).fromUIntToFloat)
       case SiToFPInst(from, value, to) =>
-        k(ss, eval(value, from, ss).fromSIntToFloat, kk)
+        k(ss, eval(value, from, ss).fromSIntToFloat)
       case PtrToIntInst(from, value, to) =>
         val v = eval(value, from, ss)
         val toSize = to.asInstanceOf[IntType].size
-        k(ss, if (ARCH_WORD_SIZE == toSize) v else v.trunc(ARCH_WORD_SIZE, toSize), kk)
+        k(ss, if (ARCH_WORD_SIZE == toSize) v else v.trunc(ARCH_WORD_SIZE, toSize))
       case IntToPtrInst(from, value, to) =>
-        k(ss, eval(value, from, ss), kk)
-      case BitCastInst(from, value, to) => k(ss, eval(value, to, ss), kk)
+        k(ss, eval(value, from, ss))
+      case BitCastInst(from, value, to) => k(ss, eval(value, to, ss))
       // Aggregate Operations
       case ExtractValueInst(ty, struct, indices) =>
         val idxList = indices.asInstanceOf[List[IntConst]].map(x => x.n)
         val idx = calculateOffsetStatic(ty, idxList)
         // v is expected to be StructV in backend
         val v = eval(struct, ty, ss)
-        k(ss, v.structAt(idx), kk)
+        k(ss, v.structAt(idx))
       // Arithm comparison operations
-      case FCmpInst(pred, ty, lhs, rhs) => k(ss, evalFloatOp2(pred.op, lhs, rhs, ty, ss), kk)
-      case ICmpInst(pred, ty, lhs, rhs) => k(ss, evalIntOp2(pred.op, lhs, rhs, ty, ss), kk)
+      case FCmpInst(pred, ty, lhs, rhs) => k(ss, evalFloatOp2(pred.op, lhs, rhs, ty, ss))
+      case ICmpInst(pred, ty, lhs, rhs) => k(ss, evalIntOp2(pred.op, lhs, rhs, ty, ss))
       // Other operations
       case CallInst(ty, f, args) =>
         val argValues: List[LLVMValue] = extractValues(args)
@@ -190,8 +190,8 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
         val fv = eval(f, VoidType, ss, Some(argTypes))
         val vs = argValues.zip(argTypes).map { case (v, t) => eval(v, t, ss) }
         val stackSize = ss.stackSize
-        ss.push(kk)
-        def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = k(s, v, s.pop(stackSize))
+        ss.push
+        def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(stackSize); k(s, v) }
         fv[Ref](ss, List(vs: _*), ContOpt(fK))
       case PhiInst(ty, incs) =>
         def selectValue(bb: Rep[BlockLabel], vs: List[() => Rep[Value]], labels: List[BlockLabel]): Rep[Value] = {
@@ -201,23 +201,23 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
         val incsValues: List[LLVMValue] = incs.map(_.value)
         val incsLabels: List[BlockLabel] = incs.map(i => Counter.block.get(ctx.withBlock(i.label)))
         val vs = incsValues.map(v => () => eval(v, ty, ss))
-        k(ss, selectValue(ss.incomingBlock, vs, incsLabels), kk)
+        k(ss, selectValue(ss.incomingBlock, vs, incsLabels))
       case SelectInst(cndTy, cndVal, thnTy, thnVal, elsTy, elsVal) if Config.iteSelect =>
-        k(ss, ITE(eval(cndVal, cndTy, ss), eval(thnVal, thnTy, ss), eval(elsVal, elsTy, ss)), kk)
+        k(ss, ITE(eval(cndVal, cndTy, ss), eval(thnVal, thnTy, ss), eval(elsVal, elsTy, ss)))
       case SelectInst(cndTy, cndVal, thnTy, thnVal, elsTy, elsVal) =>
         val cnd = eval(cndVal, cndTy, ss)
         val repK = fun(k)
         if (cnd.isConc) {
-          if (cnd.int == 1) repK(ss, eval(thnVal, thnTy, ss), kk)
-          else repK(ss, eval(elsVal, elsTy, ss), kk)
+          if (cnd.int == 1) repK(ss, eval(thnVal, thnTy, ss))
+          else repK(ss, eval(elsVal, elsTy, ss))
         } else {
           // TODO: check cond via solver
           val s1 = ss.fork
           ss.addPC(cnd)
-          repK(ss, eval(thnVal, thnTy, ss), kk)
+          repK(ss, eval(thnVal, thnTy, ss))
           s1.addPC(!cnd)
           Coverage.incPath(1)
-          repK(s1, eval(elsVal, elsTy, s1), kk)
+          repK(s1, eval(elsVal, elsTy, s1))
         }
     }
   }
@@ -308,30 +308,28 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
     }
   }
 
-  def execInst(inst: Instruction, ss: Rep[SS], k: (Rep[SS], Rep[Cont]) => Rep[Unit])(implicit ctx: Ctx, kk: Rep[Cont]): Rep[Unit] = {
+  def execInst(inst: Instruction, ss: Rep[SS], k: Rep[SS] => Rep[Unit])(implicit ctx: Ctx): Rep[Unit] = {
     inst match {
       case AssignInst(x, valInst) =>
-        execValueInst(valInst, ss, { case (s, v, kk) =>
+        execValueInst(valInst, ss) { (s, v) =>
           s.assign(x, v)
-          k(s, kk)
-        })
+          k(s)
+        }
       case StoreInst(ty1, val1, ty2, val2, align) =>
         val v1 = eval(val1, ty1, ss)
         val v2 = eval(val2, ty2, ss)
         ss.update(v2, v1, ty1.size)
-        k(ss, kk)
+        k(ss)
       case CallInst(ty, f, args) =>
-        val argValues: List[LLVMValue] = extractValues(args)
-        val argTypes: List[LLVMType] = extractTypes(args)
-        val fv = eval(f, VoidType, ss, Some(argTypes))
-        val vs = argValues.zip(argTypes).map { case (v, t) => eval(v, t, ss) }
-        val stackSize = ss.stackSize
-        ss.push(kk)
-        def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = k(s, s.pop(stackSize))
-        fv[Ref](ss, List(vs: _*), ContOpt(fK))
+        execValueInst(CallInst(ty, f, args), ss) { (s, v) => k(s) }
     }
   }
 
+  def execInsts(insts: List[Instruction], s: Rep[SS])(k: Rep[SS] => Rep[Unit])(implicit ctx: Ctx): Rep[Unit] = insts match {
+    case Nil => k(s)
+    case i::inst => execInst(i, s, s => execInsts(inst, s)(k))
+  }
+
   def execBlock(funName: String, label: String, s: Rep[SS], k: Rep[Cont]): Rep[Unit] =
     execBlock(funName, findBlock(funName, label).get, s, k)
 
@@ -341,15 +339,11 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
   }
 
   def execBlockEager(block: BB, s: Rep[SS], k: Rep[Cont])(implicit ctx: Ctx): Rep[Unit] = {
-    def runInst(insts: List[Instruction], t: Terminator, s: Rep[SS], k: Rep[Cont]): Rep[Unit] =
-      insts match {
-        case Nil =>
-          Coverage.incInst(block.ins.size+1)
-          execTerm(t, k)(s, ctx)
-        case i::inst => execInst(i, s, (s1, k1) => runInst(inst, t, s1, k1))(ctx, k)
-      }
     s.coverBlock(ctx)
-    runInst(block.ins, block.term, s, k)
+    execInsts(block.ins, s) { s =>
+      Coverage.incInst(block.ins.size+1)
+      execTerm(block.term, k)(s, ctx)
+    }
   }
 
   override def repBlockFun(b: BB)(implicit ctx: Ctx): BFTy = {
diff --git a/src/test/scala/gensym/TestGS.scala b/src/test/scala/gensym/TestGS.scala
index 01e8f702..8e309451 100644
--- a/src/test/scala/gensym/TestGS.scala
+++ b/src/test/scala/gensym/TestGS.scala
@@ -143,9 +143,9 @@ class Coreutils extends TestGS {
 class Playground extends TestGS {
   import gensym.llvm.parser.Parser._
   Config.enableOpt
-  val gs = new ImpGS
+  val gs = new ImpCPSGS
 
-  //testGS(gs, TestPrg(mergesort, "mergeSortTest1", "@main", noArg, noOpt, nPath(720)))
+  testGS(gs, TestPrg(mergesort, "mergeSortTest", "@main", noArg, noOpt, nPath(720)))
   //testGS(new PureCPSGS, TestPrg(arrayFlow, "arrayFlow", "@main", noArg, noOpt, nPath(15)++status(0)))
   //testGS(new ImpCPSGS, TestPrg(arrayFlow, "arrayFlow2", "@main", noArg, noOpt, nPath(15)++status(0)))
 

From 1150665e4a6a8e296d689cb483be90cfeea4746f Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Sat, 15 Oct 2022 17:46:08 -0400
Subject: [PATCH 02/15] simplify call inst

---
 .../scala/gensym/engines/ImpCPSEngine.scala     |  6 ++++--
 src/main/scala/gensym/engines/ImpEngine.scala   | 17 ++---------------
 .../scala/gensym/engines/PureCPSEngine.scala    |  9 ++-------
 src/main/scala/gensym/engines/PureEngine.scala  | 12 +-----------
 4 files changed, 9 insertions(+), 35 deletions(-)

diff --git a/src/main/scala/gensym/engines/ImpCPSEngine.scala b/src/main/scala/gensym/engines/ImpCPSEngine.scala
index 31ea6b88..a15d18be 100644
--- a/src/main/scala/gensym/engines/ImpCPSEngine.scala
+++ b/src/main/scala/gensym/engines/ImpCPSEngine.scala
@@ -190,7 +190,9 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
         val fv = eval(f, VoidType, ss, Some(argTypes))
         val vs = argValues.zip(argTypes).map { case (v, t) => eval(v, t, ss) }
         val stackSize = ss.stackSize
+        // ss.push(ss.stackSize, k)
         ss.push
+        //def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(s, v) }
         def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(stackSize); k(s, v) }
         fv[Ref](ss, List(vs: _*), ContOpt(fK))
       case PhiInst(ty, incs) =>
@@ -320,8 +322,8 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
         val v2 = eval(val2, ty2, ss)
         ss.update(v2, v1, ty1.size)
         k(ss)
-      case CallInst(ty, f, args) =>
-        execValueInst(CallInst(ty, f, args), ss) { (s, v) => k(s) }
+      case call@CallInst(ty, f, args) =>
+        execValueInst(call, ss) { (s, v) => k(s) }
     }
   }
 
diff --git a/src/main/scala/gensym/engines/ImpEngine.scala b/src/main/scala/gensym/engines/ImpEngine.scala
index cfd3633b..c20ff6ef 100644
--- a/src/main/scala/gensym/engines/ImpEngine.scala
+++ b/src/main/scala/gensym/engines/ImpEngine.scala
@@ -310,21 +310,8 @@ trait ImpGSEngine extends ImpSymExeDefs with EngineBase {
         val v2 = eval(val2, ty2, ss)
         ss.update(v2, v1, ty1.size)
         k(ss)
-      case CallInst(ty, f, args) =>
-        val argValues: List[LLVMValue] = extractValues(args)
-        val argTypes: List[LLVMType] = extractTypes(args)
-        val fv = eval(f, VoidType, ss, Some(argTypes))
-        val vs = argValues.zip(argTypes).map {
-          case (v, t) => eval(v, t, ss)
-        }
-        ss.push
-        val stackSize = ss.stackSize
-        val res: Rep[List[(SS, Value)]] = fv[Ref](ss, List(vs: _*))
-        res.flatMap { case sv =>
-          val s = sv._1
-          s.pop(stackSize)
-          k(s)
-        }
+      case call@CallInst(ty, f, args) =>
+        execValueInst(call, ss, { case (s, v) => k(s) })
     }
   }
 
diff --git a/src/main/scala/gensym/engines/PureCPSEngine.scala b/src/main/scala/gensym/engines/PureCPSEngine.scala
index d8192f2e..83d6efb2 100644
--- a/src/main/scala/gensym/engines/PureCPSEngine.scala
+++ b/src/main/scala/gensym/engines/PureCPSEngine.scala
@@ -306,13 +306,8 @@ trait PureCPSGSEngine extends SymExeDefs with EngineBase {
         val v1 = eval(val1, ty1, ss)
         val v2 = eval(val2, ty2, ss)
         k(ss.update(v2, v1, ty1.size))
-      case CallInst(ty, f, args) =>
-        val argValues: List[LLVMValue] = extractValues(args)
-        val argTypes: List[LLVMType] = extractTypes(args)
-        val fv = eval(f, VoidType, ss, Some(argTypes))
-        val vs = argValues.zip(argTypes).map { case (v, t) => eval(v, t, ss) }
-        def fK(s: Rep[SS], v: Rep[Value]): Rep[Unit] = k(s.pop(ss.stackSize))
-        fv[Id](ss.push, List(vs: _*), ContOpt[Id](fK))
+      case call@CallInst(ty, f, args) =>
+        execValueInst(call, ss, { case (s, v) => k(s) })
     }
   }
 
diff --git a/src/main/scala/gensym/engines/PureEngine.scala b/src/main/scala/gensym/engines/PureEngine.scala
index 4ea4956b..348b5c72 100644
--- a/src/main/scala/gensym/engines/PureEngine.scala
+++ b/src/main/scala/gensym/engines/PureEngine.scala
@@ -393,17 +393,7 @@ trait GSEngine extends StagedNondet with SymExeDefs with EngineBase {
           v2 <- eval(val2, ty2)
           _ <- updateMem(v2, v1, ty1.size)
         } yield ()
-      case CallInst(ty, f, args) =>
-        val argValues: List[LLVMValue] = extractValues(args)
-        val argTypes: List[LLVMType] = extractTypes(args)
-        for {
-          fv <- eval(f, VoidType, Some(argTypes))
-          vs <- mapM2Tup(argValues)(argTypes)(eval(_, _, None))
-          _ <- pushFrame
-          s <- getState
-          v <- reflect(fv[Id](s, List(vs:_*)))
-          _ <- popFrame(s.stackSize)
-        } yield ()
+      case call@CallInst(ty, f, args) => for { _ <- execValueInst(call) } yield ()
     }
   }
 

From ad88b1a5d0466042828411630cd602bb480a46e2 Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Sun, 16 Oct 2022 00:31:15 -0400
Subject: [PATCH 03/15] experiment with pushing cont onto stack

---
 headers/gensym/state_tsnt.hpp                 | 33 ++++++++++++-------
 src/main/scala/gensym/Codegen.scala           | 33 +++++++++++++++++--
 src/main/scala/gensym/EngineBase.scala        | 13 ++++----
 src/main/scala/gensym/GenericDefs.scala       | 27 ++++++++++++---
 src/main/scala/gensym/RunGenSym.scala         | 22 ++++++++++---
 .../scala/gensym/engines/ImpCPSEngine.scala   | 31 ++++++++++-------
 .../scala/gensym/engines/PureCPSEngine.scala  |  2 +-
 .../scala/gensym/states/ImpSymExeState.scala  |  6 +++-
 src/main/scala/lms/ListOps.scala              |  2 +-
 9 files changed, 124 insertions(+), 45 deletions(-)

diff --git a/headers/gensym/state_tsnt.hpp b/headers/gensym/state_tsnt.hpp
index 72857301..9a5f6527 100644
--- a/headers/gensym/state_tsnt.hpp
+++ b/headers/gensym/state_tsnt.hpp
@@ -163,13 +163,15 @@ class Frame {
   public:
     using Env = std::map<Id, PtrVal>;
     using Cont = std::function<std::monostate(SS&, PtrVal)>;
-    Cont cont;
+    size_t prev_stack_size;
+    Cont k;
   private:
     Env env;
   public:
-    Frame(Cont ct): cont(ct), env() {}
+    Frame() : env() {}
     Frame(Env env) : env(std::move(env)) {}
-    Frame() : env(std::map<Id, PtrVal>{}) {}
+    Frame(size_t ss, Cont k): prev_stack_size(ss), k(k), env() {}
+
     size_t size() { return env.size(); }
     PtrVal lookup_id(Id id) const { return env.at(id); }
     Frame&& assign(Id id, PtrVal v) {
@@ -205,12 +207,18 @@ class Stack {
     }
     PtrVal error_loc() { return errno_location; }
     typename Frame::Cont pop(size_t keep) {
-      auto &it = env.at(env.size() - 1);
-      auto ret = it.cont;
+      auto& f = env.at(env.size() - 1);
+      auto ret = f.k;
       mem.take(keep);
       env.take(env.size() - 1);
       return ret;
     }
+    std::monostate pop(SS& s, PtrVal v) {
+      auto f = env.at(env.size() - 1);
+      mem.take(f.prev_stack_size);
+      env.take(env.size() - 1);
+      return f.k(s, v);
+    }
     Stack&& push() {
       return push(Frame());
     }
@@ -218,8 +226,8 @@ class Stack {
       env.push_back(std::move(f));
       return std::move(*this);
     }
-    Stack&& push(std::function<std::monostate(SS&, PtrVal)> cont) {
-      return push(Frame(cont));
+    Stack&& push(size_t ss, std::function<std::monostate(SS&, PtrVal)> k) {
+      return push(Frame(ss, k));
     }
 
     Stack&& assign(Id id, PtrVal val) {
@@ -446,13 +454,16 @@ class SS {
       stack.push();
       return std::move(*this);
     }
-    SS&& push(std::function<std::monostate(SS&, PtrVal)> cont) {
-      stack.push(cont);
+    SS&& push(size_t ss, std::function<std::monostate(SS&, PtrVal)> cont) {
+      stack.push(ss, cont);
       return std::move(*this);
     }
     typename Frame::Cont pop(size_t keep) {
       return stack.pop(keep);
     }
+    std::monostate pop(PtrVal v) {
+      return stack.pop(*this, v);
+    }
     SS&& assign(Id id, PtrVal val) {
       stack.assign(id, val);
       return std::move(*this);
@@ -571,8 +582,8 @@ inline std::monostate cps_apply(PtrVal v, SS ss, List<PtrVal> args, std::functio
   ABORT("cps_apply: not applicable");
 }
 
-inline std::monostate cont_apply(std::function<std::monostate(SS&, PtrVal)> cont, SS& ss, PtrVal val) {
-  return cont(ss, val);
+inline std::monostate pop_cont_apply(SS& ss, PtrVal val) {
+  return ss.pop(val);
 }
 
 #endif
diff --git a/src/main/scala/gensym/Codegen.scala b/src/main/scala/gensym/Codegen.scala
index 56f19c1e..b0eedd37 100644
--- a/src/main/scala/gensym/Codegen.scala
+++ b/src/main/scala/gensym/Codegen.scala
@@ -106,6 +106,34 @@ trait GenericGSCodeGen extends CppSAICodeGenBase {
     case _ => super.traverse(n)
   }
 
+  // Note: this enhances the EmitHelper in lms-clean, to be merged
+  implicit class EmitHelperExt(val sc: StringContext) {
+    def realEmit(x: Any): Unit = x match {
+      case (x: Def) => shallow(x)
+      case (n: Node) => shallow(n)
+      case (s: String) => emit(s)
+      case (i: Int) => emit(i.toString)
+      case (args: List[Def]) =>
+        if (args.nonEmpty) {
+          realEmit(args(0))
+          args.tail.foreach { arg => emit(", "); realEmit(arg) }
+        }
+    }
+    def es(args: Any*): Unit = {
+      val strings = sc.parts.iterator
+      val expressions = args.iterator
+      emit(strings.next)
+      while(strings.hasNext) {
+        realEmit(expressions.next)
+        emit(strings.next)
+      }
+    }
+    def esln(args: Any*): Unit = {
+      es(args:_*)
+      emitln()
+    }
+  }
+
   override def shallow(n: Node): Unit = n match {
     case n @ Node(s, "P", List(x), _) => es"std::cout << $x << std::endl"
     case Node(s,"kStack", _, _) => emit("LocV::kStack")
@@ -137,9 +165,8 @@ trait GenericGSCodeGen extends CppSAICodeGenBase {
     case Node(s, "ss-update", List(ss, k, v, sz), _) => es"$ss.update($k, $v, $sz)"
     case Node(s, "ss-update", List(ss, k, v), _) => es"$ss.update($k, $v)"
     case Node(s, "ss-update-seq", List(ss, k, v), _) => es"$ss.update_seq($k, $v)"
-    case Node(s, "ss-push", List(ss), _) => es"$ss.push()"
-    case Node(s, "ss-push", List(ss, k), _) => es"$ss.push($k)"
-    case Node(s, "ss-pop", List(ss, n), _) => es"$ss.pop($n)"
+    case Node(s, "ss-push", ss::args, _) => es"$ss.push($args)"
+    case Node(s, "ss-pop", ss::args, _) => es"$ss.pop($args)"
     case Node(s, "ss-addpc", List(ss, e), _) => es"$ss.add_PC($e)"
     case Node(s, "add-pc", List(pc, e), _) => es"$pc.add($e)"
     case Node(s, "ss-addpcset", List(ss, es), _) => es"$ss.add_PC_set($es)"
diff --git a/src/main/scala/gensym/EngineBase.scala b/src/main/scala/gensym/EngineBase.scala
index 142c402b..18d17108 100644
--- a/src/main/scala/gensym/EngineBase.scala
+++ b/src/main/scala/gensym/EngineBase.scala
@@ -56,10 +56,12 @@ trait EngineBase extends SAIOps { self: BasicDefs with ValueDefs =>
 
   def info(msg: String) = unchecked("INFO(\"" + msg + "\")")
 
-  def getRealFunName(funName: String): String = {
-    val newFname = if (funName != "@main") "__GS_USER_"+funName.tail else "gs_main"
-    newFname.replaceAllLiterally(".","_")
-  }
+  def strippedFunName(funName: String): String = 
+    if (funName != "@main") funName.tail.replaceAllLiterally(".", "_")
+    else "gs_main"
+
+  def getRealFunName(funName: String): String = "__GS_USER_"+strippedFunName(funName)
+
   def getRealBlockFunName(ctx: Ctx): String = blockNameMap(Counter.block.get(ctx.toString))
 
   def compile(funName: String, b: BB): Unit = {
@@ -98,8 +100,7 @@ trait EngineBase extends SAIOps { self: BasicDefs with ValueDefs =>
     }
     val fn = repExternFun(f, ret, argTypes)
     val node = Unwrap(fn).asInstanceOf[Backend.Sym]
-    funNameMap(node) = "__GS_NATIVE_EXTERNAL_"+mangledName.tail
-    FunFuns(mangledName) = fn
+    funNameMap(node) = "__GS_NATIVE_"+mangledName.tail
   }
 
   // Note: the following two functions checks/retrieves functions considering aliases.
diff --git a/src/main/scala/gensym/GenericDefs.scala b/src/main/scala/gensym/GenericDefs.scala
index ca63f6d4..a9b7dab9 100644
--- a/src/main/scala/gensym/GenericDefs.scala
+++ b/src/main/scala/gensym/GenericDefs.scala
@@ -34,6 +34,7 @@ object Counter {
   val block = Counter()
   val variable = Counter()
   val function = Counter()
+  val cont = Counter()
   val branchStat: HashMap[Int, Int] = HashMap[Int, Int]()
   def setBranchNum(ctx: Ctx, n: Int): Unit = {
     val blockId = Counter.block.get(ctx.toString)
@@ -374,8 +375,15 @@ trait ValueDefs { self: SAIOps with BasicDefs with Opaques =>
     def apply(op: String, o1: Rep[Value], o2: Rep[Value]) = "float_op_2".reflectWith[Value](op, o1, o2)
   }
 
+  object ContOpt {
+    def dummyCont[W[_]](implicit m: Manifest[W[SS]]): ContOpt[W] = ContOpt[W]((s, v) => ())
+    def fromRepCont[W[_]](k: Rep[PCont[W]])(implicit m: Manifest[W[SS]]) = ContOpt[W]((s, v) => k(s, v))
+  }
+
   case class ContOpt[W[_]](k: (Rep[W[SS]], Rep[Value]) => Rep[Unit])(implicit m: Manifest[W[SS]]) {
-    lazy val repK = fun(k(_, _))
+    lazy val repK: Rep[PCont[W]] =
+      if (Config.onStackCont) unchecked[PCont[W]]("pop_cont_apply")
+      else fun(k(_, _))
     def apply(s: Rep[W[SS]], v: Rep[Value]): Rep[Unit] = k(s, v)
   }
 
@@ -417,8 +425,9 @@ trait ValueDefs { self: SAIOps with BasicDefs with Opaques =>
         case _ => "direct_apply".reflectWith[List[(SS, Value)]](v, s, args)
       }
 
-    // The CPS version
-    // W[_] is parameterized over pass-by-value (Id) or pass-by-ref (Ref) of SS
+    /*
+    // This `apply` is only needed in the top-level `exec`.
+    // W[_] is parameterized over pass-by-value (Id[_]) or pass-by-ref (Ref[_]) of SS.
     def apply[W[_]](s: Rep[W[SS]], args: Rep[List[Value]], k: Rep[PCont[W]])(implicit m: Manifest[W[SS]]): Rep[Unit] =
       v match {
         case ExternalFun("noop", ty) if Config.opt => k(s, defaultRetVal(ty))
@@ -426,18 +435,26 @@ trait ValueDefs { self: SAIOps with BasicDefs with Opaques =>
         case CPSFunV(f) => f(s, args, k)                       // direct call
         case _ => "cps_apply".reflectWith[Unit](v, s, args, k) // indirect call
       }
+    */
 
+    // This `apply` works for CPS version that takes an optimizable continuation `ContOpt`.
+    // Using `ContOpt`, we may choose to call the continuation at staging-time, or to generate
+    // the continuation function into the second stage.
+    // W[_] is parameterized over pass-by-value (Id[_]) or pass-by-ref (Ref[_]) of SS.
     def apply[W[_]](s: Rep[W[SS]], args: Rep[List[Value]], k: ContOpt[W])(implicit m: Manifest[W[SS]]): Rep[Unit] =
       v match {
+        /*
         case ExternalFun("noop", ty) if Config.opt => k(s, defaultRetVal(ty))
-        case ExternalFun(f, ty) if ExternalFun.isDeterministic(f) && !usingPureEngine =>
+        case ExternalFun(f, ty) if Config.opt && ExternalFun.isDeterministic(f) && !usingPureEngine =>
+          // This is an optimization that avoids generating CPS code for deterministic function (ie those that won't fork).
           // Be careful: since the state is not passed/returned, with the imperative backend it means
           // the state must be passed by reference to f_det! Currently not all deterministic functions
           // defined in backend works in this way (see external_shared.hpp).
           k(s, (f+"_det").reflectCtrlWith[Value](s, args))
-        case ExternalFun(f, ty) if ExternalFun.isDeterministic(f) && usingPureEngine =>
+        case ExternalFun(f, ty) if Config.opt && ExternalFun.isDeterministic(f) && usingPureEngine =>
           val sv = (f+"_det").reflectCtrlWith[(W[SS], Value)](s, args)
           k(sv._1, sv._2)
+          */
         case ExternalFun(f, ty) => f.reflectWith[Unit](s, args, k.repK)
         case CPSFunV(f) => f(s, args, k.repK)                       // direct call
         case _ => "cps_apply".reflectWith[Unit](v, s, args, k.repK) // indirect call
diff --git a/src/main/scala/gensym/RunGenSym.scala b/src/main/scala/gensym/RunGenSym.scala
index 2b7b40af..0fc1246e 100644
--- a/src/main/scala/gensym/RunGenSym.scala
+++ b/src/main/scala/gensym/RunGenSym.scala
@@ -23,21 +23,27 @@ object Config {
   val (o0, o1, o2, o3) = ("O0", "O1", "O2", "O3")
 
   /* Global compile-time configurations */
+
+  // use first-stage compile-time optimizations
   var opt: Boolean = true
+  // compile `select` instructions to ITE or branches
   var iteSelect: Boolean = true
+  // generate Makefile with -g and -DDEBUG
   var genDebug: Boolean = false
+  // generate a map from source variable names to their compiled name
   var emitVarIdMap: Boolean = true
+  // generate a map from source block labels to their compiled name
   var emitBlockIdMap: Boolean = true
+  // merge same switch targets to disjunctive conditions instead of forking
   var switchType: SwitchType = Merge
-  var runCode: Boolean = true
+  // generate code that records executed instruction numbers
   var recordInstNum: Boolean = false
+  // push continuations onto stack or pass continuations as arguments (applicable for ImpCPS engine)
+  var onStackCont: Boolean = false
 
   def disableOpt: Unit = opt = false
   def enableOpt: Unit = opt = true
 
-  def disableRunCode: Unit = runCode = false
-  def enableRunCode: Unit = runCode = true
-
   def symArg(n: Int) = Config(n, false, o3)
   def useArgv = Config(0, true, o3)
   def noArg = Config(0, false, o3)
@@ -72,6 +78,8 @@ object RunGenSym {
     |--switch-type=<string>  - compilation variants of `switch` statement (default=nonMerge)
     |  =merge                - only fork `m` paths of distinct targets
     |  =nonMerge             - fork `n` paths where `n` is the total number of feasible cases (including default)
+    |--on-stack-cont         - push continuations onto stack, otherwise continuations are passed as argument (only applicable for ImpCPS engine, default=true)
+    |                        -   push continuations onto stack avoids generating nested C++ closures
     |--help                  - print this help message
     """
 
@@ -86,6 +94,7 @@ object RunGenSym {
       case (options, r"emit-block-id-map") => options + ("blockIdMap" -> true)
       case (options, r"emit-var-id-map") => options + ("varIdMap" -> true)
       case (options, r"--switch-type=(\w+)$t") => options + ("switchType" -> SwitchType.fromString(t))
+      case (options, r"--on-stack-cont") => options + ("onStackCont" -> true)
       case (options, "--help") => println(usage.stripMargin); sys.exit(0)
       case (options, input) => options + ("input" -> input)
     }
@@ -100,6 +109,7 @@ object RunGenSym {
     val emitBlockIdMap = options.getOrElse("blockIdMap", Config.emitBlockIdMap).asInstanceOf[Boolean]
     val emitVarIdMap = options.getOrElse("varIdMap", Config.emitVarIdMap).asInstanceOf[Boolean]
     val switchType = options.getOrElse("switchType", SwitchType.NonMerge).asInstanceOf[SwitchType]
+    val onStackCont = options.getOrElse("onStackCont", Config.onStackCont).asInstanceOf[Boolean]
 
     val gensym = engine match {
       case "ImpCPS" => new ImpCPSGS
@@ -111,9 +121,11 @@ object RunGenSym {
     Config.emitBlockIdMap = emitBlockIdMap
     Config.emitVarIdMap = emitVarIdMap
     Config.switchType = switchType
+    Config.onStackCont = onStackCont
     val info = s"""Running $engine with
     |  filepath=$filepath, entrance=$entrance, output=$output,
-    |  nSym=$nSym, useArgv=$useArgv, optimize=$optimize, mainOpt=$mainOpt, switchType=$switchType"""
+    |  nSym=$nSym, useArgv=$useArgv, optimize=$optimize, mainOpt=$mainOpt, switchType=$switchType
+    |  onStackCont=$onStackCont"""
     println(info.stripMargin)
     gensym.run(parseFile(filepath), output, entrance, Config(nSym, useArgv, mainOpt))
   }
diff --git a/src/main/scala/gensym/engines/ImpCPSEngine.scala b/src/main/scala/gensym/engines/ImpCPSEngine.scala
index a15d18be..bf29630f 100644
--- a/src/main/scala/gensym/engines/ImpCPSEngine.scala
+++ b/src/main/scala/gensym/engines/ImpCPSEngine.scala
@@ -45,8 +45,11 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
     "async_exec_block".reflectWriteWith[Unit](unchecked[String](realBlockFunName), ss, k)(Adapter.CTRL)
   }
 
-  def contApply(cont: Rep[Cont], ss: Rep[SS], v: Rep[Value]): Rep[Unit] = {
-    "cont_apply".reflectWriteWith[Unit](cont, ss, v)(Adapter.CTRL)
+  def compileCont(k: (Rep[Ref[SS]], Rep[Value]) => Rep[Unit])(implicit ctx: Ctx): Rep[Cont] = {
+    val repK = fun(k)
+    val node = Unwrap(repK).asInstanceOf[Backend.Sym]
+    funNameMap(node) = strippedFunName(ctx.funName) + "_k_" + Counter.cont.fresh
+    repK
   }
 
   def eval(v: LLVMValue, ty: LLVMType, ss: Rep[SS], argTypes: Option[List[LLVMType]] = None)(implicit ctx: Ctx): Rep[Value] =
@@ -108,7 +111,7 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
   def evalFloatOp2(op: String, lhs: LLVMValue, rhs: LLVMValue, ty: LLVMType, ss: Rep[SS])(implicit ctx: Ctx): Rep[Value] =
     FloatOp2(op, eval(lhs, ty, ss), eval(rhs, ty, ss))
 
-  def execValueInst(inst: ValueInstruction, ss: Rep[SS])(k: (Rep[SS], Rep[Value]) => Rep[Unit])(implicit ctx: Ctx): Rep[Unit] = {
+  def execValueInst(inst: ValueInstruction, ss: Rep[SS])(k: (Rep[Ref[SS]], Rep[Value]) => Rep[Unit])(implicit ctx: Ctx): Rep[Unit] = {
     inst match {
       // Memory Access Instructions
       case AllocaInst(ty, align) =>
@@ -189,12 +192,16 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
         val argTypes: List[LLVMType] = extractTypes(args)
         val fv = eval(f, VoidType, ss, Some(argTypes))
         val vs = argValues.zip(argTypes).map { case (v, t) => eval(v, t, ss) }
-        val stackSize = ss.stackSize
-        // ss.push(ss.stackSize, k)
-        ss.push
-        //def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(s, v) }
-        def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(stackSize); k(s, v) }
-        fv[Ref](ss, List(vs: _*), ContOpt(fK))
+        if (Config.onStackCont) {
+          ss.push(ss.stackSize, compileCont(k))
+          def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = s.popRet(v)
+          fv[Ref](ss, List(vs: _*), ContOpt.dummyCont[Ref])
+        } else {
+          val stackSize = ss.stackSize
+          ss.push
+          def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(stackSize); k(s, v) }
+          fv[Ref](ss, List(vs: _*), ContOpt(fK))
+        }
       case PhiInst(ty, incs) =>
         def selectValue(bb: Rep[BlockLabel], vs: List[() => Rep[Value]], labels: List[BlockLabel]): Rep[Value] = {
           if (bb == labels(0) || labels.length == 1) vs(0)()
@@ -226,13 +233,13 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
 
   def execTerm(inst: Terminator, k: Rep[Cont])(implicit ss: Rep[SS], ctx: Ctx): Rep[Unit] = {
     inst match {
-      case Unreachable => contApply(k, ss, IntV(-1))
+      case Unreachable => k(ss, IntV(-1))
       case RetTerm(ty, v) =>
         val ret = v match {
           case Some(value) => eval(value, ty, ss)
           case None => NullPtr[Value]
         }
-        contApply(k, ss, ret)
+        k(ss, ret)
       case BrTerm(lab) if (cfg.pred(ctx.funName, lab).size == 1) =>
         execBlockEager(findBlock(ctx.funName, lab).get, ss, k)(Ctx(ctx.funName, lab))
       case BrTerm(lab) =>
@@ -405,6 +412,6 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
     ss.push
     ss.updateArg
     ss.initErrorLoc
-    fv[Ref](ss, args, k)
+    fv[Ref](ss, args, ContOpt.fromRepCont[Ref](k))
   }
 }
diff --git a/src/main/scala/gensym/engines/PureCPSEngine.scala b/src/main/scala/gensym/engines/PureCPSEngine.scala
index 83d6efb2..26f355c3 100644
--- a/src/main/scala/gensym/engines/PureCPSEngine.scala
+++ b/src/main/scala/gensym/engines/PureCPSEngine.scala
@@ -381,6 +381,6 @@ trait PureCPSGSEngine extends SymExeDefs with EngineBase {
     Coverage.incPath(1)
     val ss = initState(preHeap.asRepOf[Mem])
     val fv = eval(GlobalId(fname), VoidType, ss)
-    fv[Id](ss.push.updateArg.initErrorLoc, args, k)
+    fv[Id](ss.push.updateArg.initErrorLoc, args, ContOpt.fromRepCont[Id](k))
   }
 }
diff --git a/src/main/scala/gensym/states/ImpSymExeState.scala b/src/main/scala/gensym/states/ImpSymExeState.scala
index bd9c173c..e9bce4d0 100644
--- a/src/main/scala/gensym/states/ImpSymExeState.scala
+++ b/src/main/scala/gensym/states/ImpSymExeState.scala
@@ -111,9 +111,13 @@ trait ImpSymExeDefs extends SAIOps with BasicDefs with ValueDefs with Opaques wi
 
     // push before function call could be DCE-ed, due to high-level function dependency issue.
     def push: Rep[Unit] = reflectWrite[Unit]("ss-push", ss)(ss, Adapter.CTRL)
-    def push(cont: Rep[Cont]): Rep[Unit] = reflectWrite[Unit]("ss-push", ss, cont)(ss, Adapter.CTRL)
+    def push(k: Rep[Cont]): Rep[Unit] = reflectWrite[Unit]("ss-push", ss, k)(ss, Adapter.CTRL)
+    def push(stackSize: Rep[Int], k: Rep[Cont]): Rep[Unit] = reflectWrite[Unit]("ss-push", ss, stackSize, k)(ss, Adapter.CTRL)
     // XXX: since pop is used in a map, will be DCE-ed if no CTRL
     def pop(keep: Rep[Int]): Rep[Cont] = reflectWrite[Cont]("ss-pop", ss, keep)(ss, Adapter.CTRL)
+    def popRet(v: Rep[Value]): Rep[Unit] = reflectWrite[Unit]("ss-pop", ss, v)(ss, Adapter.CTRL)
+      //"pop_cont_apply".reflectUnsafeWith[Unit](ss, v) // Note: we want to inline this
+
     def addPC(e: Rep[Value]): Rep[Unit] = reflectWrite[Unit]("ss-addpc", ss, e)(ss)
     def addPCSet(es: Rep[List[Value]]): Rep[Unit] = reflectWrite[Unit]("ss-addpcset", ss, es)(ss)
     def pc: Rep[PC] = reflectRead[PC]("get-pc", ss)(ss)
diff --git a/src/main/scala/lms/ListOps.scala b/src/main/scala/lms/ListOps.scala
index 1a4e0189..c5ce9d1f 100644
--- a/src/main/scala/lms/ListOps.scala
+++ b/src/main/scala/lms/ListOps.scala
@@ -293,7 +293,7 @@ trait CppCodeGen_List extends ExtendedCCodeGen {
       emit("{")
       if (!xs.isEmpty) {
         shallow(xs.head)
-        xs.tail.map { x =>
+        xs.tail.foreach { x =>
           emit(", ")
           shallow(x)
         }

From 6381aef61bf2017fdd19cadd10088d2d1045dcd1 Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Sun, 16 Oct 2022 00:56:19 -0400
Subject: [PATCH 04/15] working version with onStackCont

---
 src/main/scala/gensym/GenericDefs.scala       | 14 -------------
 src/main/scala/gensym/RunGenSym.scala         |  2 +-
 .../scala/gensym/engines/ImpCPSEngine.scala   | 20 ++++++++++++++-----
 3 files changed, 16 insertions(+), 20 deletions(-)

diff --git a/src/main/scala/gensym/GenericDefs.scala b/src/main/scala/gensym/GenericDefs.scala
index a9b7dab9..fc2827d0 100644
--- a/src/main/scala/gensym/GenericDefs.scala
+++ b/src/main/scala/gensym/GenericDefs.scala
@@ -425,25 +425,12 @@ trait ValueDefs { self: SAIOps with BasicDefs with Opaques =>
         case _ => "direct_apply".reflectWith[List[(SS, Value)]](v, s, args)
       }
 
-    /*
-    // This `apply` is only needed in the top-level `exec`.
-    // W[_] is parameterized over pass-by-value (Id[_]) or pass-by-ref (Ref[_]) of SS.
-    def apply[W[_]](s: Rep[W[SS]], args: Rep[List[Value]], k: Rep[PCont[W]])(implicit m: Manifest[W[SS]]): Rep[Unit] =
-      v match {
-        case ExternalFun("noop", ty) if Config.opt => k(s, defaultRetVal(ty))
-        case ExternalFun(f, ty) => f.reflectWith[Unit](s, args, k)
-        case CPSFunV(f) => f(s, args, k)                       // direct call
-        case _ => "cps_apply".reflectWith[Unit](v, s, args, k) // indirect call
-      }
-    */
-
     // This `apply` works for CPS version that takes an optimizable continuation `ContOpt`.
     // Using `ContOpt`, we may choose to call the continuation at staging-time, or to generate
     // the continuation function into the second stage.
     // W[_] is parameterized over pass-by-value (Id[_]) or pass-by-ref (Ref[_]) of SS.
     def apply[W[_]](s: Rep[W[SS]], args: Rep[List[Value]], k: ContOpt[W])(implicit m: Manifest[W[SS]]): Rep[Unit] =
       v match {
-        /*
         case ExternalFun("noop", ty) if Config.opt => k(s, defaultRetVal(ty))
         case ExternalFun(f, ty) if Config.opt && ExternalFun.isDeterministic(f) && !usingPureEngine =>
           // This is an optimization that avoids generating CPS code for deterministic function (ie those that won't fork).
@@ -454,7 +441,6 @@ trait ValueDefs { self: SAIOps with BasicDefs with Opaques =>
         case ExternalFun(f, ty) if Config.opt && ExternalFun.isDeterministic(f) && usingPureEngine =>
           val sv = (f+"_det").reflectCtrlWith[(W[SS], Value)](s, args)
           k(sv._1, sv._2)
-          */
         case ExternalFun(f, ty) => f.reflectWith[Unit](s, args, k.repK)
         case CPSFunV(f) => f(s, args, k.repK)                       // direct call
         case _ => "cps_apply".reflectWith[Unit](v, s, args, k.repK) // indirect call
diff --git a/src/main/scala/gensym/RunGenSym.scala b/src/main/scala/gensym/RunGenSym.scala
index 0fc1246e..fc988879 100644
--- a/src/main/scala/gensym/RunGenSym.scala
+++ b/src/main/scala/gensym/RunGenSym.scala
@@ -39,7 +39,7 @@ object Config {
   // generate code that records executed instruction numbers
   var recordInstNum: Boolean = false
   // push continuations onto stack or pass continuations as arguments (applicable for ImpCPS engine)
-  var onStackCont: Boolean = false
+  var onStackCont: Boolean = true
 
   def disableOpt: Unit = opt = false
   def enableOpt: Unit = opt = true
diff --git a/src/main/scala/gensym/engines/ImpCPSEngine.scala b/src/main/scala/gensym/engines/ImpCPSEngine.scala
index bf29630f..38a25844 100644
--- a/src/main/scala/gensym/engines/ImpCPSEngine.scala
+++ b/src/main/scala/gensym/engines/ImpCPSEngine.scala
@@ -195,7 +195,10 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
         if (Config.onStackCont) {
           ss.push(ss.stackSize, compileCont(k))
           def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = s.popRet(v)
-          fv[Ref](ss, List(vs: _*), ContOpt.dummyCont[Ref])
+          // Note: fK is used only when we want to apply the continuation at staing-time,
+          // we will not reify it into second stage. Instead, backend-defined `pop_cont_apply`
+          // that will be used to apply the pushed continuation with returned value.
+          fv[Ref](ss, List(vs: _*), ContOpt(fK))
         } else {
           val stackSize = ss.stackSize
           ss.push
@@ -409,9 +412,16 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
     Coverage.incPath(1)
     val ss = initState(preHeap.asRepOf[Mem])
     val fv = eval(GlobalId(fname), VoidType, ss)
-    ss.push
-    ss.updateArg
-    ss.initErrorLoc
-    fv[Ref](ss, args, ContOpt.fromRepCont[Ref](k))
+    if (Config.onStackCont) {
+      ss.push(0, k)
+      ss.updateArg
+      ss.initErrorLoc
+      fv[Ref](ss, args, ContOpt.dummyCont[Ref])
+    } else {
+      ss.push
+      ss.updateArg
+      ss.initErrorLoc
+      fv[Ref](ss, args, ContOpt.fromRepCont[Ref](k))
+    }
   }
 }

From 33fb7c36e54eccb257356bf33fac57e8c6ec4fe8 Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Sun, 16 Oct 2022 01:13:00 -0400
Subject: [PATCH 05/15] a few tweak

---
 headers/gensym/state_tsnt.hpp                     | 11 +++++------
 src/main/scala/gensym/engines/ImpCPSEngine.scala  | 11 +++++++----
 src/main/scala/gensym/states/ImpSymExeState.scala |  2 +-
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/headers/gensym/state_tsnt.hpp b/headers/gensym/state_tsnt.hpp
index 9a5f6527..9b0af34e 100644
--- a/headers/gensym/state_tsnt.hpp
+++ b/headers/gensym/state_tsnt.hpp
@@ -206,12 +206,10 @@ class Stack {
       return std::move(*this);
     }
     PtrVal error_loc() { return errno_location; }
-    typename Frame::Cont pop(size_t keep) {
-      auto& f = env.at(env.size() - 1);
-      auto ret = f.k;
+    Stack&& pop(size_t keep) {
       mem.take(keep);
       env.take(env.size() - 1);
-      return ret;
+      return std::move(*this);
     }
     std::monostate pop(SS& s, PtrVal v) {
       auto f = env.at(env.size() - 1);
@@ -458,8 +456,9 @@ class SS {
       stack.push(ss, cont);
       return std::move(*this);
     }
-    typename Frame::Cont pop(size_t keep) {
-      return stack.pop(keep);
+    SS&& pop(size_t keep) {
+      stack.pop(keep);
+      return std::move(*this);
     }
     std::monostate pop(PtrVal v) {
       return stack.pop(*this, v);
diff --git a/src/main/scala/gensym/engines/ImpCPSEngine.scala b/src/main/scala/gensym/engines/ImpCPSEngine.scala
index 38a25844..99e5c6b1 100644
--- a/src/main/scala/gensym/engines/ImpCPSEngine.scala
+++ b/src/main/scala/gensym/engines/ImpCPSEngine.scala
@@ -192,17 +192,20 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
         val argTypes: List[LLVMType] = extractTypes(args)
         val fv = eval(f, VoidType, ss, Some(argTypes))
         val vs = argValues.zip(argTypes).map { case (v, t) => eval(v, t, ss) }
+        val stackSize = ss.stackSize
+        def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(stackSize); k(s, v) }
         if (Config.onStackCont) {
-          ss.push(ss.stackSize, compileCont(k))
-          def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = s.popRet(v)
+          ss.push(stackSize, compileCont(k))
           // Note: fK is used only when we want to apply the continuation at staing-time,
           // we will not reify it into second stage. Instead, backend-defined `pop_cont_apply`
           // that will be used to apply the pushed continuation with returned value.
+          // TODO: this is still not idea, since `compileCont(k)` will anyway generate/push
+          // a continuation function, even though we can statically reduce it some time.
+          // See an example in @merge_%while.end64 of merge_sort; we also want to eliminate dummy
+          // push/pop pair, which benefits the case if not using onStackCont.
           fv[Ref](ss, List(vs: _*), ContOpt(fK))
         } else {
-          val stackSize = ss.stackSize
           ss.push
-          def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(stackSize); k(s, v) }
           fv[Ref](ss, List(vs: _*), ContOpt(fK))
         }
       case PhiInst(ty, incs) =>
diff --git a/src/main/scala/gensym/states/ImpSymExeState.scala b/src/main/scala/gensym/states/ImpSymExeState.scala
index e9bce4d0..c3725b12 100644
--- a/src/main/scala/gensym/states/ImpSymExeState.scala
+++ b/src/main/scala/gensym/states/ImpSymExeState.scala
@@ -114,7 +114,7 @@ trait ImpSymExeDefs extends SAIOps with BasicDefs with ValueDefs with Opaques wi
     def push(k: Rep[Cont]): Rep[Unit] = reflectWrite[Unit]("ss-push", ss, k)(ss, Adapter.CTRL)
     def push(stackSize: Rep[Int], k: Rep[Cont]): Rep[Unit] = reflectWrite[Unit]("ss-push", ss, stackSize, k)(ss, Adapter.CTRL)
     // XXX: since pop is used in a map, will be DCE-ed if no CTRL
-    def pop(keep: Rep[Int]): Rep[Cont] = reflectWrite[Cont]("ss-pop", ss, keep)(ss, Adapter.CTRL)
+    def pop(keep: Rep[Int]): Rep[Unit] = reflectWrite[Unit]("ss-pop", ss, keep)(ss, Adapter.CTRL)
     def popRet(v: Rep[Value]): Rep[Unit] = reflectWrite[Unit]("ss-pop", ss, v)(ss, Adapter.CTRL)
       //"pop_cont_apply".reflectUnsafeWith[Unit](ss, v) // Note: we want to inline this
 

From 34ccec942c2526b7140074697d3345ad5a823c55 Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Sun, 16 Oct 2022 01:15:11 -0400
Subject: [PATCH 06/15] minor

---
 src/main/scala/gensym/engines/ImpCPSEngine.scala | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/main/scala/gensym/engines/ImpCPSEngine.scala b/src/main/scala/gensym/engines/ImpCPSEngine.scala
index 99e5c6b1..74392cba 100644
--- a/src/main/scala/gensym/engines/ImpCPSEngine.scala
+++ b/src/main/scala/gensym/engines/ImpCPSEngine.scala
@@ -196,18 +196,18 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
         def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(stackSize); k(s, v) }
         if (Config.onStackCont) {
           ss.push(stackSize, compileCont(k))
-          // Note: fK is used only when we want to apply the continuation at staing-time,
-          // we will not reify it into second stage. Instead, backend-defined `pop_cont_apply`
+          // Note: in this case, even we have compiled cont k, fK still will be used only when 
+          // we want to apply the continuation at staing-time (as an optimization);
+          // we will not reify `fK` into second stage. Instead, backend-defined `pop_cont_apply`
           // that will be used to apply the pushed continuation with returned value.
           // TODO: this is still not idea, since `compileCont(k)` will anyway generate/push
           // a continuation function, even though we can statically reduce it some time.
           // See an example in @merge_%while.end64 of merge_sort; we also want to eliminate dummy
           // push/pop pair, which benefits the case if not using onStackCont.
-          fv[Ref](ss, List(vs: _*), ContOpt(fK))
         } else {
           ss.push
-          fv[Ref](ss, List(vs: _*), ContOpt(fK))
         }
+        fv[Ref](ss, List(vs: _*), ContOpt(fK))
       case PhiInst(ty, incs) =>
         def selectValue(bb: Rep[BlockLabel], vs: List[() => Rep[Value]], labels: List[BlockLabel]): Rep[Value] = {
           if (bb == labels(0) || labels.length == 1) vs(0)()

From 52f140ab506716878b65c44818e0818a5e270240 Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Sun, 16 Oct 2022 01:37:35 -0400
Subject: [PATCH 07/15] misc

---
 src/main/scala/gensym/engines/ImpCPSEngine.scala | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/src/main/scala/gensym/engines/ImpCPSEngine.scala b/src/main/scala/gensym/engines/ImpCPSEngine.scala
index 74392cba..7cf3185c 100644
--- a/src/main/scala/gensym/engines/ImpCPSEngine.scala
+++ b/src/main/scala/gensym/engines/ImpCPSEngine.scala
@@ -415,16 +415,10 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
     Coverage.incPath(1)
     val ss = initState(preHeap.asRepOf[Mem])
     val fv = eval(GlobalId(fname), VoidType, ss)
-    if (Config.onStackCont) {
-      ss.push(0, k)
-      ss.updateArg
-      ss.initErrorLoc
-      fv[Ref](ss, args, ContOpt.dummyCont[Ref])
-    } else {
-      ss.push
-      ss.updateArg
-      ss.initErrorLoc
-      fv[Ref](ss, args, ContOpt.fromRepCont[Ref](k))
-    }
+    if (Config.onStackCont) ss.push(0, k)
+    else ss.push
+    ss.updateArg
+    ss.initErrorLoc
+    fv[Ref](ss, args, ContOpt.fromRepCont[Ref](k))
   }
 }

From 83cb8581013984b854d4751b244fb1a5f24c5e55 Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Sun, 16 Oct 2022 14:02:33 -0400
Subject: [PATCH 08/15] fix

---
 src/main/scala/gensym/EngineBase.scala           | 1 +
 src/main/scala/gensym/engines/ImpCPSEngine.scala | 4 ++--
 src/test/scala/gensym/TestGS.scala               | 3 ++-
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/main/scala/gensym/EngineBase.scala b/src/main/scala/gensym/EngineBase.scala
index 18d17108..91e72189 100644
--- a/src/main/scala/gensym/EngineBase.scala
+++ b/src/main/scala/gensym/EngineBase.scala
@@ -101,6 +101,7 @@ trait EngineBase extends SAIOps { self: BasicDefs with ValueDefs =>
     val fn = repExternFun(f, ret, argTypes)
     val node = Unwrap(fn).asInstanceOf[Backend.Sym]
     funNameMap(node) = "__GS_NATIVE_"+mangledName.tail
+    FunFuns(mangledName) = fn
   }
 
   // Note: the following two functions checks/retrieves functions considering aliases.
diff --git a/src/main/scala/gensym/engines/ImpCPSEngine.scala b/src/main/scala/gensym/engines/ImpCPSEngine.scala
index 7cf3185c..2892872c 100644
--- a/src/main/scala/gensym/engines/ImpCPSEngine.scala
+++ b/src/main/scala/gensym/engines/ImpCPSEngine.scala
@@ -193,20 +193,20 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
         val fv = eval(f, VoidType, ss, Some(argTypes))
         val vs = argValues.zip(argTypes).map { case (v, t) => eval(v, t, ss) }
         val stackSize = ss.stackSize
-        def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(stackSize); k(s, v) }
         if (Config.onStackCont) {
           ss.push(stackSize, compileCont(k))
           // Note: in this case, even we have compiled cont k, fK still will be used only when 
           // we want to apply the continuation at staing-time (as an optimization);
           // we will not reify `fK` into second stage. Instead, backend-defined `pop_cont_apply`
           // that will be used to apply the pushed continuation with returned value.
-          // TODO: this is still not idea, since `compileCont(k)` will anyway generate/push
+          // TODO: this is still not ideal, since `compileCont(k)` will anyway generate/push
           // a continuation function, even though we can statically reduce it some time.
           // See an example in @merge_%while.end64 of merge_sort; we also want to eliminate dummy
           // push/pop pair, which benefits the case if not using onStackCont.
         } else {
           ss.push
         }
+        def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(stackSize); k(s, v) }
         fv[Ref](ss, List(vs: _*), ContOpt(fK))
       case PhiInst(ty, incs) =>
         def selectValue(bb: Rep[BlockLabel], vs: List[() => Rep[Value]], labels: List[BlockLabel]): Rep[Value] = {
diff --git a/src/test/scala/gensym/TestGS.scala b/src/test/scala/gensym/TestGS.scala
index 8e309451..c07b75a4 100644
--- a/src/test/scala/gensym/TestGS.scala
+++ b/src/test/scala/gensym/TestGS.scala
@@ -145,7 +145,8 @@ class Playground extends TestGS {
   Config.enableOpt
   val gs = new ImpCPSGS
 
-  testGS(gs, TestPrg(mergesort, "mergeSortTest", "@main", noArg, noOpt, nPath(720)))
+  //testGS(gs, TestPrg(mergesort, "mergeSortTest", "@main", noArg, noOpt, nPath(720)))
+  //testGS(gs, TestPrg(kmpmatcher, "kmp", "@main", noArg, noOpt, nPath(1287)))
   //testGS(new PureCPSGS, TestPrg(arrayFlow, "arrayFlow", "@main", noArg, noOpt, nPath(15)++status(0)))
   //testGS(new ImpCPSGS, TestPrg(arrayFlow, "arrayFlow2", "@main", noArg, noOpt, nPath(15)++status(0)))
 

From 297f81576158d0900ac00eee3fc1c426ee8d7f24 Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Sun, 16 Oct 2022 15:26:21 -0400
Subject: [PATCH 09/15] fix

---
 src/main/scala/gensym/GenericDefs.scala | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/main/scala/gensym/GenericDefs.scala b/src/main/scala/gensym/GenericDefs.scala
index fc2827d0..257246cd 100644
--- a/src/main/scala/gensym/GenericDefs.scala
+++ b/src/main/scala/gensym/GenericDefs.scala
@@ -382,7 +382,7 @@ trait ValueDefs { self: SAIOps with BasicDefs with Opaques =>
 
   case class ContOpt[W[_]](k: (Rep[W[SS]], Rep[Value]) => Rep[Unit])(implicit m: Manifest[W[SS]]) {
     lazy val repK: Rep[PCont[W]] =
-      if (Config.onStackCont) unchecked[PCont[W]]("pop_cont_apply")
+      if (Config.onStackCont && !usingPureEngine) unchecked[PCont[W]]("pop_cont_apply")
       else fun(k(_, _))
     def apply(s: Rep[W[SS]], v: Rep[Value]): Rep[Unit] = k(s, v)
   }

From d684d6aedbaf61acb10ae44794656955fb79454b Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Sun, 16 Oct 2022 15:45:32 -0400
Subject: [PATCH 10/15] comment

---
 src/main/scala/gensym/GenericDefs.scala | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/main/scala/gensym/GenericDefs.scala b/src/main/scala/gensym/GenericDefs.scala
index 257246cd..95cdf18c 100644
--- a/src/main/scala/gensym/GenericDefs.scala
+++ b/src/main/scala/gensym/GenericDefs.scala
@@ -425,20 +425,23 @@ trait ValueDefs { self: SAIOps with BasicDefs with Opaques =>
         case _ => "direct_apply".reflectWith[List[(SS, Value)]](v, s, args)
       }
 
-    // This `apply` works for CPS version that takes an optimizable continuation `ContOpt`.
+    // This `apply` works for the CPS version that takes an optimizable continuation `ContOpt`.
     // Using `ContOpt`, we may choose to call the continuation at staging-time, or to generate
     // the continuation function into the second stage.
     // W[_] is parameterized over pass-by-value (Id[_]) or pass-by-ref (Ref[_]) of SS.
     def apply[W[_]](s: Rep[W[SS]], args: Rep[List[Value]], k: ContOpt[W])(implicit m: Manifest[W[SS]]): Rep[Unit] =
       v match {
-        case ExternalFun("noop", ty) if Config.opt => k(s, defaultRetVal(ty))
+        case ExternalFun("noop", ty) if Config.opt =>
+          // Avoids generating continuations for the `noop` function.
+          k(s, defaultRetVal(ty))
         case ExternalFun(f, ty) if Config.opt && ExternalFun.isDeterministic(f) && !usingPureEngine =>
-          // This is an optimization that avoids generating CPS code for deterministic function (ie those that won't fork).
+          // Avoids generating continuations for the _imperative_ CPS engine if the function is deterministic.
           // Be careful: since the state is not passed/returned, with the imperative backend it means
           // the state must be passed by reference to f_det! Currently not all deterministic functions
           // defined in backend works in this way (see external_shared.hpp).
           k(s, (f+"_det").reflectCtrlWith[Value](s, args))
         case ExternalFun(f, ty) if Config.opt && ExternalFun.isDeterministic(f) && usingPureEngine =>
+          // Avoids generating continuations for the _pure_ engine if the function is deterministic.
           val sv = (f+"_det").reflectCtrlWith[(W[SS], Value)](s, args)
           k(sv._1, sv._2)
         case ExternalFun(f, ty) => f.reflectWith[Unit](s, args, k.repK)

From ff1af289f47fdab66a0775424c690260785e6fcb Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Mon, 17 Oct 2022 14:11:03 -0400
Subject: [PATCH 11/15] bring back at/update

---
 headers/gensym/state_tsnt.hpp      | 19 +++++++++++++++++--
 src/test/scala/gensym/TestGS.scala |  7 +++----
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/headers/gensym/state_tsnt.hpp b/headers/gensym/state_tsnt.hpp
index 9b0af34e..16764be9 100644
--- a/headers/gensym/state_tsnt.hpp
+++ b/headers/gensym/state_tsnt.hpp
@@ -387,7 +387,13 @@ class SS {
       // TODO: should we modify the pc to add the in-bound constraints
       return read_res;
     }
-    PtrVal at(PtrVal addr, size_t size = 1) {
+    PtrVal at(PtrVal addr) {
+      auto loc = addr->to_LocV();
+      ASSERT(loc != nullptr, "Lookup an non-address value");
+      if (loc->k == LocV::kStack) return stack.at(loc->l);
+      return heap.at(loc->l);
+	  }
+    PtrVal at(PtrVal addr, size_t size) {
       if (auto loc = addr->to_LocV()) {
         if (loc->k == LocV::kStack) return stack.at(loc->l, size);
         return heap.at(loc->l, size);
@@ -435,7 +441,16 @@ class SS {
       heap.alloc(size);
       return std::move(*this);
     }
-    SS&& update(PtrVal addr, PtrVal val, size_t size = 1) {
+    SS&& update(PtrVal addr, PtrVal val) {
+      auto loc = addr->to_LocV();
+      ASSERT(loc != nullptr, "Lookup an non-address value");
+      if (loc->k == LocV::kStack)
+        stack.update(loc->l, val);
+      else
+        heap.update(loc->l, val);
+      return std::move(*this);
+    }
+    SS&& update(PtrVal addr, PtrVal val, size_t size) {
       auto loc = addr->to_LocV();
       ASSERT(loc != nullptr, "Lookup an non-address value");
       if (loc->k == LocV::kStack) stack.update(loc->l, val, size);
diff --git a/src/test/scala/gensym/TestGS.scala b/src/test/scala/gensym/TestGS.scala
index c07b75a4..604a7359 100644
--- a/src/test/scala/gensym/TestGS.scala
+++ b/src/test/scala/gensym/TestGS.scala
@@ -126,7 +126,6 @@ class TestImpCPSGS_Z3 extends TestGS {
   testGS(gs, cases)
 }
 
-/*
 class Coreutils extends TestGS {
   import gensym.llvm.parser.Parser._
   Config.enableOpt
@@ -134,11 +133,11 @@ class Coreutils extends TestGS {
   val cases = TestCases.coreutils.map { t =>
     t.copy(runOpt = runtimeOptions ++ t.runOpt, runCode = false)
   }
-  testGS(new ImpCPSGS, cases)
+  //testGS(new ImpCPSGS, cases)
 
-  //testGS(new ImpCPSGS, TestPrg(cat_linked, "cat_linked_posix", "@main", noMainFileOpt, "--argv=./cat.bc --sym-stdout --sym-stdin 2 --sym-arg 2", nPath(28567)++status(0)))
+  lazy val cat_linked = parseFile("benchmarks/gs_posix/cat_linked.ll")
+  testGS(new ImpCPSGS, TestPrg(cat_linked, "cat_linked_posix", "@main", noMainFileOpt, "--argv=./cat.bc --sym-stdout --sym-stdin 2 --sym-arg 2", nPath(28567)++status(0), runCode = false))
 }
-*/
 
 class Playground extends TestGS {
   import gensym.llvm.parser.Parser._

From 33a13f53735f1e692d18a5a7fe480bcc49543764 Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Tue, 18 Oct 2022 21:59:19 -0400
Subject: [PATCH 12/15] further refactor

---
 src/main/scala/gensym/EngineBase.scala        |  1 -
 src/main/scala/gensym/GenericDefs.scala       | 55 ++++++++++---------
 .../scala/gensym/engines/ImpCPSEngine.scala   | 30 +++++-----
 3 files changed, 44 insertions(+), 42 deletions(-)

diff --git a/src/main/scala/gensym/EngineBase.scala b/src/main/scala/gensym/EngineBase.scala
index 91e72189..ed4d7edf 100644
--- a/src/main/scala/gensym/EngineBase.scala
+++ b/src/main/scala/gensym/EngineBase.scala
@@ -148,7 +148,6 @@ trait EngineBase extends SAIOps { self: BasicDefs with ValueDefs =>
         // "When indexing into a (optionally packed) structure, only i32 integer
         //  constants are allowed"
         // TODO: the align argument for getTySize
-        // TODO: test this
         val indexCst: List[Long] = index.map { case IntV(n, _) => n.toLong }
         IntV(calculateOffsetStatic(ty, indexCst), DEFAULT_INDEX_BW)
       case PackedStruct(types) =>
diff --git a/src/main/scala/gensym/GenericDefs.scala b/src/main/scala/gensym/GenericDefs.scala
index 95cdf18c..a7113b8d 100644
--- a/src/main/scala/gensym/GenericDefs.scala
+++ b/src/main/scala/gensym/GenericDefs.scala
@@ -324,16 +324,16 @@ trait ValueDefs { self: SAIOps with BasicDefs with Opaques =>
   }
 
   object IntOp2 {
-    def applyNoOpt(op: String, o1: Rep[Value], o2: Rep[Value]): Rep[Value] =
+    def primOp2(op: String, o1: Rep[Value], o2: Rep[Value]): Rep[Value] =
       "int_op_2".reflectWith[Value](op, o1, o2)
-    def apply(op: String, o1: Rep[Value], o2: Rep[Value]): Rep[Value] =
-      if (!Config.opt) applyNoOpt(op, o1, o2)
-      else op match {
-        case "neq" => neq(o1, o2)
-        case "eq" => eq(o1, o2)
-        case "add" => add(o1, o2)
-        case _ => applyNoOpt(op, o1, o2)
-      }
+
+    def apply(op: String, o1: Rep[Value], o2: Rep[Value]): Rep[Value] = op match {
+      case "neq" => neq(o1, o2)
+      case "eq" => eq(o1, o2)
+      case "add" => add(o1, o2)
+      case "mul" => mul(o1, o2)
+      case _ => primOp2(op, o1, o2)
+    }
 
     def unapply(v: Rep[Value]): Option[(String, Rep[Value], Rep[Value])] = Unwrap(v) match {
       case gNode("int_op_2", bConst(x: String)::(o1: bSym)::(o2: bSym)::_) =>
@@ -342,32 +342,32 @@ trait ValueDefs { self: SAIOps with BasicDefs with Opaques =>
     }
 
     def add(v1: Rep[Value], v2: Rep[Value]): Rep[Value] = (v1, v2) match {
-      case (IntV(n1, bw1), IntV(n2, bw2)) if (bw1 == bw2) => IntV(n1 + n2, bw1)
-      case _ => applyNoOpt("add", v1, v2)
+      case (IntV(n1, bw1), IntV(n2, bw2)) if (bw1 == bw2) && Config.opt => IntV(n1 + n2, bw1)
+      case _ => primOp2("add", v1, v2)
     }
 
     def mul(v1: Rep[Value], v2: Rep[Value]): Rep[Value] = (v1, v2) match {
-      case (IntV(n1, bw1), IntV(n2, bw2)) if (bw1 == bw2) => IntV(n1 * n2, bw1)
-      case _ => applyNoOpt("mul", v1, v2)
+      case (IntV(n1, bw1), IntV(n2, bw2)) if (bw1 == bw2) && Config.opt => IntV(n1 * n2, bw1)
+      case _ => primOp2("mul", v1, v2)
     }
 
     def neq(o1: Rep[Value], o2: Rep[Value]): Rep[Value] = (Unwrap(o1), Unwrap(o2)) match {
       case (gNode("bv_sext", (e1: bExp)::bConst(bw1: Int)::_),
-            gNode("bv_sext", (e2: bExp)::bConst(bw2: Int)::_)) if bw1 == bw2 =>
+            gNode("bv_sext", (e2: bExp)::bConst(bw2: Int)::_)) if bw1 == bw2 && Config.opt =>
         val v1 = Wrap[Value](e1)
         val v2 = Wrap[Value](e2)
-        if (v1.bw == v2.bw) applyNoOpt("neq", v1, v2)
-        else applyNoOpt("neq", o1, o2)
-      case _ => applyNoOpt("neq", o1, o2)
+        if (v1.bw == v2.bw) neq(v1, v2)
+        else primOp2("neq", o1, o2)
+      case _ => primOp2("neq", o1, o2)
     }
     def eq(o1: Rep[Value], o2: Rep[Value]): Rep[Value] = (Unwrap(o1), Unwrap(o2)) match {
       case (gNode("bv_sext", (e1: bExp)::bConst(bw1: Int)::_),
-            gNode("bv_sext", (e2: bExp)::bConst(bw2: Int)::_)) if bw1 == bw2 =>
+            gNode("bv_sext", (e2: bExp)::bConst(bw2: Int)::_)) if bw1 == bw2 && Config.opt =>
         val v1 = Wrap[Value](e1)
         val v2 = Wrap[Value](e2)
-        if (v1.bw == v2.bw) applyNoOpt("eq", v1, v2)
-        else applyNoOpt("eq", o1, o2)
-      case _ => applyNoOpt("eq", o1, o2)
+        if (v1.bw == v2.bw) eq(v1, v2)
+        else primOp2("eq", o1, o2)
+      case _ => primOp2("eq", o1, o2)
     }
   }
 
@@ -454,22 +454,22 @@ trait ValueDefs { self: SAIOps with BasicDefs with Opaques =>
     val foldableOp = StaticSet[String]("make_SymV", "make_IntV", "bv_sext", "bv_zext")
 
     def sExt(bw: Int): Rep[Value] = Unwrap(v) match {
-      case gNode(s, (v1: bExp)::bConst(bw1: Int)::_) if (foldableOp(s) && (bw1 == bw)) => v
-      case gNode("make_IntV", (v1: bExp)::bConst(bw1: Int)::_) if bw > bw1 =>
+      case gNode(s, (v1: bExp)::bConst(bw1: Int)::_) if foldableOp(s) && (bw1 == bw) && Config.opt => v
+      case gNode("make_IntV", (v1: bExp)::bConst(bw1: Int)::_) if bw > bw1 && Config.opt =>
         // sExt(IntV(n, bw1), bw) ⇒ IntV(n, bw)
         IntV(Wrap[Long](v1), bw)
-      case gNode("bv_sext", (v1: bExp)::bConst(bw1: Int)::_) if bw > bw1=>
+      case gNode("bv_sext", (v1: bExp)::bConst(bw1: Int)::_) if bw > bw1 && Config.opt =>
         // sExt(sExt(n, bw1), bw) ⇒ sExt(n, bw)
         Wrap[Value](v1).sExt(bw)
       case _ => "bv_sext".reflectWith[Value](v, bw)
     }
 
     def zExt(bw: Int): Rep[Value] = Unwrap(v) match {
-      case gNode(s, (v1: bExp)::bConst(bw1: Int)::_) if (foldableOp(s) && (bw1 == bw)) => v
-      case gNode("make_IntV", (v1: bExp)::bConst(bw1: Int)::_) if bw > bw1 =>
+      case gNode(s, (v1: bExp)::bConst(bw1: Int)::_) if foldableOp(s) && (bw1 == bw) && Config.opt => v
+      case gNode("make_IntV", (v1: bExp)::bConst(bw1: Int)::_) if bw > bw1 && Config.opt =>
         // zExt(IntV(n, bw1), bw) ⇒ IntV(n, bw)
         IntV(Wrap[Long](v1), bw)
-      case gNode("bv_zext", (v1: bExp)::bConst(bw1: Int)::_) if bw > bw1=>
+      case gNode("bv_zext", (v1: bExp)::bConst(bw1: Int)::_) if bw > bw1 && Config.opt =>
         // zExt(sExt(n, bw1), bw) ⇒ zExt(n, bw)
         Wrap[Value](v1).zExt(bw)
       case _ => "bv_zext".reflectWith[Value](v, bw)
@@ -492,6 +492,7 @@ trait ValueDefs { self: SAIOps with BasicDefs with Opaques =>
     def *(rhs: Rep[Value]): Rep[Value] = IntOp2.mul(v, rhs)
     def &(rhs: Rep[Value]): Rep[Value] = IntOp2("and", v, rhs)
     def |(rhs: Rep[Value]): Rep[Value] = IntOp2("or", v, rhs)
+    def ≡(rhs: Rep[Value]): Rep[Value] = IntOp2.eq(v, rhs)
     def unary_! : Rep[Value] = IntOp1.neg(v)
     def unary_~ : Rep[Value] = IntOp1.bvnot(v)
 
diff --git a/src/main/scala/gensym/engines/ImpCPSEngine.scala b/src/main/scala/gensym/engines/ImpCPSEngine.scala
index 2892872c..9642edcf 100644
--- a/src/main/scala/gensym/engines/ImpCPSEngine.scala
+++ b/src/main/scala/gensym/engines/ImpCPSEngine.scala
@@ -111,6 +111,10 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
   def evalFloatOp2(op: String, lhs: LLVMValue, rhs: LLVMValue, ty: LLVMType, ss: Rep[SS])(implicit ctx: Ctx): Rep[Value] =
     FloatOp2(op, eval(lhs, ty, ss), eval(rhs, ty, ss))
 
+  def selectValue(bb: Rep[BlockLabel], vs: List[() => Rep[Value]], labels: List[BlockLabel]): Rep[Value] =
+    if (bb == labels(0) || labels.length == 1) vs(0)()
+    else selectValue(bb, vs.tail, labels.tail)
+
   def execValueInst(inst: ValueInstruction, ss: Rep[SS])(k: (Rep[Ref[SS]], Rep[Value]) => Rep[Unit])(implicit ctx: Ctx): Rep[Unit] = {
     inst match {
       // Memory Access Instructions
@@ -209,10 +213,6 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
         def fK(s: Rep[Ref[SS]], v: Rep[Value]): Rep[Unit] = { s.pop(stackSize); k(s, v) }
         fv[Ref](ss, List(vs: _*), ContOpt(fK))
       case PhiInst(ty, incs) =>
-        def selectValue(bb: Rep[BlockLabel], vs: List[() => Rep[Value]], labels: List[BlockLabel]): Rep[Value] = {
-          if (bb == labels(0) || labels.length == 1) vs(0)()
-          else selectValue(bb, vs.tail, labels.tail)
-        }
         val incsValues: List[LLVMValue] = incs.map(_.value)
         val incsLabels: List[BlockLabel] = incs.map(i => Counter.block.get(ctx.withBlock(i.label)))
         val vs = incsValues.map(v => () => eval(v, ty, ss))
@@ -255,7 +255,6 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
         Counter.setBranchNum(ctx, 2)
         ss.addIncomingBlock(ctx)
         val cndVal = eval(cnd, ty, ss)
-        //branch(ss, cndVal.toSym, cndVal.toSymNeg, thnLab, elsLab, funName, k)
         if (cndVal.isConc) {
           if (cndVal.int == 1) {
             Coverage.incBranch(ctx, 0)
@@ -268,30 +267,33 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
           symExecBr(ss, cndVal, !cndVal, thnLab, elsLab, k)
         }
       case SwitchTerm(cndTy, cndVal, default, swTable) =>
-        type MergedCase = (String, (Rep[Boolean], Rep[Value]))
+        // Given a case, GSCase stores its label, concrete condition expr,
+        // and symbolic condition expr.
+        case class GSCase(label: String, concCond: Rep[Boolean], symCond: Rep[Value])
         val v = eval(cndVal, cndTy, ss)
-        val mergedSwTable: StaticList[MergedCase] =
+        val mergedSwTable: StaticList[GSCase] =
           swTable.foldLeft(StaticMap[String, (Rep[Boolean], Rep[Value])]())({
             case (m, LLVMCase(_, n, tgt)) if (Config.switchType == Merge) && m.contains(tgt) =>
-              m + (tgt -> (m(tgt)._1 || v.int == n, IntOp2("or", m(tgt)._2, IntOp2("eq", v, IntV(n)))))
+              m + (tgt -> (m(tgt)._1 || v.int == n, m(tgt)._2 | v ≡ IntV(n)))
             case (m, LLVMCase(_, n, tgt)) =>
-              m + (tgt -> (v.int == n, IntOp2("eq", v, IntV(n))))
-          }).toList
+              m + (tgt -> (v.int == n, v ≡ IntV(n)))
+          }).map({ case (l, (c, s)) => GSCase(l, c, s) }).toList
+
         Counter.setBranchNum(ctx, mergedSwTable.size+1)
         System.out.println(s"Shrinking switch table from ${swTable.size+1} cases to ${mergedSwTable.size+1}")
 
         val nPath: Var[Int] = var_new(0)
-        def switch(s: Rep[SS], table: List[MergedCase]): Rep[Unit] = table match {
+        def switch(s: Rep[SS], table: List[GSCase]): Rep[Unit] = table match {
           case Nil =>
             Coverage.incBranch(ctx, mergedSwTable.size) // Note: the default case has the last branch ID
             execBlock(ctx.funName, default, s, k)
-          case (tgt, (concCnd, _))::rest =>
+          case GSCase(tgt, concCnd, _)::rest =>
             if (concCnd) {
               Coverage.incBranch(ctx, mergedSwTable.size - table.size)
               execBlock(ctx.funName, tgt, s, k)
             } else switch(s, table.tail)
         }
-        def switchSym(s: Rep[SS], table: List[MergedCase]): Rep[Unit] = table match {
+        def switchSym(s: Rep[SS], table: List[GSCase]): Rep[Unit] = table match {
           case Nil =>
             if (checkPC(s.pc)) {
               nPath += 1
@@ -299,7 +301,7 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
               Coverage.incBranch(ctx, mergedSwTable.size)
               execBlock(ctx.funName, default, newState, k)
             }
-          case (tgt, (_, symCnd))::rest =>
+          case GSCase(tgt, _, symCnd)::rest =>
             val st = s.copy
             s.addPC(symCnd)
             if (checkPC(s.pc)) {

From 2cb84ff1fab61e34a02c44c9f5e26aa432b30ee1 Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Wed, 19 Oct 2022 13:06:10 -0400
Subject: [PATCH 13/15] misc

---
 src/main/scala/gensym/EngineBase.scala           | 2 +-
 src/main/scala/gensym/Transform.scala            | 2 +-
 src/main/scala/gensym/engines/ImpCPSEngine.scala | 1 +
 src/test/scala/gensym/TestGS.scala               | 5 ++++-
 4 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/src/main/scala/gensym/EngineBase.scala b/src/main/scala/gensym/EngineBase.scala
index ed4d7edf..66d60915 100644
--- a/src/main/scala/gensym/EngineBase.scala
+++ b/src/main/scala/gensym/EngineBase.scala
@@ -158,7 +158,7 @@ trait EngineBase extends SAIOps { self: BasicDefs with ValueDefs =>
   }
 
   // Note: we can also assign symbolic values here
-  def uninitValue: Rep[Value] = IntV(0, 8) //NullPtr()
+  def uninitValue: Rep[Value] = IntV(0, 8)
 
   def evalHeapAtomicConst(v: Constant, ty: LLVMType): Rep[Value] = v match {
     case BoolConst(b) => IntV(if (b) 1 else 0, 1)
diff --git a/src/main/scala/gensym/Transform.scala b/src/main/scala/gensym/Transform.scala
index 1d9acff5..2a05d414 100644
--- a/src/main/scala/gensym/Transform.scala
+++ b/src/main/scala/gensym/Transform.scala
@@ -7,7 +7,7 @@ import lms.core.stub.{While => _, _}
 import gensym.imp.Mut
 
 import scala.collection.immutable.{List => StaticList}
-import scala.collection.mutable.{HashMap,HashSet}
+import scala.collection.mutable.{HashMap, HashSet}
 
 object AssignElim {
   type Subst = HashMap[Sym, Exp]
diff --git a/src/main/scala/gensym/engines/ImpCPSEngine.scala b/src/main/scala/gensym/engines/ImpCPSEngine.scala
index 9642edcf..6b6bd1e6 100644
--- a/src/main/scala/gensym/engines/ImpCPSEngine.scala
+++ b/src/main/scala/gensym/engines/ImpCPSEngine.scala
@@ -100,6 +100,7 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
       case InlineASM() => NullPtr[Value]
       case ZeroInitializerConst =>
         System.out.println("Warning: Evaluate zeroinitialize in body")
+        uninitValue
         NullPtr[Value] // FIXME: use uninitValue
       case NullConst => NullLoc()
       case NoneConst => NullPtr[Value]
diff --git a/src/test/scala/gensym/TestGS.scala b/src/test/scala/gensym/TestGS.scala
index 604a7359..cb0bfacd 100644
--- a/src/test/scala/gensym/TestGS.scala
+++ b/src/test/scala/gensym/TestGS.scala
@@ -126,6 +126,7 @@ class TestImpCPSGS_Z3 extends TestGS {
   testGS(gs, cases)
 }
 
+/*
 class Coreutils extends TestGS {
   import gensym.llvm.parser.Parser._
   Config.enableOpt
@@ -133,16 +134,18 @@ class Coreutils extends TestGS {
   val cases = TestCases.coreutils.map { t =>
     t.copy(runOpt = runtimeOptions ++ t.runOpt, runCode = false)
   }
-  //testGS(new ImpCPSGS, cases)
+  testGS(new ImpCPSGS, cases)
 
   lazy val cat_linked = parseFile("benchmarks/gs_posix/cat_linked.ll")
   testGS(new ImpCPSGS, TestPrg(cat_linked, "cat_linked_posix", "@main", noMainFileOpt, "--argv=./cat.bc --sym-stdout --sym-stdin 2 --sym-arg 2", nPath(28567)++status(0), runCode = false))
 }
+*/
 
 class Playground extends TestGS {
   import gensym.llvm.parser.Parser._
   Config.enableOpt
   val gs = new ImpCPSGS
+  //testGS(gs, TestPrg(assumeTest, "assumeTest", "@main", noArg, noOpt, nPath(1)++status(0)))
 
   //testGS(gs, TestPrg(mergesort, "mergeSortTest", "@main", noArg, noOpt, nPath(720)))
   //testGS(gs, TestPrg(kmpmatcher, "kmp", "@main", noArg, noOpt, nPath(1287)))

From c63ea2a6d0b380d1d9ef494f27d30ae55a18858a Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Wed, 19 Oct 2022 16:35:32 -0400
Subject: [PATCH 14/15] fix

---
 src/main/scala/gensym/engines/ImpCPSEngine.scala | 1 -
 src/test/scala/gensym/TestGS.scala               | 2 +-
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/main/scala/gensym/engines/ImpCPSEngine.scala b/src/main/scala/gensym/engines/ImpCPSEngine.scala
index 6b6bd1e6..f72f70ff 100644
--- a/src/main/scala/gensym/engines/ImpCPSEngine.scala
+++ b/src/main/scala/gensym/engines/ImpCPSEngine.scala
@@ -101,7 +101,6 @@ trait ImpCPSGSEngine extends ImpSymExeDefs with EngineBase {
       case ZeroInitializerConst =>
         System.out.println("Warning: Evaluate zeroinitialize in body")
         uninitValue
-        NullPtr[Value] // FIXME: use uninitValue
       case NullConst => NullLoc()
       case NoneConst => NullPtr[Value]
     }
diff --git a/src/test/scala/gensym/TestGS.scala b/src/test/scala/gensym/TestGS.scala
index cb0bfacd..7d427cd6 100644
--- a/src/test/scala/gensym/TestGS.scala
+++ b/src/test/scala/gensym/TestGS.scala
@@ -147,7 +147,7 @@ class Playground extends TestGS {
   val gs = new ImpCPSGS
   //testGS(gs, TestPrg(assumeTest, "assumeTest", "@main", noArg, noOpt, nPath(1)++status(0)))
 
-  //testGS(gs, TestPrg(mergesort, "mergeSortTest", "@main", noArg, noOpt, nPath(720)))
+  testGS(gs, TestPrg(mergesort, "mergeSortTest", "@main", noArg, noOpt, nPath(720)))
   //testGS(gs, TestPrg(kmpmatcher, "kmp", "@main", noArg, noOpt, nPath(1287)))
   //testGS(new PureCPSGS, TestPrg(arrayFlow, "arrayFlow", "@main", noArg, noOpt, nPath(15)++status(0)))
   //testGS(new ImpCPSGS, TestPrg(arrayFlow, "arrayFlow2", "@main", noArg, noOpt, nPath(15)++status(0)))

From d84e9e90b32cfd8f4f1fd1f13d1d7ff17efb869c Mon Sep 17 00:00:00 2001
From: Guannan Wei <kiss.kraks@gmail.com>
Date: Sat, 22 Oct 2022 23:38:50 -0400
Subject: [PATCH 15/15] fix

---
 src/main/scala/gensym/EngineBase.scala | 2 +-
 third-party/immer                      | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/main/scala/gensym/EngineBase.scala b/src/main/scala/gensym/EngineBase.scala
index b7b5f202..e5c55994 100644
--- a/src/main/scala/gensym/EngineBase.scala
+++ b/src/main/scala/gensym/EngineBase.scala
@@ -60,7 +60,7 @@ trait EngineBase extends SAIOps { self: BasicDefs with ValueDefs =>
   val mainRename = "gs_main"
   val gsPrefix = "__GS_USER_"
 
-  def getRealFunName(funName: String, prefix: String): String = 
+  def getRealFunName(funName: String, prefix: String = gsPrefix): String =
     if (funName != "@main") gsPrefix + funName.tail.replaceAllLiterally(".", "_")
     else mainRename
     
diff --git a/third-party/immer b/third-party/immer
index e5d79ed8..93ab7151 160000
--- a/third-party/immer
+++ b/third-party/immer
@@ -1 +1 @@
-Subproject commit e5d79ed80ec74d511cc4f52fb68feeac66507f2c
+Subproject commit 93ab7151cd7b3482e59fe0e4d43549ac3f1520b4