From f5783f2b2beffc1f825f83878f73a5f139761dc9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Th=C3=A9ophile=20Bastian?= <contact@tobast.fr>
Date: Fri, 16 Feb 2018 15:29:36 +0100
Subject: [PATCH] Light CPS: implement lambdas

---
 src/CPS.ml | 25 +++++++++++++++----------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/src/CPS.ml b/src/CPS.ml
index 68f1cc9..77eaca0 100644
--- a/src/CPS.ml
+++ b/src/CPS.ml
@@ -29,9 +29,9 @@ let letCont name varName body next =
 
 let rec has_calls (t: S.term): bool = match t with
 | S.Var _ | S.Lit _ | S.BinOp _ -> false
-| S.Lam _ -> true (* TODO *)
-    (* A lambda itself may contain calls, but this call is not evaluated at
-     * declaration time *)
+| S.Lam _ -> false
+        (* A lambda itself may contain calls, but this call is not evaluated at
+         * declaration time *)
 | S.App _ -> true
 | S.IfZero (cond, tIf, tElse) ->
         (* Cannot optimize continuation creation
@@ -60,12 +60,10 @@ let rec cps_term_inner (t: S.term) (cont: T.variable) (nameHint: string option)
 | S.Var _ -> cps_value_as_term t cont
 | S.Lit _ -> cps_value_as_term t cont
 | S.BinOp _ -> cps_value_as_term t cont
-| S.Lam (self, var, term) ->
-        let fName = freshBlockVarHinted nameHint
-        and innerCont = freshBlockVar () in
-        T.LetBlo(fName,
-            T.Lam(self, [var; innerCont], cps_term_inner term innerCont None),
-            T.TailCall(T.vvar cont, T.vvars [fName]))
+| S.Lam _ as lambda ->
+        let fName = freshBlockVarHinted nameHint in
+        light_term fName lambda None @@
+            T.TailCall(T.vvar cont, T.vvars [fName])
 | S.App (f, x) ->
         let xVal = freshVarHinted nameHint
         and fVal = freshVar () in
@@ -108,7 +106,14 @@ and light_term varName valExpr valHint next =
                     subLetVar,
                     cps_value subLetVal,
                     light_term varName subLetNext valHint next)
-        | S.Lam _ | S.App _ | S.Print _ | S.IfZero _ -> assert false
+        | S.Lam(self, lamVar, lamBody) ->
+                let lamCont = freshBlockVar () in
+                T.LetBlo (
+                    varName, T.Lam(
+                        self, [lamVar; lamCont],
+                        cps_term_inner lamBody lamCont None),
+                    next)
+        | S.App _ | S.Print _ | S.IfZero _ -> assert false
     )