Relax modes for probe (#2968)

Author: riaqn

ocaml/lambda/translcore.ml

@@ -1116,24 +1116,28 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
         let assume_zero_alloc = get_assume_zero_alloc ~scopes in
         let scopes = enter_value_definition ~scopes ~assume_zero_alloc funcid in
         lfunction
-          ~kind:(Curried {nlocal=0})
+          (* We conservatively assume that all arguments are local. This doesn't
+             hurt performance as probe handlers are always applied fully. *)
+          ~kind:(Curried {nlocal=List.length param_idents})
           (* CR layouts: Adjust param layouts when we allow other things in
              probes. *)
-          ~params:(List.map (fun name -> { name; layout = layout_probe_arg; attributes = Lambda.default_param_attribute; mode = alloc_heap }) param_idents)
+          ~params:(List.map (fun name -> { name; layout = layout_probe_arg; attributes = Lambda.default_param_attribute; mode = alloc_local }) param_idents)
           ~return:return_layout
-          ~body:(maybe_region_layout return_layout body)
+          ~body:body
           ~loc:(of_location ~scopes exp.exp_loc)
           ~attr
           ~mode:alloc_heap
-          ~ret_mode:alloc_heap
+          ~ret_mode:alloc_local
+          (* CR zqian: the handler function doesn't have a region. However, the
+             [region] field is currently broken. *)
           ~region:true
       in
       let app =
         { ap_func = Lvar funcid;
           ap_args = List.map (fun id -> Lvar id) arg_idents;
           ap_result_layout = return_layout;
           ap_region_close = Rc_normal;
-          ap_mode = alloc_heap;
+          ap_mode = alloc_local;
           ap_loc = of_location e.exp_loc ~scopes;
           ap_tailcall = Default_tailcall;
           ap_inlined = Never_inlined;

ocaml/testsuite/tests/typing-modes/probe.ml

@@ -0,0 +1,12 @@
+(* TEST
+    flags += "-extension unique";
+    expect;
+*)
+
+(* probe can refer to local,nonportable,once values *)
+
+let f (x @ local nonportable once) =
+    [%probe "a" (let _ = x in ())]
+[%%expect{|
+val f : local_ once_ 'a -> unit = <fun>
+|}]

ocaml/typing/typecore.ml

@@ -6296,8 +6296,6 @@ and type_expect_
     | 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