Refine Stdlib.Option metadata handling

Author: cristianoc

analysis/reanalyze/src/Arnold.ml

@@ -534,7 +534,7 @@ module FindFunctionsCalled = struct
     let super = Tast_mapper.default in
     let expr (self : Tast_mapper.mapper) (e : Typedtree.expression) =
       (match e.exp_desc with
-      | Texp_apply {funct = {exp_desc = Texp_ident (callee, _, _)}} ->
+      | Texp_apply {funct = {exp_desc = Texp_ident (callee, _, _)}; _} ->
         let functionName = Path.name callee in
         callees := !callees |> StringSet.add functionName
       | _ -> ());
@@ -579,7 +579,8 @@ module ExtendFunctionTable = struct
     | Some
         {
           exp_desc =
-            Texp_apply {funct = {exp_desc = Texp_ident (path, {loc}, _)}; args};
+            Texp_apply
+              {funct = {exp_desc = Texp_ident (path, {loc}, _)}; args; _};
         }
       when kindOpt <> None ->
       let checkArg ((argLabel : Asttypes.Noloc.arg_label), _argOpt) =
@@ -620,7 +621,7 @@ module ExtendFunctionTable = struct
                             calls a progress function"
                            functionName printPos id_pos;
                      })))
-      | Texp_apply {funct = {exp_desc = Texp_ident (callee, _, _)}; args}
+      | Texp_apply {funct = {exp_desc = Texp_ident (callee, _, _)}; args; _}
         when callee |> FunctionTable.isInFunctionInTable ~functionTable ->
         let functionName = Path.name callee in
         args
@@ -668,8 +669,8 @@ module CheckExpressionWellFormed = struct
       | Texp_ident (path, {loc}, _) ->
         checkIdent ~path ~loc;
         e
-      | Texp_apply {funct = {exp_desc = Texp_ident (functionPath, _, _)}; args}
-        ->
+      | Texp_apply
+          {funct = {exp_desc = Texp_ident (functionPath, _, _)}; args; _} ->
         let functionName = Path.name functionPath in
         args
         |> List.iter (fun ((argLabel : Asttypes.Noloc.arg_label), argOpt) ->
@@ -850,7 +851,7 @@ module Compile = struct
              and create a function call with the appropriate arguments *)
           assert false
         | None -> expr |> expression ~ctx |> evalArgs ~args ~ctx)
-    | Texp_apply {funct = expr; args} ->
+    | Texp_apply {funct = expr; args; _} ->
       expr |> expression ~ctx |> evalArgs ~args ~ctx
     | Texp_let
         ( Recursive,

analysis/reanalyze/src/Exception.ml

@@ -304,7 +304,8 @@ let traverseAst () =
       let exceptions = [arg] |> raiseArgs in
       currentEvents := {Event.exceptions; loc; kind = Raises} :: !currentEvents;
       arg |> snd |> iterExprOpt self
-    | Texp_apply {funct = {exp_desc = Texp_ident (callee, _, _)} as e; args} ->
+    | Texp_apply {funct = {exp_desc = Texp_ident (callee, _, _)} as e; args; _}
+      ->
       let calleeName = Path.name callee in
       if calleeName |> isRaise then
         let exceptions = args |> raiseArgs in

analysis/reanalyze/src/SideEffects.ml

@@ -26,7 +26,7 @@ let rec exprNoSideEffects (expr : Typedtree.expression) =
   | Texp_ident _ | Texp_constant _ -> true
   | Texp_construct (_, _, el) -> el |> List.for_all exprNoSideEffects
   | Texp_function _ -> true
-  | Texp_apply {funct = {exp_desc = Texp_ident (path, _, _)}; args}
+  | Texp_apply {funct = {exp_desc = Texp_ident (path, _, _)}; args; _}
     when path |> pathIsWhitelistedForSideEffects ->
     args |> List.for_all (fun (_, eo) -> eo |> exprOptNoSideEffects)
   | Texp_apply _ -> false

compiler/ml/printtyped.ml

@@ -260,16 +260,6 @@ 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; payload_not_nested; _} ->
-    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 payload_not_nested=%b\n" kind
-      payload_not_nested;
-    attributes i ppf attrs
 
 and expression i ppf x =
   line i ppf "expression %a\n" fmt_location x.exp_loc;
@@ -298,11 +288,12 @@ and expression i ppf x =
     line i ppf "%a" Ident.print param;
     arg_label i ppf p;
     case i ppf case_
-  | Texp_apply {funct = e; args = l; partial} ->
+  | Texp_apply {funct = e; args = l; partial; stdlib_option_call; _} ->
     if partial then line i ppf "partial\n";
     line i ppf "Texp_apply\n";
     expression i ppf e;
-    list i label_x_expression ppf l
+    list i label_x_expression ppf l;
+    if stdlib_option_call then line i ppf "stdlib_option_call\n" else ()
   | Texp_match (e, l1, l2, _partial) ->
     line i ppf "Texp_match\n";
     expression i ppf e;

compiler/ml/rec_check.ml

@@ -199,7 +199,7 @@ let rec classify_expression : Typedtree.expression -> sd =
   | Texp_while _ | Texp_pack _ | Texp_function _ | Texp_extension_constructor _
     ->
     Static
-  | Texp_apply {funct = {exp_desc = Texp_ident (_, _, vd)}} when is_ref vd ->
+  | Texp_apply {funct = {exp_desc = Texp_ident (_, _, vd)}; _} when is_ref vd ->
     Static
   | Texp_apply _ | Texp_match _ | Texp_ifthenelse _ | Texp_send _ | Texp_field _
   | Texp_assert _ | Texp_try _ ->
@@ -233,10 +233,10 @@ let rec expression : Env.env -> Typedtree.expression -> Use.t =
         (discard (expression env e3)))
   | Texp_constant _ -> Use.empty
   | Texp_apply
-      {funct = {exp_desc = Texp_ident (_, _, vd)}; args = [(_, Some arg)]}
+      {funct = {exp_desc = Texp_ident (_, _, vd)}; args = [(_, Some arg)]; _}
     when is_ref vd ->
     Use.guard (expression env arg)
-  | Texp_apply {funct = e; args} ->
+  | Texp_apply {funct = e; args; _} ->
     let arg env (_, eo) = option expression env eo in
     Use.(join (inspect (expression env e)) (inspect (list arg env args)))
   | Texp_tuple exprs -> Use.guard (list expression env exprs)

compiler/ml/tast_iterator.ml

@@ -137,16 +137,11 @@ 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 rec extra = function
+  let 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;
@@ -157,7 +152,7 @@ let expr sub {exp_extra; exp_desc; exp_env; _} =
     sub.value_bindings sub (rec_flag, list);
     sub.expr sub exp
   | Texp_function {case; _} -> sub.case sub case
-  | Texp_apply {funct = exp; args = list} ->
+  | Texp_apply {funct = exp; args = list; _} ->
     sub.expr sub exp;
     List.iter (fun (_, o) -> Option.iter (sub.expr sub) o) list
   | Texp_match (exp, list1, list2, _) ->

compiler/ml/tast_mapper.ml

@@ -180,27 +180,12 @@ 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; payload_not_nested} =
-    {
-      callback = map_stdlib_option_callback sub callback;
-      call_kind;
-      payload_not_nested;
-    }
-  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
@@ -213,11 +198,14 @@ let expr sub x =
     | Texp_function {arg_label; arity; param; case; partial; async} ->
       Texp_function
         {arg_label; arity; param; case = sub.case sub case; partial; async}
-    | Texp_apply {funct = exp; args = list; partial; transformed_jsx} ->
+    | Texp_apply
+        {funct = exp; args = list; partial; transformed_jsx; stdlib_option_call}
+      ->
       Texp_apply
         {
           funct = sub.expr sub exp;
           args = List.map (tuple2 id (opt (sub.expr sub))) list;
+          stdlib_option_call;
           partial;
           transformed_jsx;
         }

compiler/ml/translcore.ml

@@ -646,13 +646,111 @@ 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
+(* Recompute metadata needed for inlining Stdlib.Option helpers at translation
+   time; the typed tree only marks such applications with a boolean flag. *)
+type stdlib_option_call_kind =
+  | Stdlib_option_forEach
+  | Stdlib_option_map of {result_cannot_contain_undefined: bool}
+  | Stdlib_option_flatMap
+
+type stdlib_option_callback =
+  | Stdlib_option_inline_lambda of {param: Ident.t; body: expression}
+  | Stdlib_option_inline_ident of expression
+
+type stdlib_option_call = {
+  callback: stdlib_option_callback;
+  call_kind: stdlib_option_call_kind;
+  payload_not_nested: bool;
+}
+
+type stdlib_option_fun_kind =
+  | Stdlib_option_fun_forEach
+  | Stdlib_option_fun_map
+  | Stdlib_option_fun_flatMap
+
+let stdlib_option_fun_of_path env path =
+  let canonical = Env.normalize_path_prefix None env path in
+  match Path.flatten canonical with
+  | `Contains_apply -> None
+  | `Ok (head, rest) ->
+    let segments = Ident.name head :: rest in
+    let matches fname =
+      match segments with
+      | ["Stdlib"; "Option"; name] -> String.equal name fname
+      | ["Stdlib_Option"; name] -> String.equal name fname
+      | _ -> false
+    in
+    if matches "forEach" then Some Stdlib_option_fun_forEach
+    else if matches "map" then Some Stdlib_option_fun_map
+    else if matches "flatMap" then Some Stdlib_option_fun_flatMap
+    else None
+
+let inline_lambda_callback (expr : expression) : stdlib_option_callback option =
+  match expr.exp_desc with
+  | Texp_function {arg_label = Nolabel; case; partial = Total; async = false; _}
+    when Option.is_none case.c_guard -> (
+    match case.c_lhs.pat_desc with
+    | Tpat_var (param, _) ->
+      Some (Stdlib_option_inline_lambda {param; body = case.c_rhs})
+    | _ -> None)
+  | _ -> None
+
+let inline_ident_callback (expr : expression) : stdlib_option_callback option =
+  match expr.exp_desc with
+  | Texp_ident _ -> Some (Stdlib_option_inline_ident expr)
+  | _ -> None
+
+let callback_return_type env (expr : expression) =
+  match (Ctype.expand_head env expr.exp_type).desc with
+  | Tarrow (_, ret_ty, _, _) -> Some ret_ty
+  | _ -> None
+
+let detect_stdlib_option_call env (funct : expression)
+    (args : (Noloc.arg_label * expression option) list) :
+    stdlib_option_call option =
+  match funct.exp_desc with
+  | Texp_ident (path, _, _) -> (
+    match stdlib_option_fun_of_path env path with
+    | None -> None
+    | Some fun_kind -> (
+      match args with
+      | [(Nolabel, Some opt_expr); (Nolabel, Some callback_expr)] -> (
+        let callback_info =
+          match inline_lambda_callback callback_expr with
+          | Some info -> Some info
+          | None -> inline_ident_callback callback_expr
+        in
+        match callback_info with
+        | None -> None
+        | Some callback ->
+          let payload_not_nested =
+            match (Ctype.expand_head env opt_expr.exp_type).desc with
+            | Tconstr (path, [payload_ty], _)
+              when Path.same path Predef.path_option ->
+              Typeopt.type_cannot_contain_undefined payload_ty env
+            | _ -> false
+          in
+          let call_kind =
+            match fun_kind with
+            | Stdlib_option_fun_forEach -> Stdlib_option_forEach
+            | Stdlib_option_fun_map ->
+              let result_cannot_contain_undefined =
+                match callback with
+                | Stdlib_option_inline_lambda {body; _} ->
+                  Typeopt.type_cannot_contain_undefined body.exp_type
+                    body.exp_env
+                | Stdlib_option_inline_ident cb -> (
+                  match callback_return_type cb.exp_env cb with
+                  | Some ret_ty ->
+                    Typeopt.type_cannot_contain_undefined ret_ty cb.exp_env
+                  | None -> false)
+              in
+              Stdlib_option_map {result_cannot_contain_undefined}
+            | Stdlib_option_fun_flatMap -> Stdlib_option_flatMap
+          in
+          Some {callback; call_kind; payload_not_nested})
+      | _ -> None))
+  | _ -> None
 
 let lambda_none = Lconst (Const_pointer (0, Pt_shape_none))
 
@@ -765,11 +863,15 @@ 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; stdlib_option_call} -> (
     let inlined, funct =
       Translattribute.get_and_remove_inlined_attribute funct
     in
-    let option_call_info = stdlib_option_call_extra e in
+    let option_call_info =
+      if stdlib_option_call then detect_stdlib_option_call e.exp_env funct oargs
+      else None
+    in
     let uncurried_partial_application =
       (* In case of partial application foo(args, ...) when some args are missing,
          get the arity *)