(* Opening GT *)
open GT

(* Some useful idioms: identity, fixpoint *)
let id x = x
let rec fix f x = f (fix f) x;;

(* The first example: decorated expressions *)
@type 'expr a_expr =
| Const of int
| Add   of 'expr * 'expr with show, gmap, foldl

(* Must replace with @type expr = expr a_expr with ... *)
type expr = expr a_expr
let show_expr = fix (fun s e -> show(a_expr) s e)

(* Must replace with @type 'a decorated = ('a decorated a_expr * 'a) with ...*)
type 'a decorated = ('a decorated a_expr * 'a)
let show_decorated sa = fix (fun s (e, a) -> Printf.sprintf "(%s, %s)"(show(a_expr) s e) (sa a))

(* Simple --- one-level --- decoration *)
module Simple =
  struct

    (* Decoration *)
    let decorate : (unit -> 'a) -> expr -> 'a decorated = fun fd e ->
      fix (fun s e -> let d = fd () in gmap(a_expr) s e, d) e

    (* Stripping *)
    let strip : 'a decorated -> expr = fun e ->
      fix (fun s e -> gmap(a_expr) s (fst e)) e

    (* Testing *)
    let _ =
      let e = Add (Add (Const 1, Const 2), Const 3) in
      Printf.printf "%s\n" (show_expr e);
      let e = decorate (let i = ref 0 in fun () -> let n = !i in incr i; n) e in
      Printf.printf "%s\n" (show_decorated string_of_int e);
      Printf.printf "%s\n" (show_expr @@ strip e)

  end

(* Advanced --- multilayer decorations *)
module Advanced =
  struct

    (* Lifting --- attaching unit decoration *)
    let lift = Simple.decorate (fun _ -> ())

    (* Redecorate *)
    let redecorate : ('a -> 'b) -> 'a decorated -> 'b decorated = fun ab e ->
      fix (fun s (e, a) -> gmap(a_expr) s e, ab a) e

    (* Testing *)
    let _ =
      let e = Add (Add (Const 1, Const 2), Const 3) in
      Printf.printf "%s\n" (show_expr e);
      let e = lift e in
      Printf.printf "%s\n" (show_decorated (fun _ -> "()") e);
      let e = redecorate (let i = ref 0 in fun () -> let n = !i in incr i; n) e in
      Printf.printf "%s\n" (show_decorated (show int) e);
      let e = redecorate (let i = ref 0 in fun a -> let n = !i in incr i; (a, n)) e in
      Printf.printf "%s\n" (show_decorated (show(pair) (show int) (show int)) e);
      let e = redecorate snd e in
      Printf.printf "%s\n" (show_decorated (show int) e)

  end

(* Custom --- no use of gmap *)
module Custom =
  struct

    (* OO transformer for redecorate *)
    class virtual ['a, 'b] decorator =
      object ('self)
        inherit [] @pair
      end

  end