Implement typed Option stdlib optimizations

Author: cristianoc

compiler/frontend/ast_option_optimizations.ml

@@ -112,8 +112,7 @@ let expr_mapper ~async_context ~in_function_def (self : mapper)
             body;
         pexp_attributes;
       })
-  | Pexp_apply _ ->
-    Ast_exp_apply.app_exp_mapper e self |> Ast_option_optimizations.transform
+  | Pexp_apply _ -> Ast_exp_apply.app_exp_mapper e self
   | Pexp_match
       ( b,
         [

compiler/frontend/ast_option_optimizations.mli

@@ -260,6 +260,15 @@ and expression_extra i ppf x attrs =
   | Texp_newtype s ->
     line i ppf "Texp_newtype \"%s\"\n" s;
     attributes i ppf attrs
+  | Texp_stdlib_option_call {call_kind; _} ->
+    let kind =
+      match call_kind with
+      | Stdlib_option_forEach -> "forEach"
+      | Stdlib_option_map {result_cannot_contain_undefined = _} -> "map"
+      | Stdlib_option_flatMap -> "flatMap"
+    in
+    line i ppf "Texp_stdlib_option_call %s\n" kind;
+    attributes i ppf attrs
 
 and expression i ppf x =
   line i ppf "expression %a\n" fmt_location x.exp_loc;

compiler/frontend/bs_builtin_ppx.ml

@@ -137,11 +137,16 @@ let pat sub {pat_extra; pat_desc; pat_env; _} =
   | Tpat_alias (p, _, _) -> sub.pat sub p
 
 let expr sub {exp_extra; exp_desc; exp_env; _} =
-  let extra = function
+  let rec extra = function
     | Texp_constraint cty -> sub.typ sub cty
     | Texp_coerce cty2 -> sub.typ sub cty2
     | Texp_newtype _ -> ()
     | Texp_open (_, _, _, _) -> ()
+    | Texp_stdlib_option_call info -> stdlib_option_call info
+  and stdlib_option_call {callback; _} = stdlib_option_callback callback
+  and stdlib_option_callback = function
+    | Stdlib_option_inline_lambda {body; _} -> sub.expr sub body
+    | Stdlib_option_inline_ident expr -> sub.expr sub expr
   in
   List.iter (fun (e, _, _) -> extra e) exp_extra;
   sub.env sub exp_env;

compiler/ml/printtyped.ml

@@ -180,12 +180,23 @@ let pat sub x =
   {x with pat_extra; pat_desc; pat_env}
 
 let expr sub x =
+  let map_stdlib_option_callback sub = function
+    | Stdlib_option_inline_lambda {param; body} ->
+      Stdlib_option_inline_lambda {param; body = sub.expr sub body}
+    | Stdlib_option_inline_ident expr ->
+      Stdlib_option_inline_ident (sub.expr sub expr)
+  in
+  let map_stdlib_option_call sub {callback; call_kind} =
+    {callback = map_stdlib_option_callback sub callback; call_kind}
+  in
   let extra = function
     | Texp_constraint cty -> Texp_constraint (sub.typ sub cty)
     | Texp_coerce cty2 -> Texp_coerce (sub.typ sub cty2)
     | Texp_open (ovf, path, loc, env) ->
       Texp_open (ovf, path, loc, sub.env sub env)
     | Texp_newtype _ as d -> d
+    | Texp_stdlib_option_call info ->
+      Texp_stdlib_option_call (map_stdlib_option_call sub info)
   in
   let exp_extra = List.map (tuple3 extra id id) x.exp_extra in
   let exp_env = sub.env sub x.exp_env in

compiler/ml/tast_iterator.ml

@@ -646,6 +646,16 @@ let rec cut n l =
 
 let try_ids = Hashtbl.create 8
 
+let stdlib_option_call_extra exp =
+  let rec aux = function
+    | [] -> None
+    | (Texp_stdlib_option_call info, _, _) :: _ -> Some info
+    | _ :: rest -> aux rest
+  in
+  aux exp.exp_extra
+
+let lambda_none = Lconst (Const_pointer (0, Pt_shape_none))
+
 let extract_directive_for_fn exp =
   exp.exp_attributes
   |> List.find_map (fun ({txt}, payload) ->
@@ -755,10 +765,11 @@ and transl_exp0 (e : Typedtree.expression) : Lambda.lambda =
           (Lprim
              (Pccall (set_transformed_jsx d ~transformed_jsx), argl, e.exp_loc))
       | _ -> wrap (Lprim (prim, argl, e.exp_loc))))
-  | Texp_apply {funct; args = oargs; partial; transformed_jsx} ->
+  | Texp_apply {funct; args = oargs; partial; transformed_jsx} -> (
     let inlined, funct =
       Translattribute.get_and_remove_inlined_attribute funct
     in
+    let option_call_info = stdlib_option_call_extra e in
     let uncurried_partial_application =
       (* In case of partial application foo(args, ...) when some args are missing,
          get the arity *)
@@ -771,8 +782,17 @@ and transl_exp0 (e : Typedtree.expression) : Lambda.lambda =
         | None -> None
       else None
     in
-    transl_apply ~inlined ~uncurried_partial_application ~transformed_jsx
-      (transl_exp funct) oargs e.exp_loc
+    match option_call_info with
+    | Some info when not partial -> (
+      match oargs with
+      | (Nolabel, Some opt_expr) :: _ ->
+        transl_stdlib_option_call e opt_expr info oargs
+      | _ ->
+        transl_apply ~inlined ~uncurried_partial_application ~transformed_jsx
+          (transl_exp funct) oargs e.exp_loc)
+    | _ ->
+      transl_apply ~inlined ~uncurried_partial_application ~transformed_jsx
+        (transl_exp funct) oargs e.exp_loc)
   | Texp_match (arg, pat_expr_list, exn_pat_expr_list, partial) ->
     transl_match e arg pat_expr_list exn_pat_expr_list partial
   | Texp_try (body, pat_expr_list) ->
@@ -924,6 +944,53 @@ and transl_exp0 (e : Typedtree.expression) : Lambda.lambda =
     if !Clflags.noassert then lambda_unit
     else Lifthenelse (transl_exp cond, lambda_unit, assert_failed e)
 
+and bind_option_value opt_var opt_loc callback =
+  let value_expr = Lprim (Pval_from_option, [opt_var], opt_loc) in
+  match callback with
+  | Stdlib_option_inline_lambda {param; body} ->
+    bind Strict param value_expr (transl_exp body)
+  | Stdlib_option_inline_ident expr ->
+    let func = transl_exp expr in
+    let value_id = Ident.create "__res_option_value" in
+    let apply =
+      Lapply
+        {
+          ap_func = func;
+          ap_args = [Lvar value_id];
+          ap_inlined = Default_inline;
+          ap_loc = expr.exp_loc;
+          ap_transformed_jsx = false;
+        }
+    in
+    bind Strict value_id value_expr apply
+
+and transl_stdlib_option_call exp opt_expr info oargs =
+  match oargs with
+  | (Nolabel, Some _) :: (Nolabel, Some _) :: _ | (Nolabel, Some _) :: [] ->
+    let opt_lam = transl_exp opt_expr in
+    let opt_id = Ident.create "__res_option_opt" in
+    let opt_var = Lvar opt_id in
+    let callback_result = bind_option_value opt_var exp.exp_loc info.callback in
+    let some_branch =
+      match info.call_kind with
+      | Stdlib_option_forEach -> callback_result
+      | Stdlib_option_map {result_cannot_contain_undefined} ->
+        let tag =
+          if result_cannot_contain_undefined then Blk_some_not_nested
+          else Blk_some
+        in
+        Lprim (Pmakeblock tag, [callback_result], exp.exp_loc)
+      | Stdlib_option_flatMap -> callback_result
+    in
+    let none_branch =
+      match info.call_kind with
+      | Stdlib_option_forEach -> lambda_unit
+      | Stdlib_option_map _ | Stdlib_option_flatMap -> lambda_none
+    in
+    let cond = Lprim (Pis_not_none, [opt_var], exp.exp_loc) in
+    bind Strict opt_id opt_lam (Lifthenelse (cond, some_branch, none_branch))
+  | _ -> assert false
+
 and transl_list expr_list = List.map transl_exp expr_list
 
 and transl_guard guard rhs =