Source file `ocaml_monadic_ppx.ml`

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
open Ppxlib

let bind_expand ~ctxt:_ expr =
  let loc = expr.pexp_loc in
  match expr.pexp_desc with
  | Pexp_let (Nonrecursive, value_bindings, body) ->
      (* This is a let%bind expression!  It's of the form
                   let%bind $p1 = $e1 and ... and $pn = $en in $e0
                 and we want it to take the form
                   bind $e1 (fun $p1 -> ... bind $en (fun $pn -> ...) ...)
      *)
      let rec bind_wrap value_bindings' =
        match value_bindings' with
        | {
            pvb_pat = bind_pattern;
            pvb_expr = bind_expr;
            pvb_attributes = [];
            pvb_loc = _bind_loc;
          }
          :: value_bindings'' ->
            (* Recurse and then wrap the resulting body. *)
            let body' = bind_wrap value_bindings'' in
            let cont_function = [%expr fun [%p bind_pattern] -> [%e body']] in
            [%expr bind [%e bind_expr] [%e cont_function]]
        | _ ->
            (* Nothing left to do.  Just return the body. *)
            body
      in
      bind_wrap value_bindings
  | Pexp_match (expr_match, cases) ->
      let f = Ast_helper.Exp.function_ cases in
      [%expr bind [%e expr_match] [%e f]]
  | Pexp_ifthenelse (expr_if, expr_then, expr_else) ->
      let expr_else =
        match expr_else with None -> [%expr ()] | Some case -> case
      in
      let cases =
        [
          Ast_helper.Exp.case [%pat? true] expr_then;
          Ast_helper.Exp.case [%pat? false] expr_else;
        ]
      in
      let f = Ast_helper.Exp.function_ cases in
      [%expr bind [%e expr_if] [%e f]]
  | Pexp_sequence (expr_seq_l, expr_seq_r) ->
      [%expr bind [%e expr_seq_l] (fun () -> [%e expr_seq_r])]
  | _ -> expr

let orzero_expand ~ctxt:_ expr =
  let loc = expr.pexp_loc in
  match expr.pexp_desc with
  | Pexp_let (Nonrecursive, value_bindings, body) ->
      (* This is a let%orzero expression.  It's of the form
           let%orzero $p1 = $e1 and ... and $pn = $en in $e0
         and we want it to take the form
           match $e1 with
           | $p1 -> (match $e2 with
                     | $p2 -> ...
                              (match $en with
                               | $pn -> $e0
                               | _ -> zero ())
                     | _ -> zero ())
           | _ -> zero ()
      *)
      let rec orzero_wrap value_bindings' =
        match value_bindings' with
        | {
            pvb_pat = orzero_pattern;
            pvb_expr = orzero_expr;
            pvb_attributes = [];
            pvb_loc = _orzero_loc;
          }
          :: value_bindings'' ->
            (* Recurse and then wrap the resulting body. *)
            let body' = orzero_wrap value_bindings'' in
            [%expr
              match [%e orzero_expr] with
              | [%p orzero_pattern] -> [%e body']
              | _ -> zero ()]
        | _ ->
            (* Nothing left to do.  Just return the body. *)
            body
      in
      orzero_wrap value_bindings
  | _ -> expr

let bind_extension =
  Extension.V3.declare "bind" Extension.Context.expression
    Ast_pattern.(single_expr_payload __)
    bind_expand

let orzero_extension =
  Extension.V3.declare "orzero" Extension.Context.expression
    Ast_pattern.(single_expr_payload __)
    orzero_expand

let orzero_rule = Ppxlib.Context_free.Rule.extension orzero_extension

let bind_rule = Ppxlib.Context_free.Rule.extension bind_extension

let expr_mapper =
  object
    inherit Ast_traverse.map as super

    method! expression e =
      let e = super#expression e in
      match e with
      | [%expr
          [%guard [%e? guard_expr]];
          [%e? body_expr]] ->
          (* This is a sequenced expression with a [%guard ...] extension.  It
             takes the form
               [%guard expr']; expr
             and we want it to take the form
               if expr' then expr else zero ()
          *)
          let loc = e.pexp_loc in
          [%expr if [%e guard_expr] then [%e body_expr] else zero ()]
      | _ -> e
  end

let () =
  Driver.register_transformation "ocaml-monadic"
    ~rules:[ bind_rule; orzero_rule ] ~impl:expr_mapper#structure