let expressions

3 files changed, 108 insertions, 0 deletions

diff --git a/discocaml/ast.ml b/discocaml/ast.ml
index 9436b69..025fa24 100644
--- a/discocaml/ast.ml
+++ b/discocaml/ast.ml
@@ -23,6 +23,7 @@ and expr =
   | Cons of expr index * expr index
   | Int of int
   | Lam of string * expr index
+  | Let of bool * string * expr index * expr index
   | Nil
   | Prim : 'a prim * 'a -> expr
   | Var of string
@@ -85,6 +86,20 @@ let add_expr_to_ast (ast : 'a ast) : Parsetree.expression -> expr index =
     | Pexp_ident { txt = Lident "-"; _ } -> binop Sub
     | Pexp_ident { txt = Lident "*"; _ } -> binop Mul
     | Pexp_ident { txt = Lident name; _ } -> add (Var name)
+    | Pexp_let
+        ( recursive,
+          [
+            {
+              pvb_pat = { ppat_desc = Ppat_var { txt = name; _ }; _ };
+              pvb_expr;
+              _;
+            };
+          ],
+          body ) ->
+        let recursive =
+          match recursive with Nonrecursive -> false | Recursive -> true
+        in
+        add (Let (recursive, name, loop pvb_expr, loop body))
     | _ -> raise_unsupported_expr expr
   in
   loop
@@ -154,6 +169,20 @@ let parsetree_of_subexpr (ast : 'a ast) : expr -> Parsetree.expression =
                None,
                Wrap.pattern (Ppat_var (Wrap.var x)),
                loop (subexpr b) ))
+    | Let (recursive, name, bound, body) ->
+        Wrap.expression
+          (Pexp_let
+             ( (if recursive then Recursive else Nonrecursive),
+               [
+                 {
+                   pvb_pat = Wrap.pattern (Ppat_var (Wrap.var name));
+                   pvb_expr = loop (subexpr bound);
+                   pvb_constraint = None;
+                   pvb_attributes = [];
+                   pvb_loc = Location.none;
+                 };
+               ],
+               loop (subexpr body) ))
     | Nil -> Wrap.expression (Pexp_construct (Wrap.ident "[]", None))
     | Prim (Add, (l, r)) ->
         Wrap.expression
diff --git a/discocaml/draw_tree.ml b/discocaml/draw_tree.ml
index 4b1688b..3429a71 100644
--- a/discocaml/draw_tree.ml
+++ b/discocaml/draw_tree.ml
@@ -13,6 +13,9 @@ let add_node (fmt : Format.formatter) (i : expr index) (expr : expr) : unit =
         i.index
   | Int n -> Format.fprintf fmt "  expr%d [label=\"%d\"];\n" i.index n
   | Lam _ -> Format.fprintf fmt "  expr%d [label=\"λ\"];\n" i.index
+  | Let (recursive, _, _, _) ->
+      Format.fprintf fmt "  expr%d [label=\"%s\"];\n" i.index
+        (if recursive then "letrec" else "let")
   | Nil ->
       Format.fprintf fmt
         "  expr%d [fontname=\"CMU Typewriter Text Bold\", label=\"[]\"];\n"
@@ -41,6 +44,11 @@ let add_expr_edges (ast : 'a ast) (fmt : Format.formatter)
         Format.fprintf fmt "  expr%d -> expr%d_var;\n" i.index i.index;
         Format.fprintf fmt "  expr%d_var [label=%S];\n" i.index x;
         edge_to b
+    | Let (_, name, bound, body) ->
+        Format.fprintf fmt "  expr%d -> expr%d_var;\n" i.index i.index;
+        Format.fprintf fmt "  expr%d_var [label=%S];\n" i.index name;
+        edge_to bound;
+        edge_to body
     | Prim (Add, (l, r)) ->
         edge_to l;
         edge_to r
diff --git a/discocaml/eval.ml b/discocaml/eval.ml
index 3957531..81709e6 100644
--- a/discocaml/eval.ml
+++ b/discocaml/eval.ml
@@ -2,10 +2,69 @@ open Ast
 
 exception FoundRedex of expr index
 
+let has_free_var (ast : expr ast) (i : expr index) (name : string) : bool =
+  let exception FoundIt in
+  let rec loop (i : expr index) =
+    match get_subexpr ast i with
+    | App (f, x) ->
+        loop f;
+        loop x
+    | Cons (hd, tl) ->
+        loop hd;
+        loop tl
+    | Int _ | Nil -> ()
+    | Lam (x, b) -> if not (String.equal name x) then loop b
+    | Let (recursive, name', bound, body) ->
+        if not (String.equal name name') then loop body;
+        if (not recursive) || not (String.equal name name') then loop bound
+    | Prim (Add, (l, r)) ->
+        loop l;
+        loop r
+    | Prim (Sub, (l, r)) ->
+        loop l;
+        loop r
+    | Prim (Mul, (l, r)) ->
+        loop l;
+        loop r
+    | Var x -> if String.equal name x then raise FoundIt
+  in
+  try
+    loop i;
+    false
+  with FoundIt -> true
+
+let%expect_test "has_free_var shadowing tests" =
+  let ast = new_ast () in
+  let add (expr : expr) : expr index =
+    let index = Arraylist.length ast.subexprs in
+    Arraylist.push ast.subexprs expr;
+    { index }
+  in
+
+  let x = add (Var "x") in
+  let three = add (Int 3) in
+
+  let letx3inx = add (Let (false, "x", three, x)) in
+  Printf.printf "let     x = 3 in x: %b\n" (has_free_var ast letx3inx "x");
+
+  let letxxinx = add (Let (false, "x", x, x)) in
+  Printf.printf "let     x = x in x: %b\n" (has_free_var ast letxxinx "x");
+
+  let letrecxxinx = add (Let (true, "x", x, x)) in
+  Printf.printf "let rec x = x in x: %b\n" (has_free_var ast letrecxxinx "x");
+
+  [%expect
+    {|
+    let     x = 3 in x: false
+    let     x = x in x: true
+    let rec x = x in x: false
+  |}]
+
 let check_redex (ast : expr ast) (i : expr index) : unit =
   match get_subexpr ast i with
   | App (f, _) -> (
       match get_subexpr ast f with Lam _ -> raise (FoundRedex i) | _ -> ())
+  | Let (_, _, _, _) -> raise (FoundRedex i)
   | Prim (Add, (l, r)) | Prim (Sub, (l, r)) | Prim (Mul, (l, r)) -> (
       match (get_subexpr ast l, get_subexpr ast r) with
       | Int _, Int _ -> raise (FoundRedex i)
@@ -19,6 +78,7 @@ let find_redex_cbv_in (ast : expr ast) : expr index -> unit =
     | App (f, x) ->
         loop f;
         loop x
+    | Let (_, _, bound, _) -> loop bound
     | Prim (Add, (l, r)) | Prim (Sub, (l, r)) | Prim (Mul, (l, r)) ->
         loop l;
         loop r
@@ -35,6 +95,7 @@ let find_redex_cbn_in (ast : expr ast) : expr index -> unit =
     | App (f, x) ->
         loop f;
         loop x
+    | Let (_, _, _, body) -> loop body
     | Prim (Add, (l, r)) | Prim (Sub, (l, r)) | Prim (Mul, (l, r)) ->
         loop l;
         loop r
@@ -57,6 +118,7 @@ let find_redex_cbn (ast : expr ast) : expr index option =
 
 exception NotARedex of expr ast
 
+(* TODO: Be capture-avoiding! *)
 let subst (ast : expr ast) (from : string) (to_ : expr index) :
     expr index -> expr index =
   let rec loop (i : expr index) : expr index =
@@ -72,6 +134,11 @@ let subst (ast : expr ast) (from : string) (to_ : expr index) :
     | Cons (hd, tl) -> add (Cons (loop hd, loop tl))
     | Int _ | Nil -> i
     | Lam (x, b) -> if String.equal from x then i else add (Lam (x, loop b))
+    | Let (false, name, bound, body) ->
+        add
+          (Let (false, name, loop bound, if name = from then body else loop body))
+    | Let (true, name, bound, body) ->
+        if name = from then i else add (Let (true, name, loop bound, loop body))
     | Prim (Add, (l, r)) -> add (Prim (Add, (loop l, loop r)))
     | Prim (Sub, (l, r)) -> add (Prim (Sub, (loop l, loop r)))
     | Prim (Mul, (l, r)) -> add (Prim (Mul, (loop l, loop r)))
@@ -89,6 +156,10 @@ let reduce (ast : expr ast) (i : expr index) : expr ast =
         match get_subexpr ast f with
         | Lam (x', b) -> get_subexpr ast (subst ast x' x b)
         | _ -> fail ())
+    | Let (_, name, _, body) when not (has_free_var ast body name) ->
+        get_subexpr ast body
+    | Let (recursive, name, bound, body) ->
+        Let (recursive, name, bound, subst ast name bound body)
     | Prim (Add, (l, r)) -> Int (must_int l + must_int r)
     | Prim (Sub, (l, r)) -> Int (must_int l - must_int r)
     | Prim (Mul, (l, r)) -> Int (must_int l * must_int r)