map-as-a-recursion-scheme.ml

(*
 * This file can be executed by running:
 * $ ocaml -rectypes map-as-a-recursion-scheme.ml
 *
 *)

open Printf


(*

e -> ()
   | (e)
   | true | false
   | 0 | 1 | 2 | …
   | id
   | e / e
   | e; e
   | let e = e in e
   | if e then e else e

*)

(* 1. Naive recursion *)

module Syntax = struct
  type t =
    | Unit
    | Boolean of bool
    | Number of int
    | Id of string
    | Divide of t * t
    | Sequence of t * t
    | Let of {id: string; value: t; body: t}
    | If of {conditional: t; consequence: t; alternative: t}
end

open Syntax

module Dead_code_elimination = struct
  let rec pass = function
    | Unit | Boolean _ | Number _ | Id _ as t ->
        t
    | Divide (left, right) ->
        Divide (pass left, pass right)
    | Sequence (left, right) ->
        Sequence (pass left, pass right)
    | Let {id; value; body} ->
        Let {id; value=pass value; body=pass body}
    | If {conditional=Boolean true; consequence; _} ->
        pass consequence
    | If {conditional=Boolean false; alternative; _} ->
        pass alternative
    | If {conditional; consequence; alternative} ->
        let conditional = pass conditional in
        let consequence = pass consequence in
        let alternative = pass alternative in
        If {conditional; consequence; alternative}
end

(* 2. Factored recursion *)

let map f = function
  | Unit | Boolean _ | Number _ | Id _ as t ->
      t
  | Divide (left, right) ->
      Divide (f left, f right)
  | Sequence (left, right) ->
      Sequence (f left, f right)
  | Let {id; value; body} ->
      Let {id; value=f value; body=f body}
  | If {conditional; consequence; alternative} ->
      let conditional = f conditional in
      let consequence = f consequence in
      let alternative = f alternative in
      If {conditional; consequence; alternative}

module Dead_code_elimination_2 = struct
  let rec pass = function
    | If {conditional=Boolean true; consequence; _} ->
        pass consequence
    | If {conditional=Boolean false; alternative; _} ->
        pass alternative
    | other -> map pass other
end

(* 3. Factored recursion for free *)

module Syntax2 = struct
  module Open = struct
    type 'a t =
      | Unit
      | Boolean of bool
      | Number of int
      | Id of string
      | Divide of 'a * 'a
      | Sequence of 'a * 'a
      | Let of {id: string; value: 'a; body: 'a}
      | If of {conditional: 'a; consequence: 'a; alternative: 'a}
      [@@deriving map]

    let map f = function
      | Unit | Boolean _ | Number _ | Id _ as t ->
          t
      | Divide (left, right) ->
          Divide (f left, f right)
      | Sequence (left, right) ->
          Sequence (f left, f right)
      | Let {id; value; body} ->
          Let {id; value=f value; body=f body}
      | If {conditional; consequence; alternative} ->
          let conditional = f conditional in
          let consequence = f consequence in
          let alternative = f alternative in
          If {conditional; consequence; alternative}
  end

  type t = t Open.t

  let map: (t -> t) -> t -> t = Open.map
end

open Syntax2.Open



module Dead_code_elimination_3 = struct
  let rec pass = function
    | If {conditional=Boolean true; consequence; _} ->
        pass consequence
    | If {conditional=Boolean false; alternative; _} ->
        pass alternative
    | other -> map pass other
end

(* 4. Monadic passes *)

module Identity = struct
  let return x = x

  let (>>=) t f = f t
end

module Result = struct
  let return x = Ok x

  let (>>=) = function
    | Ok ok -> fun f -> f ok
    | error -> fun _ -> error
end

open Result

module Check_literal_division_by_zero = struct
  let rec pass = function
    | Unit | Boolean _ | Number _ | Id _ as t ->
        return t
    | Divide (_, Number 0) ->
        Error `Literal_division_by_zero
    | Divide (left, right) ->
        pass left >>= fun left ->
        pass right >>= fun right ->
        return (Divide (left, right))
    | Sequence (left, right) ->
        pass left >>= fun left ->
        pass right >>= fun right ->
        return (Sequence (left, right))
    | Let {id; value; body} ->
        pass value >>= fun value ->
        pass body >>= fun body ->
        return (Let {id; value; body})
    | If {conditional; consequence; alternative} ->
        pass conditional >>= fun conditional ->
        pass consequence >>= fun consequence ->
        pass alternative >>= fun alternative ->
        return (If {conditional; consequence; alternative})
end

module Syntax3 = struct
  module Open = struct
    type 'a t =
      | Unit
      | Boolean of bool
      | Number of int
      | Id of string
      | Divide of 'a * 'a
      | Sequence of 'a * 'a
      | Let of {id: string; value: 'a; body: 'a}
      | If of {conditional: 'a; consequence: 'a; alternative: 'a}

    let map_monad ~return ~bind:(>>=) f = function
      | Unit | Boolean _ | Number _ | Id _ as t ->
          return t
      | Divide (left, right) ->
          f left >>= fun left ->
          f right >>= fun right ->
          return (Divide (left, right))
      | Sequence (left, right) ->
          f left >>= fun left ->
          f right >>= fun right ->
          return (Sequence (left, right))
      | Let {id; value; body} ->
          f value >>= fun value ->
          f body >>= fun body ->
          return (Let {id; value; body})
      | If {conditional; consequence; alternative} ->
          f conditional >>= fun conditional ->
          f consequence >>= fun consequence ->
          f alternative >>= fun alternative ->
          return (If {conditional; consequence; alternative})
  end

  type t = t Open.t

end

(* open Syntax *)
open Syntax3.Open

let map = map_monad ~return:Identity.return ~bind:Identity.(>>=)

module Dead_code_elimination_4 = struct
  let rec pass = function
    | If {conditional=Boolean true; consequence; _} ->
        pass consequence
    | If {conditional=Boolean false; alternative; _} ->
        pass alternative
    | other -> map pass other
end


let map_result = map_monad ~return:Result.return ~bind:Result.(>>=)

module Check_literal_division_by_zero_2 = struct
  let rec pass = function
    | Divide (_, Number 0) -> Error "`Literal_division_by_zero"
    | other -> map_result pass other
end


module Environment = struct
  type t = {defined: string list; undefined: string list}

  let initial = {defined=[]; undefined=[]}

  let print {defined; undefined} =
    let list = String.concat "; " in
    printf "{defined=[%s]; undefined=[%s]}\n" (list defined) (list undefined)
end

open Environment

module Monad = struct
  type env = t
  type 'a t = env -> 'a * env

  let return a env = a, env

  let (>>=) t callback env =
    let a, env = t env in
    (* print env; *)
    callback a env

  let with_defined id t {defined; undefined} =
    let a, env = t {undefined; defined=id :: defined} in
    a, {env with defined}

  let check_id id env =
    if List.mem id env.defined then
      (), env
    else
      (), {env with undefined=id :: env.undefined}

  let undefined t = (snd (t initial)).undefined
end

let map_environment =
  map_monad ~return:Monad.return ~bind:Monad.(>>=)

open Monad

module Collect_undefined_variables = struct
  let rec pass = function
    | Let {id; value; body} ->
        pass value >>= fun value ->
        with_defined id (pass body) >>= fun body ->
        return (Let {id; value; body})
    | Id id as t ->
        check_id id >>= fun () ->
        return t
    | other -> map_environment pass other

  module Test = struct

    let data = (* let hai = bye in hai *)
      Let {id="hai"; value=Id "bye"; body=Id "hai"}

    let () = printf "1.\n"; assert (undefined (pass data) = ["bye"])

    let data = (* (let hai = () in ()); hai *)
      Sequence (
        Let {id="hai"; value=Unit; body=Unit},
        Id "hai"
      )

    let () = printf "2.\n"; assert (undefined (pass data) = ["hai"])

    let data = (* let hai = () in
                  let bye = yyy in
                  xxx / hai *)
      Let {id="hai"; value=Unit; body=
        Let {id="bye"; value=Id "yyy"; body=Divide (Id "xxx", Id "hai")}}

    let () = printf "3.\n"; assert (undefined (pass data) = ["xxx"; "yyy"])


    let data = (* let hai = () in
                  bye; let bye = () in
                  hai *)
      Let {id="hai"; value=Unit; body=
        Sequence (
          Id "bye",
          Let {id="bye"; value=Unit; body=Id "hai"}
        )}


    let () = printf "4.\n"; assert (undefined (pass data) = ["bye"])

    (*

    foo;
    let foo = () in
    let bar = bar in
    foo; bar; baz;
    (let baz = () in ()); baz

    *)
    let tree =
      let (%) left right = Sequence (left, right) in
      Id "foo" %
        Let {id="foo"; value=Unit; body=
          Let {id="bar"; value=Id "bar"; body=
            Id "foo" % Id "bar" % Id "baz" %
              Let {id="baz"; value=Unit; body=Unit} % Id "baz"}}

    let () = printf "5.\n";
      assert (undefined (pass tree) = ["baz"; "baz"; "bar"; "foo"])

    (*

    x;
    let x = () in
    let a = () in
    a;
    let b = a in
    let y = y in
    (let z = () in ());
    a; b; z

    *)

    let tree =
      let (%) left right = Sequence (left, right) in
      Id "x" %
        Let {id="x"; value=Unit; body=
          Let {id="a"; value=Unit; body=
            Id "a" %
              Let {id="b"; value=Id "a"; body=
                Let {id="y"; value=Id "y"; body=
                  Let {id="z"; value=Unit; body=Unit} %
                    Id "a" % Id "b" % Id "z"}}}}

    let () = printf "6.\n";
      assert (undefined (pass tree) = ["z"; "y"; "x"])
  end
end