path: root/discocaml/eval.ml

open Ast

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
    | Bool _ | Int _ | Nil -> ()
    | Cons (hd, tl) ->
        loop hd;
        loop tl
    | If (cond, then_, else_) ->
        loop cond;
        loop then_;
        loop else_
    | 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
    | Prim (RelOp, (_, 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 bump_var (name : string) : string =
  match Util.break_to_subscript name with
  | Some (base, i) -> Format.sprintf "%s%d" base (i + 1)
  | None -> name ^ "0"

let%test _ = bump_var "x" = "x0"
let%test _ = bump_var "x0" = "x1"
let%test _ = bump_var "x123" = "x124"

let get_fresh_in (ast : expr ast) (is : expr index list) : string -> string =
  let rec loop (name : string) : string =
    if List.exists (fun i -> has_free_var ast i name) is then
      loop (bump_var name)
    else name
  in
  loop

let%expect_test "get_fresh_in" =
  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") and x0 = add (Var "x0") and x1 = add (Var "x1") in

  let xx0 = add (Prim (Add, (x, x0))) in
  let xx0x1 = add (Prim (Add, (xx0, x1))) in

  Printf.printf "%S\n" (get_fresh_in ast [ xx0x1 ] "x");
  Printf.printf "%S\n" (get_fresh_in ast [ xx0 ] "x");
  Printf.printf "%S\n" (get_fresh_in ast [ x; x1 ] "x");
  Printf.printf "%S\n" (get_fresh_in ast [ x; x1 ] "y");

  [%expect {|
    "x2"
    "x1"
    "x0"
    "y"
  |}]

exception FoundRedex of expr index

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) | _ -> ())
  | If (cond, _, _) -> (
      match get_subexpr ast cond with Bool _ -> 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)
      | _ -> ())
  | Prim (RelOp, (_, l, r)) -> (
      match (get_subexpr ast l, get_subexpr ast r) with
      | Bool _, Bool _ -> raise (FoundRedex i)
      | Int _, Int _ -> raise (FoundRedex i)
      | _ -> ())
  | Var _ -> ()
  | Bool _ | Cons _ | Int _ | Lam _ | Nil -> ()

let find_redex_cbv_in (ast : expr ast) : expr index -> unit =
  let rec loop (i : expr index) : unit =
    (match get_subexpr ast i with
    | App (f, x) ->
        loop f;
        loop x
    | If (cond, _, _) -> loop cond
    | Let (false, _, bound, _) -> loop bound
    | Let (true, _, bound, body) ->
        loop bound;
        loop body
    | Prim (Add, (l, r)) | Prim (Sub, (l, r)) | Prim (Mul, (l, r)) ->
        loop l;
        loop r
    | Prim (RelOp, (_, l, r)) ->
        loop l;
        loop r
    | Var _ -> ()
    | Bool _ | Cons _ | Int _ | Lam _ | Nil -> ());
    check_redex ast i
  in
  loop

let find_redex_cbn_in (ast : expr ast) : expr index -> unit =
  let rec loop (i : expr index) : unit =
    check_redex ast i;
    match get_subexpr ast i with
    | App (f, x) ->
        loop f;
        loop x
    | If (cond, _, _) -> loop cond
    | Let (_, _, _, body) -> loop body
    | Prim (Add, (l, r)) | Prim (Sub, (l, r)) | Prim (Mul, (l, r)) ->
        loop l;
        loop r
    | Prim (RelOp, (_, l, r)) ->
        loop l;
        loop r
    | Var _ -> ()
    | Bool _ | Cons _ | Int _ | Lam _ | Nil -> ()
  in
  loop

let find_redex_cbv (ast : expr ast) : expr index option =
  try
    find_redex_cbv_in ast ast.root;
    None
  with FoundRedex i -> Some i

let find_redex_cbn (ast : expr ast) : expr index option =
  try
    find_redex_cbn_in ast ast.root;
    None
  with FoundRedex i -> Some i

exception NotARedex of expr ast

let rec subst (ast : expr ast) (from : string) (to_ : expr index) :
    expr index -> expr index =
  let rec loop (i : expr index) : expr index =
    let add (expr : expr) : expr index =
      if get_subexpr ast i = expr then i
      else
        let index = Arraylist.length ast.subexprs in
        Arraylist.push ast.subexprs expr;
        { index }
    in
    match get_subexpr ast i with
    | App (f, x) -> add (App (loop f, loop x))
    | Bool _ | Int _ | Nil -> i
    | Cons (hd, tl) -> add (Cons (loop hd, loop tl))
    | If (cond, then_, else_) -> add (If (loop cond, loop then_, loop else_))
    | Lam (x, b) ->
        if String.equal from x then i
        else if has_free_var ast to_ x then
          let x' = get_fresh_in ast [ b; to_ ] x in
          add (Lam (x', loop (subst ast x (add (Var x')) b)))
        else add (Lam (x, loop b))
    | Let (false, name, bound, body) ->
        if String.equal from name then add (Let (false, name, loop bound, body))
        else if has_free_var ast to_ name then
          let name' = get_fresh_in ast [ body; to_ ] name in
          let body' = subst ast name (add (Var name')) body in
          add (Let (false, name', loop bound, loop body'))
        else add (Let (false, name, loop bound, loop body))
    | Let (true, name, bound, body) ->
        if String.equal from name then i
        else if has_free_var ast to_ name then
          let name' = get_fresh_in ast [ bound; body; to_ ] name in
          let bound' = subst ast name (add (Var name')) bound
          and body' = subst ast name (add (Var name')) body in
          add (Let (true, name', loop bound', loop body'))
        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)))
    | Prim (RelOp, (op, l, r)) -> add (Prim (RelOp, (op, loop l, loop r)))
    | Var x -> if String.equal from x then to_ else i
  in
  loop

let reduce (ast : expr ast) (i : expr index) : expr ast =
  let fail () = raise (NotARedex { ast with root = i }) in
  let ast = copy ast in
  let must_int j = match get_subexpr ast j with Int n -> n | _ -> fail () in
  Arraylist.set ast.subexprs i.index
    (match get_subexpr ast i with
    | App (f, x) -> (
        match get_subexpr ast f with
        | Lam (x', b) -> get_subexpr ast (subst ast x' x b)
        | _ -> fail ())
    | If (cond, then_, else_) -> (
        match get_subexpr ast cond with
        | Bool true -> get_subexpr ast then_
        | Bool false -> get_subexpr ast else_
        | _ -> 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)
    | Prim (RelOp, (op, l, r)) ->
        let l, r =
          match (get_subexpr ast l, get_subexpr ast r) with
          | Bool l, Bool r -> (Obj.repr l, Obj.repr r)
          | Int l, Int r -> (Obj.repr l, Obj.repr r)
          | _ -> fail ()
        in
        Bool
          (match op with
          | `LT -> l < r
          | `LTE -> l <= r
          | `GT -> l > r
          | `GTE -> l >= r
          | `EQ -> l = r
          | `NE -> l <> r)
    | _ -> fail ());
  ast