Enable Stdlib.Effect + add Flambda 2 support (ocaml-flambda#2205)

backend/cmm_helpers.ml

@@ -4108,3 +4108,45 @@ let allocate_unboxed_nativeint_array =
 let block_header x y = block_header x y
 
 let dls_get ~dbg = Cop (Cdls_get, [], dbg)
+
+let perform ~dbg eff =
+  let cont =
+    make_alloc dbg Runtimetags.cont_tag
+      [int_const dbg 0]
+      ~mode:Lambda.alloc_heap
+  in
+  (* Rc_normal means "allow tailcalls". Preventing them here by using Rc_nontail
+     improves backtraces of paused fibers. *)
+  Cop
+    ( Capply (typ_val, Rc_nontail),
+      [Cconst_symbol (Cmm.global_symbol "caml_perform", dbg); eff; cont],
+      dbg )
+
+let run_stack ~dbg ~stack ~f ~arg =
+  (* Rc_normal would be fine here, but this is unlikely to ever be a tail call
+     (usages of this primitive shouldn't be generated in tail position), so we
+     use Rc_nontail for clarity. *)
+  Cop
+    ( Capply (typ_val, Rc_nontail),
+      [Cconst_symbol (Cmm.global_symbol "caml_runstack", dbg); stack; f; arg],
+      dbg )
+
+let resume ~dbg ~stack ~f ~arg =
+  (* Rc_normal is required here, because there are some uses of effects with
+     repeated resumes, and these should consume O(1) stack space by tail-calling
+     caml_resume. *)
+  Cop
+    ( Capply (typ_val, Rc_normal),
+      [Cconst_symbol (Cmm.global_symbol "caml_resume", dbg); stack; f; arg],
+      dbg )
+
+let reperform ~dbg ~eff ~cont ~last_fiber =
+  (* Rc_normal is required here, this is used in tail position and should tail
+     call. *)
+  Cop
+    ( Capply (typ_val, Rc_normal),
+      [ Cconst_symbol (Cmm.global_symbol "caml_reperform", dbg);
+        eff;
+        cont;
+        last_fiber ],
+      dbg )

backend/cmm_helpers.mli

@@ -964,6 +964,29 @@ val atomic_compare_and_set :
 
 val emit_gc_roots_table : symbols:symbol list -> phrase list -> phrase list
 
+val perform : dbg:Debuginfo.t -> expression -> expression
+
+val run_stack :
+  dbg:Debuginfo.t ->
+  stack:expression ->
+  f:expression ->
+  arg:expression ->
+  expression
+
+val resume :
+  dbg:Debuginfo.t ->
+  stack:expression ->
+  f:expression ->
+  arg:expression ->
+  expression
+
+val reperform :
+  dbg:Debuginfo.t ->
+  eff:expression ->
+  cont:expression ->
+  last_fiber:expression ->
+  expression
+
 (** Allocate a block to hold an unboxed float32 array for the given number of
     elements. *)
 val allocate_unboxed_float32_array :

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -404,10 +404,11 @@ module Inlining = struct
     let callee = Apply.callee apply in
     let region_inlined_into =
       match Apply.call_kind apply with
-      | Function { alloc_mode; _ } | Method { alloc_mode; _ } -> alloc_mode
-      | C_call _ ->
+      | Function { alloc_mode; _ } -> alloc_mode
+      | Method _ | C_call _ | Effect _ ->
         Misc.fatal_error
-          "Trying to call [Closure_conversion.Inlining.inline] on a C call."
+          "Trying to call [Closure_conversion.Inlining.inline] on a non-OCaml \
+           function call."
     in
     let args = Apply.args apply in
     let apply_return_continuation = Apply.continuation apply in
@@ -508,10 +509,7 @@ let close_c_call acc env ~loc ~let_bound_ids_with_kinds
   in
   let cost_metrics_of_body, free_names_of_body, acc, body =
     Acc.measure_cost_metrics acc ~f:(fun acc ->
-        k acc
-          (List.map
-             (fun var -> Named.create_simple (Simple.var var))
-             let_bound_vars))
+        k acc (List.map Named.create_var let_bound_vars))
   in
   let alloc_mode =
     match Lambda.alloc_mode_of_primitive_description prim_desc with
@@ -752,6 +750,82 @@ let close_raise acc env ~raise_kind ~arg ~dbg exn_continuation =
   let acc, arg = find_simple acc env arg in
   close_raise0 acc env ~raise_kind ~arg ~dbg exn_continuation
 
+let close_effect_primitive acc env ~dbg exn_continuation
+    (prim : Lambda.primitive) ~args ~let_bound_ids_with_kinds
+    (k : Acc.t -> Named.t list -> Expr_with_acc.t) : Expr_with_acc.t =
+  if not Config.runtime5
+  then Misc.fatal_error "Effect primitives are only supported on runtime5";
+  (* CR mshinwell: share with close_c_call, above *)
+  let _env, let_bound_vars =
+    List.fold_left_map
+      (fun env (id, kind) -> Env.add_var_like env id Not_user_visible kind)
+      env let_bound_ids_with_kinds
+  in
+  let let_bound_var =
+    match let_bound_vars with
+    | [let_bound_var] -> let_bound_var
+    | [] | _ :: _ :: _ ->
+      Misc.fatal_errorf
+        "close_effect_primitive: expected singleton return for primitive %a, \
+         but got: [%a]"
+        Printlambda.primitive prim
+        (Format.pp_print_list ~pp_sep:Format.pp_print_space Variable.print)
+        let_bound_vars
+  in
+  let continuation = Continuation.create () in
+  let return_kind = Flambda_kind.With_subkind.any_value in
+  let params =
+    [BP.create let_bound_var return_kind] |> Bound_parameters.create
+  in
+  let close call_kind =
+    let apply acc =
+      Apply_expr.create ~callee:None ~continuation:(Return continuation)
+        exn_continuation ~args:[] ~args_arity:Flambda_arity.nullary
+        ~return_arity:
+          (Flambda_arity.create_singletons
+             [Flambda_kind.With_subkind.any_value])
+        ~call_kind dbg ~inlined:Never_inlined
+        ~inlining_state:(Inlining_state.default ~round:0)
+        ~probe:None ~position:Normal
+        ~relative_history:Inlining_history.Relative.empty
+      |> Expr_with_acc.create_apply acc
+    in
+    Let_cont_with_acc.build_non_recursive acc continuation
+      ~handler_params:params
+      ~handler:(fun acc ->
+        let cost_metrics_of_body, free_names_of_body, acc, code_after_call =
+          Acc.measure_cost_metrics acc ~f:(fun acc ->
+              k acc (List.map Named.create_var let_bound_vars))
+        in
+        let acc =
+          Acc.with_cost_metrics
+            (Cost_metrics.( + ) (Acc.cost_metrics acc) cost_metrics_of_body)
+            (Acc.with_free_names free_names_of_body acc)
+        in
+        acc, code_after_call)
+      ~body:apply ~is_exn_handler:false ~is_cold:false
+  in
+  let module C = Call_kind in
+  let module E = C.Effect in
+  match[@ocaml.warning "-fragile-match"] prim, args with
+  | Pperform, [[eff]] ->
+    let call_kind = C.effect (E.perform ~eff) in
+    close call_kind
+  | Prunstack, [[stack]; [f]; [arg]] ->
+    let call_kind = C.effect (E.run_stack ~stack ~f ~arg) in
+    close call_kind
+  | Presume, [[stack]; [f]; [arg]] ->
+    let call_kind = C.effect (E.resume ~stack ~f ~arg) in
+    close call_kind
+  | Preperform, [[eff]; [cont]; [last_fiber]] ->
+    let call_kind = C.effect (E.reperform ~eff ~cont ~last_fiber) in
+    close call_kind
+  | _ ->
+    Misc.fatal_errorf
+      "close_effect_primitive: Wrong primitive and/or number of arguments: %a \
+       (%d args)"
+      Printlambda.primitive prim (List.length args)
+
 let close_primitive acc env ~let_bound_ids_with_kinds named
     (prim : Lambda.primitive) ~args loc
     (exn_continuation : IR.exn_continuation option) ~current_region
@@ -879,6 +953,17 @@ let close_primitive acc env ~let_bound_ids_with_kinds named
         assert false
     in
     k acc [Named.create_simple (Simple.symbol sym)]
+  | (Pperform | Prunstack | Presume | Preperform), args ->
+    let exn_continuation =
+      match exn_continuation with
+      | None ->
+        Misc.fatal_errorf
+          "Effect primitive is missing exception continuation: %a"
+          IR.print_named named
+      | Some exn_continuation -> exn_continuation
+    in
+    close_effect_primitive acc env ~dbg exn_continuation prim ~args
+      ~let_bound_ids_with_kinds k
   | prim, args ->
     Lambda_to_flambda_primitives.convert_and_bind acc exn_continuation
       ~big_endian:(Env.big_endian env) ~register_const0 prim ~args dbg

middle_end/flambda2/from_lambda/closure_conversion_aux.ml

@@ -957,8 +957,7 @@ module Expr_with_acc = struct
             { function_call = Indirect_unknown_arity | Indirect_known_arity; _ }
           ->
           false
-        | Method _ -> false
-        | C_call _ -> false)
+        | Method _ | C_call _ | Effect _ -> false)
     in
     let acc =
       match Apply.callee apply with

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -716,9 +716,7 @@ let primitive_can_raise (prim : Lambda.primitive) =
   | Punboxed_product_field _ | Pget_header _ ->
     false
   | Patomic_exchange | Patomic_cas | Patomic_fetch_add | Patomic_load _ -> false
-  | Prunstack | Pperform | Presume | Preperform ->
-    Misc.fatal_errorf "Primitive %a is not yet supported by Flambda 2"
-      Printlambda.primitive prim
+  | Prunstack | Pperform | Presume | Preperform -> true (* XXX! *)
   | Pdls_get -> false
 
 type non_tail_continuation =

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -1993,7 +1993,8 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
        %a (%a)"
       Printlambda.primitive prim H.print_list_of_lists_of_simple_or_prim args
   | ( ( Pignore | Psequand | Psequor | Pbytes_of_string | Pbytes_to_string
-      | Parray_of_iarray | Parray_to_iarray ),
+      | Parray_of_iarray | Parray_to_iarray | Prunstack | Pperform | Presume
+      | Preperform ),
       _ ) ->
     Misc.fatal_errorf
       "[%a] should have been removed by [Lambda_to_flambda.transform_primitive]"
@@ -2002,9 +2003,6 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
     Misc.fatal_errorf
       "[%a] should have been handled by [Closure_conversion.close_primitive]"
       Printlambda.primitive prim
-  | (Prunstack | Pperform | Presume | Preperform), _ ->
-    Misc.fatal_errorf "Primitive %a is not yet supported by Flambda 2"
-      Printlambda.primitive prim
 
 module Acc = Closure_conversion_aux.Acc
 module Expr_with_acc = Closure_conversion_aux.Expr_with_acc

middle_end/flambda2/kinds/flambda_arity.mli

@@ -55,7 +55,7 @@ val create : 'uc Component_for_creation.t list -> 'uc t
 val create_singletons : Flambda_kind.With_subkind.t list -> [> ] t
 
 (** "No parameters".  (Not e.g. "one parameter of type void".) *)
-val nullary : [> `Unarized] t
+val nullary : [> ] t
 
 (** How many parameters, potentially of unboxed product layout, the given
     arity describes. *)

middle_end/flambda2/parser/flambda_to_fexpr.ml

@@ -1040,6 +1040,7 @@ and apply_expr env (app : Apply_expr.t) : Fexpr.expr =
       Function (Indirect alloc)
     | C_call { needs_caml_c_call; _ } -> C_call { alloc = needs_caml_c_call }
     | Method _ -> Misc.fatal_error "TODO: Method call kind"
+    | Effect _ -> Misc.fatal_error "TODO: Effect call kind"
   in
   let param_arity = Apply_expr.args_arity app in
   let return_arity = Apply_expr.return_arity app in
@@ -1063,9 +1064,9 @@ and apply_expr env (app : Apply_expr.t) : Fexpr.expr =
       let params_arity = Some (complex_arity param_arity) in
       let ret_arity = arity return_arity in
       Some { params_arity; ret_arity }
-    | Function { function_call = Indirect_unknown_arity; alloc_mode = _ }
-    | Method _ ->
+    | Function { function_call = Indirect_unknown_arity; alloc_mode = _ } ->
       None
+    | Method _ | Effect _ -> assert false
   in
   let inlined : Fexpr.inlined_attribute option =
     if Flambda2_terms.Inlined_attribute.is_default (Apply_expr.inlined app)

middle_end/flambda2/simplify/inlining/inlining_transforms.ml

@@ -95,10 +95,11 @@ let inline dacc ~apply ~unroll_to ~was_inline_always function_decl =
   let callee = Apply.callee apply in
   let region_inlined_into =
     match Apply.call_kind apply with
-    | Function { alloc_mode; _ } | Method { alloc_mode; _ } -> alloc_mode
-    | C_call _ ->
+    | Function { alloc_mode; _ } -> alloc_mode
+    | Method _ | C_call _ | Effect _ ->
       Misc.fatal_error
-        "Trying to call [Inlining_transforms.inline] on a C call."
+        "Trying to call [Inlining_transforms.inline] on something other than \
+         an OCaml function call."
   in
   let args = Apply.args apply in
   let apply_return_continuation = Apply.continuation apply in