flambda-backend: Statically enabled probes (#1388)

* Support [~enabled_at_init] argument of %probe

The following new constructs are accepted:
[%probe "name" ~enabled_at_init:true handler]
[%probe "name" ~enabled_at_init:false handler]
The default is
[%probe "name" handler]
and implies [~enabled_at_init:false]

* Propagate [enabled_at_init] to Emit

* emit call instead of jmp opcode for enabled_at_init probes

* Add test

* Initialize semaphore according to enabled_at_init

Check that all probes with the same name in the current compilation
unit have consistent [enabled_at_init].

* Improve error message

* Add a test for inconsistent enabled_at_init for %probe

* Refactor parsing

* Add a module for Probe in Flambda2

* Rename variables [probe_name] to [probe] in middle_end/flambda2 code

* Improve code style and make [desc] private

lambda/lambda.ml

@@ -411,7 +411,7 @@ let equal_inlined_attribute (x : inlined_attribute) (y : inlined_attribute) =
     | Hint_inlined | Unroll _ | Default_inlined), _ ->
     false
 
-type probe_desc = { name: string }
+type probe_desc = { name: string; enabled_at_init: bool; }
 type probe = probe_desc option
 
 type specialise_attribute =

lambda/lambda.mli

@@ -298,7 +298,7 @@ type inlined_attribute =
 val equal_inline_attribute : inline_attribute -> inline_attribute -> bool
 val equal_inlined_attribute : inlined_attribute -> inlined_attribute -> bool
 
-type probe_desc = { name: string }
+type probe_desc = { name: string; enabled_at_init: bool; }
 type probe = probe_desc option
 
 type specialise_attribute =

lambda/translcore.ml

@@ -863,7 +863,7 @@ and transl_exp0 ~in_new_scope ~scopes e =
           Llet(pure, Lambda.layout_module, oid,
                !transl_module ~scopes Tcoerce_none None od.open_expr, body)
       end
-  | Texp_probe {name; handler=exp} ->
+  | Texp_probe {name; handler=exp; enabled_at_init} ->
     if !Clflags.native_code && !Clflags.probes then begin
       let lam = transl_exp ~scopes exp in
       let map =
@@ -946,7 +946,7 @@ and transl_exp0 ~in_new_scope ~scopes e =
           ap_tailcall = Default_tailcall;
           ap_inlined = Never_inlined;
           ap_specialised = Always_specialise;
-          ap_probe = Some {name};
+          ap_probe = Some {name; enabled_at_init};
         }
       in
       begin match Config.flambda || Config.flambda2 with

typing/tast_mapper.ml

@@ -439,8 +439,8 @@ let expr sub x =
         e
     | Texp_open (od, e) ->
         Texp_open (sub.open_declaration sub od, sub.expr sub e)
-    | Texp_probe {name; handler} ->
-      Texp_probe {name; handler = sub.expr sub handler }
+    | Texp_probe {name; handler; enabled_at_init;} ->
+      Texp_probe {name; handler = sub.expr sub handler; enabled_at_init}
     | Texp_probe_is_enabled _ as e -> e
     | Texp_exclave exp ->
         Texp_exclave (sub.expr sub exp)

typing/typecore.ml

@@ -5498,7 +5498,8 @@ and type_expect_
           raise (Error (loc, env, Invalid_extension_constructor_payload))
       end
   | Pexp_extension ({ txt = ("probe" | "ocaml.probe"); _ }, payload) ->
-      begin match payload with
+    let name, name_loc, args =
+      match payload with
       | PStr
           ([{ pstr_desc =
                 Pstr_eval
@@ -5508,22 +5509,37 @@ and type_expect_
                                (Pexp_constant (Pconst_string(name,_,None)));
                              pexp_loc = name_loc;
                              _ }
-                          , [Nolabel, arg]))
+                          , args))
                    ; _ }
-                  , _)}]) ->
-        check_probe_name name name_loc env;
-        let env = Env.add_lock Alloc_mode.global env in
-        Env.add_probe name;
-        let exp = type_expect env mode_global arg
-                    (mk_expected Predef.type_unit) in
-        rue {
-          exp_desc = Texp_probe {name; handler=exp};
-          exp_loc = loc; exp_extra = [];
-          exp_type = instance Predef.type_unit;
-          exp_attributes = sexp.pexp_attributes;
-          exp_env = env }
+                  , _)}]) -> name, name_loc, args
       | _ -> raise (Error (loc, env, Probe_format))
-    end
+    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_lock Alloc_mode.global env in
+    Env.add_probe name;
+    let exp = type_expect env mode_global 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 }
   | Pexp_extension ({ txt = ("probe_is_enabled"
                             |"ocaml.probe_is_enabled"); _ }, payload) ->
       begin match payload with
@@ -8041,8 +8057,9 @@ let report_error ~loc env = function
         name name
   | Probe_format ->
       Location.errorf ~loc
-        "Probe points must consist of a name, as a string \
-         literal, followed by a single expression of type unit."
+        "Probe points must consist of a name, as a string literal, \
+         optionally followed by ~enabled_at_init:true or ~enabled_at_init:false, \
+         followed by a single expression of type unit."
   | Probe_is_enabled_format ->
       Location.errorf ~loc
         "%%probe_is_enabled points must specify a single probe name as a \

typing/typedtree.ml

@@ -174,7 +174,7 @@ and expression_desc =
   | Texp_unreachable
   | Texp_extension_constructor of Longident.t loc * Path.t
   | Texp_open of open_declaration * expression
-  | Texp_probe of { name:string; handler:expression; }
+  | Texp_probe of { name:string; handler:expression; enabled_at_init:bool; }
   | Texp_probe_is_enabled of { name:string }
   | Texp_exclave of expression
 

typing/typedtree.mli

@@ -351,7 +351,7 @@ and expression_desc =
   | Texp_extension_constructor of Longident.t loc * Path.t
   | Texp_open of open_declaration * expression
         (** let open[!] M in e *)
-  | Texp_probe of { name:string; handler:expression; }
+  | Texp_probe of { name:string; handler:expression; enabled_at_init:bool }
   | Texp_probe_is_enabled of { name:string }
   | Texp_exclave of expression