Fix interaction between probes and ocamldep

ocaml/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 }

ocaml/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

ocaml/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 -> ()
 

ocaml/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

tests/backend/probes/dune

@@ -54,3 +54,17 @@
        (= %{architecture} "amd64")))
  (deps t3.expected t3.output)
  (action (diff t3.expected t3.output)))
+
+(rule
+ (target test_ocamldep.output)
+ (deps for_test_ocamldep1.ml for_test_ocamldep2.ml test_ocamldep.ml)
+ (action (with-outputs-to %{target} (run %{bin:ocamlopt.opt} -depend %{deps}))))
+
+(rule
+ (alias   runtest)
+ (enabled_if
+  (and (= %{context_name} "main")
+       (= %{system} "linux")
+       (= %{architecture} "amd64")))
+ (deps test_ocamldep.expected test_ocamldep.output)
+ (action (diff test_ocamldep.expected test_ocamldep.output)))

tests/backend/probes/for_test_ocamldep1.ml

@@ -0,0 +1,3 @@
+(* See [test_ocamldep.ml] for more context. *)
+
+let x = x

tests/backend/probes/for_test_ocamldep2.ml

@@ -0,0 +1,3 @@
+(* See [test_ocamldep.ml] for more context. *)
+
+let x = x

tests/backend/probes/test_ocamldep.expected

@@ -0,0 +1,10 @@
+for_test_ocamldep1.cmo :
+for_test_ocamldep1.cmx :
+for_test_ocamldep2.cmo :
+for_test_ocamldep2.cmx :
+test_ocamldep.cmo : \
+    for_test_ocamldep2.cmo \
+    for_test_ocamldep1.cmo
+test_ocamldep.cmx : \
+    for_test_ocamldep2.cmx \
+    for_test_ocamldep1.cmx

tests/backend/probes/test_ocamldep.ml

@@ -0,0 +1,8 @@
+(* Regression test for a bug where ocamldep wouldn't register a dependency on
+   the payload of a probe expression.
+*)
+
+let () =
+  [%probe "probe1" For_test_ocamldep1.x];
+  [%probe "probe2" ~enabled_at_init:true For_test_ocamldep2.x];
+  ()