free.ml

let (=>) left right = print_char (if left = right then '.' else 'F')


module Env = struct
  include Set.Make (String)
  let to_list t = fold List.cons t []
end



type id = string


type t =
  | Var of id
  | Const of int
  | Apply of t * t
  | Function of id * t
  | Let of id * t * t
  | If of t * t * t
  | Sequence of t * t

(*type 'a t =
  | Var of id
  | Const of int
  | Apply of 'a * 'a
  | Function of id * 'a
  | Let of id * 'a * 'a
  | If of 'a * 'a * 'a
  | Sequence of 'a * 'a*)
  [@@deriving map, fold]

let map f = function
  | Var _ | Const _ as t -> t
  | Apply (t1, t2) -> Apply (f t1, f t2)
  | Function (id, t) -> Function (id, f t)
  | Let (id, t1, t2) -> Let (id, f t1, f t2)
  | If (t1, t2, t3) -> If (f t1, f t2, f t3)
  | Sequence (t1, t2) -> Sequence (f t1, f t2)

(*let rec fold_right ~init ~f = function
  | Var _ | Const _ -> init
  | Apply (t1, t2) -> f t1 t2
  | Function (_, t) -> t
  | Let (_, t1, t2) -> f t1 t2
  | If (t1, t2, t3) -> f t1 (f t2 t3)
  | Sequence (t1, t2) -> f t1 t2*)

(* ('a -> 'b -> 'a) -> 'a -> 'b btree -> 'a *)

let rec fold f init = function
  | Var _ | Const _ -> init
  | Apply (t1, t2) -> f (f init t1) t2
  | Function (_, t) -> f init t
  | Let (_, t1, t2) -> f (f init t1) t2
  | If (t1, t2, t3) -> f (f (f init t1) t2) t3
  | Sequence (t1, t2) -> f (f init t1) t2

let rec fold_right ~init ~f = function
  | Var _ | Const _ -> init
  | Apply (t1, t2) -> f t1 (f t2 init)
  | Function (_, t) -> f t init
  | Let (_, t1, t2) -> f t1 (f t2 init)
  | If (t1, t2, t3) -> f t1 (f t2 (f t3 init))
  | Sequence (t1, t2) -> f t1 (f t2 init)

let (%) left right = Sequence (left, right)

let rec cata f t = f (map (cata f) t)

module With_map = struct
  let rec pass = function
    | If (Const 0, _, t) -> pass t
    | If (Const 1, t, _) -> pass t
    | other -> map pass other
end

module With_map_wrong = struct
  let rec pass = map @@ function
    | If (Const 0, _, t) -> t
    | If (Const 1, t, _) -> t
    | other -> other
end

module With_cata = struct
  let pass = cata @@ function
    | If (Const 0, _, t) -> t
    | If (Const 1, t, _) -> t
    | other -> other
end

module With_cata_wrong = struct
  let rec pass = function
    | If (Const 0, _, t) -> t
    | If (Const 1, t, _) -> t
    | other -> cata pass other
end


module Test = struct
  let input =
    If (Const 0, Var "dead",
      Sequence (Var "live",
        If (Const 1, Var "live", Var "dead")))  in
  let expected = Sequence (Var "live", Var "live") in

  With_map.pass input => expected;
(*  With_map_wrong.pass input => expected; *)
  With_cata.pass input => expected;
(*  With_cata_wrong.pass input => expected; *)

end




let rec free = function
  | Var id -> Env.singleton id
  | Const _ -> Env.empty
  | Apply (f, x) -> Env.union (free f) (free x)
  | Function (id, t) -> Env.remove id (free t)
  | Let (id, term, body) ->
      Env.union (free term) (Env.remove id (free body))
  | If (cond, cons, alt) ->
      Env.union (free cond) (Env.union (free cons) (free alt))
  | Sequence (left, right) ->
      Env.union (free left) (free right)

(*let free' = cata @@ function
  | Var id -> Env.singleton id
  | Function (id, t) -> Env.remove id t
  | Let (id, term, body) -> Env.union term (Env.remove id body)
  | other -> fold_right ~init:Env.empty ~f:Env.union other*)

module Free_with_map = struct
  let rec pass = function
    | Var id -> Env.singleton id
    | Function (id, t) -> Env.remove id (pass t)
    | Let (id, term, body) -> Env.union (pass term) (Env.remove id (pass body))
    | other ->
        fold (fun env t -> Env.union (pass t) env) Env.empty other
end



module Test_free = struct
  let input =
    Var "x" %
      Let ("x", Const 0,
        Let ("a", Const 0,
          Var "a" %
            Let ("b", Var "a",
              Let ("y", Var "y",
                Let ("z", Const 0, Const 0) %
                  Var "a" % Var "b" % Var "z")))) in
  Env.to_list (free input) => ["z"; "y"; "x"];
  (*Env.to_list (free' input) => ["z"; "y"; "x"];*)
  Env.to_list (Free_with_map.pass input) => ["z"; "y"; "x"];

end



let () = print_endline "."









(*
let rec fold ~nil ~cons = function
  | Var _id -> nil
  | Const _ -> nil
  | Apply (f, x) -> cons (fold ~nil ~cons f) (fold ~nil ~cons x)
  | Function (id, t) -> (fold ~nil ~cons t)
  | Let (_id, term, body) ->
      cons (fold ~nil ~cons term) (fold ~nil ~cons body)
  | If (condition, consequence, alternative) ->
      cons (fold ~nil ~cons condition)
        (cons (fold ~nil ~cons consequence) (fold ~nil ~cons alternative))
  | Sequence (left, right) ->
      cons (fold ~nil ~cons left) (fold ~nil ~cons right)*)