flambda-backend: Fix interaction between probes and ocamldep (#2616)

Author: ncik-roberts

parsing/builtin_attributes.ml

@@ -937,3 +937,47 @@ let assume_zero_alloc ~is_check_allowed check : Zero_alloc_utils.Assume_info.t =
       Location.prerr_warning loc (Warnings.Attribute_payload (name, msg))
     end;
     Zero_alloc_utils.Assume_info.none
+
+type tracing_probe =
+  { name : string;
+    name_loc : Location.t;
+    enabled_at_init : bool;
+    arg : Parsetree.expression;
+  }
+
+let get_tracing_probe_payload (payload : Parsetree.payload) =
+  let ( let* ) = Result.bind in
+  let* name, name_loc, args =
+    match payload with
+    | PStr
+        ([{ pstr_desc =
+              Pstr_eval
+                ({ pexp_desc =
+                      (Pexp_apply
+                        ({ pexp_desc=
+                              (Pexp_constant (Pconst_string(name,_,None)));
+                            pexp_loc = name_loc;
+                            _ }
+                        , args))
+                  ; _ }
+                , _)}]) -> Ok (name, name_loc, args)
+    | _ -> Error ()
+  in
+  let bool_of_string = function
+    | "true" -> Ok true
+    | "false" -> Ok false
+    | _ -> Error ()
+  in
+  let* arg, enabled_at_init =
+    match args with
+    | [Nolabel, arg] -> Ok (arg, false)
+    | [Labelled "enabled_at_init",
+        { pexp_desc =
+            Pexp_construct({ txt = Longident.Lident b; _ },
+                          None); _ };
+        Nolabel, arg] ->
+          let* enabled_at_init = bool_of_string b in
+          Ok (arg, enabled_at_init)
+    | _ -> Error ()
+  in
+  Ok { name; name_loc; enabled_at_init; arg }

parsing/builtin_attributes.mli

@@ -272,3 +272,23 @@ val get_zero_alloc_attribute :
 
 val assume_zero_alloc :
   is_check_allowed:bool -> zero_alloc_attribute -> Zero_alloc_utils.Assume_info.t
+
+type tracing_probe =
+  { name : string;
+    name_loc : Location.t;
+    enabled_at_init : bool;
+    arg : Parsetree.expression;
+  }
+
+(* Gets the payload of a [probe] extension node. Example syntax of a probe
+   that's disabled by default:
+
+   [%probe "my_probe" arg]
+
+   You can use [enabled_at_init] to control whether the probe is enabled
+   by default:
+
+   [%probe "my_probe" ~enabled_at_init:true arg]
+*)
+val get_tracing_probe_payload :
+  Parsetree.payload -> (tracing_probe, unit) result

parsing/depend.ml

@@ -322,6 +322,11 @@ let rec add_expr bv exp =
       | Pstr_eval ({ pexp_desc = Pexp_construct (c, None) }, _) -> add bv c
       | _ -> handle_extension e
       end
+  | Pexp_extension (({ txt = ("probe"|"ocaml.probe"); _ }, payload) as e) ->
+      begin match Builtin_attributes.get_tracing_probe_payload payload with
+      | Error () -> handle_extension e
+      | Ok { arg; _ } -> add_expr bv arg
+      end
   | Pexp_extension e -> handle_extension e
   | Pexp_unreachable -> ()
 

typing/typecore.ml

@@ -6358,49 +6358,22 @@ and type_expect_
           raise (Error (loc, env, Invalid_extension_constructor_payload))
       end
   | Pexp_extension ({ txt = ("probe" | "ocaml.probe"); _ }, payload) ->
-    let name, name_loc, args =
-      match payload with
-      | PStr
-          ([{ pstr_desc =
-                Pstr_eval
-                  ({ pexp_desc =
-                       (Pexp_apply
-                          ({ pexp_desc=
-                               (Pexp_constant (Pconst_string(name,_,None)));
-                             pexp_loc = name_loc;
-                             _ }
-                          , args))
-                   ; _ }
-                  , _)}]) -> name, name_loc, args
-      | _ -> raise (Error (loc, env, Probe_format))
-    in
-    let bool_of_string = function
-      | "true" -> true
-      | "false" -> false
-      | _ -> raise (Error (loc, env, Probe_format))
-    in
-    let arg, enabled_at_init =
-      match args with
-      | [Nolabel, arg] -> arg, false
-      | [Labelled "enabled_at_init",
-         { pexp_desc =
-             Pexp_construct({ txt = Longident.Lident b; _ },
-                            None); _ };
-         Nolabel, arg] -> arg, bool_of_string b
-      | _ -> raise (Error (loc, env, Probe_format))
-    in
-    check_probe_name name name_loc env;
-    let env = Env.add_escape_lock Probe env in
-    let env = Env.add_share_lock Probe env in
-    Env.add_probe name;
-    let exp = type_expect env mode_legacy arg
-                (mk_expected Predef.type_unit) in
-    rue {
-      exp_desc = Texp_probe {name; handler=exp; enabled_at_init};
-      exp_loc = loc; exp_extra = [];
-      exp_type = instance Predef.type_unit;
-      exp_attributes = sexp.pexp_attributes;
-      exp_env = env }
+    begin match Builtin_attributes.get_tracing_probe_payload payload with
+    | Error () -> raise (Error (loc, env, Probe_format))
+    | Ok { name; name_loc; enabled_at_init; arg; } ->
+        check_probe_name name name_loc env;
+        let env = Env.add_escape_lock Probe env in
+        let env = Env.add_share_lock Probe env in
+        Env.add_probe name;
+        let exp = type_expect env mode_legacy arg
+                    (mk_expected Predef.type_unit) in
+        rue {
+          exp_desc = Texp_probe {name; handler=exp; enabled_at_init};
+          exp_loc = loc; exp_extra = [];
+          exp_type = instance Predef.type_unit;
+          exp_attributes = sexp.pexp_attributes;
+          exp_env = env }
+    end
   | Pexp_extension ({ txt = ("probe_is_enabled"
                             |"ocaml.probe_is_enabled"); _ }, payload) ->
       begin match payload with