Use Function0 instead of Suspend

keywords-Await/src/test/scala/com/thoughtworks/dsl/keywords/AwaitTest.scala

@@ -17,6 +17,37 @@ import scala.util.Failure
 import scala.util.Success
 
 class AwaitTest extends AsyncFreeSpec with Matchers with Inside {
+  "div by zero" ignore {
+    val future = *[Future] {
+      !Pure(())
+      !Suspend(() => Pure(0 / 0))
+    }
+    inside(future.value) {
+      case Some(Failure(e)) =>
+        e should be(an[ArithmeticException])
+    }
+  }
+  "Suspend(div by zero)" ignore {
+    val future = *[Future] {
+      !Suspend(() => Pure(()))
+      !Suspend(() => Pure(0 / 0))
+    }
+    inside(future.value) {
+      case Some(Failure(e)) =>
+        e should be(an[ArithmeticException])
+    }
+  }
+
+  "Shift(div by zero)" ignore {
+    val future = *[Future] {
+      !Shift[Future[Int], Unit](_(()))
+      !Suspend(() => Pure(0 / 0))
+    }
+    inside(future.value) {
+      case Some(Failure(e)) =>
+        e should be(an[ArithmeticException])
+    }
+  }
 
   "testReturnIf" in {
     val reified = reify {
@@ -328,15 +359,17 @@ class AwaitTest extends AsyncFreeSpec with Matchers with Inside {
     summon[
       reified.type <:<
         Typed[
-          keywords.TryCatchFinally[Suspend[
-            Await[Future[String]]
-          ], Match.WithIndex[0, FlatMap[
-            Await[Future[Int]],
-            FlatMap[Await[
-              Future[Int]
-            ], Pure[String]]
-          ]]
-            +: Nothing, Suspend[Pure[Unit]]],
+          keywords.TryCatchFinally[
+            Await[Future[String]],
+            Match.WithIndex[0, FlatMap[
+              Await[Future[Int]],
+              FlatMap[Await[
+                Future[Int]
+              ], Pure[String]]
+            ]]
+              +: Nothing,
+            Pure[Unit]
+          ],
           String
         ]
     ]
@@ -356,9 +389,7 @@ class AwaitTest extends AsyncFreeSpec with Matchers with Inside {
 
     summon[
       reified.type <:<
-        Typed[keywords.TryFinally[Suspend[
-          Await[Future[Int]]
-        ], Suspend[Pure[Unit]]], Int]
+        Typed[keywords.TryFinally[Await[Future[Int]], Pure[Unit]], Int]
     ]
 
     reified.as[Future[Int]].transform { t =>
@@ -380,13 +411,11 @@ class AwaitTest extends AsyncFreeSpec with Matchers with Inside {
 
     summon[
       reified.type <:<
-        Typed[keywords.TryCatch[Suspend[
-          FlatMap[
-            Await[Future[Int]],
-            FlatMap[Await[
-              Future[Int]
-            ], Pure[String]]
-          ]
+        Typed[keywords.TryCatch[FlatMap[
+          Await[Future[Int]],
+          FlatMap[Await[
+            Future[Int]
+          ], Pure[String]]
         ], Match.WithIndex[0, FlatMap[
           Await[Future[Int]],
           FlatMap[Await[
@@ -411,14 +440,10 @@ class AwaitTest extends AsyncFreeSpec with Matchers with Inside {
     summon[
       reified.type <:<
         Typed[FlatMap[While[
-          Suspend[
-            Await[Future[Boolean]]
-          ],
-          Suspend[
-            FlatMap[Await[
-              Future[Int]
-            ], Pure[Unit]]
-          ]
+          Await[Future[Boolean]],
+          FlatMap[Await[
+            Future[Int]
+          ], Pure[Unit]]
         ], Pure[Long]], Long]
     ]
     *[Future] {

keywords-If/src/main/scala/com/thoughtworks/dsl/keywords/If.scala

@@ -5,8 +5,8 @@ import Dsl.cpsApply
 
 final case class If[+ConditionKeyword, +ThenKeyword, +ElseKeyword](
     cond: ConditionKeyword,
-    thenp: ThenKeyword,
-    elsep: ElseKeyword
+    thenp: () => ThenKeyword,
+    elsep: () => ElseKeyword
 ) extends Dsl.Keyword.Trait
 
 object If {
@@ -32,9 +32,9 @@ object If {
     ) =>
       keyword.cond.cpsApply {
         case true =>
-          keyword.thenp.cpsApply(handler)
+          keyword.thenp().cpsApply(handler)
         case false =>
-          keyword.elsep.cpsApply(handler)
+          keyword.elsep().cpsApply(handler)
       }
   }
 }

keywords-TryCatch/src/main/scala/com/thoughtworks/dsl/keywords/TryCatch.scala

@@ -6,21 +6,24 @@ import scala.util.control.Exception.Catcher
 import scala.concurrent._
 import scala.util.control.NonFatal
 
-case class TryCatch[+BlockKeyword, +CaseKeyword](block: BlockKeyword, cases: Catcher[CaseKeyword]) extends Dsl.Keyword.Trait
+case class TryCatch[+BlockKeyword, +CaseKeyword](
+    block: () => BlockKeyword,
+    cases: Catcher[CaseKeyword]
+) extends Dsl.Keyword.Trait
 object TryCatch {
 
-  given [Value, OuterDomain, BlockKeyword, BlockDomain, CaseKeyword](
-      using
+  given [Value, OuterDomain, BlockKeyword, BlockDomain, CaseKeyword](using
       dslTryCatch: Dsl.TryCatch[Value, OuterDomain, BlockDomain],
       blockDsl: Dsl.Searching[BlockKeyword, BlockDomain, Value],
-      caseDsl: Dsl.Searching[CaseKeyword, BlockDomain, Value],
-  ): Dsl.Composed[TryCatch[BlockKeyword, CaseKeyword], OuterDomain, Value] = Dsl.Composed {
-    case (TryCatch(blockKeyword, cases), handler) =>
+      caseDsl: Dsl.Searching[CaseKeyword, BlockDomain, Value]
+  ): Dsl.Composed[TryCatch[BlockKeyword, CaseKeyword], OuterDomain, Value] =
+    Dsl.Composed { case (TryCatch(blockKeyword, cases), handler) =>
       dslTryCatch.tryCatch(
-        blockDsl(blockKeyword, _),
+        // TODO: Use Suspend to catch the exception
+        blockDsl(blockKeyword(), _),
         cases.andThen { caseKeyword => caseDsl(caseKeyword, _) },
         handler
       )
-  }
+    }
 
 }

keywords-TryCatchFinally/src/main/scala/com/thoughtworks/dsl/keywords/TryCatchFinally.scala

@@ -5,9 +5,9 @@ import Dsl.IsKeyword
 import scala.util.control.Exception.Catcher
 
 case class TryCatchFinally[+BlockKeyword, +CaseKeyword, +FinalizerKeyword](
-    block: BlockKeyword,
+    block: () => BlockKeyword,
     cases: Catcher[CaseKeyword],
-    finalizer: FinalizerKeyword
+    finalizer: () => FinalizerKeyword
 ) extends Dsl.Keyword.Trait
 object TryCatchFinally {
 
@@ -20,9 +20,10 @@ object TryCatchFinally {
   ): Dsl.Composed[TryCatchFinally[BlockKeyword, CaseKeyword, FinalizerKeyword], OuterDomain, Value] = Dsl.Composed {
     case (TryCatchFinally(blockKeyword, cases, finalizerKeyword), handler) =>
       dslTryCatchFinally.tryCatchFinally(
-        blockDsl(blockKeyword, _),
+        // TODO: Use Suspend to catch the exception
+        blockDsl(blockKeyword(), _),
         cases.andThen { caseKeyword => caseDsl(caseKeyword, _) },
-        finalizerDsl(finalizerKeyword, _),
+        finalizerDsl(finalizerKeyword(), _),
         handler
       )
   }

keywords-TryFinally/src/main/scala/com/thoughtworks/dsl/keywords/TryFinally.scala

@@ -7,8 +7,8 @@ import scala.concurrent._
 import scala.util.control.NonFatal
 
 case class TryFinally[+TryKeyword, +FinalizerKeyword](
-    block: TryKeyword,
-    finalizer: FinalizerKeyword
+    block: () => TryKeyword,
+    finalizer: () => FinalizerKeyword
 ) extends Dsl.Keyword.Trait
 
 object TryFinally {
@@ -20,8 +20,9 @@ object TryFinally {
   ): Dsl.Composed[TryFinally[BlockKeyword, FinalizerKeyword], OuterDomain, Value] = Dsl.Composed {
     case (TryFinally(blockKeyword, finalizerKeyword), handler) =>
       dslTryFinally.tryFinally(
-        blockDsl(blockKeyword, _),
-        finalizerDsl(finalizerKeyword, _),
+        // TODO: Use Suspend to catch the exception
+        blockDsl(blockKeyword(), _),
+        finalizerDsl(finalizerKeyword(), _),
         handler
       )
   }

keywords-While/src/main/scala/com/thoughtworks/dsl/keywords/While.scala

@@ -7,8 +7,8 @@ case class While[
     +ConditionKeyword,
     +BodyKeyword
 ](
-    condition: ConditionKeyword,
-    body: BodyKeyword
+    condition: () => ConditionKeyword,
+    body: () => BodyKeyword
 ) extends Dsl.Keyword.Trait
 object While {
   given [
@@ -24,9 +24,10 @@ object While {
     Unit
   ] = Dsl.Composed {
     (keyword: While[ConditionKeyword, BodyKeyword], handler: Unit => Domain) =>
-      keyword.condition.cpsApply {
+      // TODO: Use Suspend to rewind the stack in case of stack overflow
+      keyword.condition().cpsApply {
         case true =>
-          keyword.body.cpsApply { _ =>
+          keyword.body().cpsApply { _ =>
             keyword.cpsApply(handler)
           }
         case false =>