Add new primitives to create blocks for lazy values

This adds a new primitive in Lambda and Flambda2 to create a block
for a (not short-circuited) lazy value. These primitives allow to
more systematically handle lazy blocks correctly, for instance with
respect to the value-let-rec size computation, and in flambda to better
ensure that we do not track block appromixations for lazy values
(without relying on the frontend to insert opaque identities).

Author: Gbury

bug.ml

@@ -0,0 +1,9 @@
+
+type t =
+  | C of t Lazy.t Lazy.t
+
+let rec x =
+  let y = (lazy (C x)) in
+  lazy y
+
+

bytecomp/bytegen.ml

@@ -153,7 +153,7 @@ let preserve_tailcall_for_prim = function
   | Pget_header _
   | Pignore
   | Pgetglobal _ | Psetglobal _ | Pgetpredef _
-  | Pmakeblock _ | Pmakefloatblock _ | Pmakeufloatblock _ | Pmakemixedblock _
+  | Pmakeblock _ | Pmakefloatblock _ | Pmakeufloatblock _ | Pmakemixedblock _ | Pmakelazyblock _
   | Pfield _ | Pfield_computed _ | Psetfield _
   | Psetfield_computed _ | Pfloatfield _ | Psetfloatfield _ | Pduprecord _
   | Pufloatfield _ | Psetufloatfield _ | Pmixedfield _ | Psetmixedfield _
@@ -741,6 +741,7 @@ let comp_primitive stack_info p sz args =
   | Pmakefloatblock _
   | Pmakeufloatblock _
   | Pmakemixedblock _
+  | Pmakelazyblock _
   | Pprobe_is_enabled _
   | Punbox_float _ | Pbox_float (_, _) | Punbox_int _ | Pbox_int _
     ->
@@ -1137,6 +1138,14 @@ and comp_expr stack_info env exp sz cont =
       let cont = add_pseudo_event loc !compunit_name cont in
       comp_args stack_info env args sz
         (Kmakeblock(List.length args, tag) :: cont)
+  | Lprim (Pmakelazyblock (Lazy_tag, _), args, loc) ->
+      let cont = add_pseudo_event loc !compunit_name cont in
+      comp_args stack_info env args sz
+        (Kmakeblock(List.length args, Config.lazy_tag) :: cont)
+  | Lprim (Pmakelazyblock (Forward_tag, _), args, loc) ->
+      let cont = add_pseudo_event loc !compunit_name cont in
+      comp_args stack_info env args sz
+        (Kmakeblock(List.length args, Obj.forward_tag) :: cont)
   | Lprim(Pmake_unboxed_product _, args, loc) ->
       let cont = add_pseudo_event loc !compunit_name cont in
       comp_args stack_info env args sz

lambda/lambda.ml

@@ -124,6 +124,10 @@ type region_close =
   | Rc_nontail
   | Rc_close_at_apply
 
+type lazy_block_tag =
+  | Lazy_tag
+  | Forward_tag
+
 type primitive =
   | Pbytes_to_string
   | Pbytes_of_string
@@ -137,6 +141,7 @@ type primitive =
   | Pmakefloatblock of mutable_flag * locality_mode
   | Pmakeufloatblock of mutable_flag * locality_mode
   | Pmakemixedblock of int * mutable_flag * mixed_block_shape * locality_mode
+  | Pmakelazyblock of lazy_block_tag * locality_mode
   | Pfield of int * immediate_or_pointer * field_read_semantics
   | Pfield_computed of field_read_semantics
   | Psetfield of int * immediate_or_pointer * initialization_or_assignment
@@ -1823,6 +1828,7 @@ let primitive_may_allocate : primitive -> locality_mode option = function
   | Pmakefloatblock (_, m) -> Some m
   | Pmakeufloatblock (_, m) -> Some m
   | Pmakemixedblock (_, _, _, m) -> Some m
+  | Pmakelazyblock (_, m) -> Some m
   | Pfield _ | Pfield_computed _ | Psetfield _ | Psetfield_computed _ -> None
   | Pfloatfield (_, _, m) -> Some m
   | Pufloatfield _ -> None
@@ -2056,7 +2062,7 @@ let primitive_can_raise prim =
   | Pmakefloatblock _ | Pfield _ | Pfield_computed _ | Psetfield _
   | Psetfield_computed _ | Pfloatfield _ | Psetfloatfield _ | Pduprecord _
   | Pmakeufloatblock _ | Pufloatfield _ | Psetufloatfield _ | Psequand | Psequor
-  | Pmixedfield _ | Psetmixedfield _ | Pmakemixedblock _ | Pnot | Pnegint
+  | Pmixedfield _ | Psetmixedfield _ | Pmakemixedblock _ | Pmakelazyblock _ | Pnot | Pnegint
   | Paddint | Psubint | Pmulint | Pandint | Porint | Pxorint | Plslint | Plsrint
   | Pasrint | Pintcomp _ | Pcompare_ints | Pcompare_floats _ | Pcompare_bints _
   | Poffsetint _ | Poffsetref _ | Pintoffloat _
@@ -2254,7 +2260,7 @@ let primitive_result_layout (p : primitive) =
     -> layout_unit
   | Pgetglobal _ | Psetglobal _ | Pgetpredef _ -> layout_module_field
   | Pmakeblock _ | Pmakefloatblock _ | Pmakearray _ | Pmakearray_dynamic _
-  | Pduprecord _ | Pmakeufloatblock _ | Pmakemixedblock _
+  | Pduprecord _ | Pmakeufloatblock _ | Pmakemixedblock _ | Pmakelazyblock _
   | Pduparray _ | Pbigarraydim _ | Pobj_dup -> layout_block
   | Pfield _ | Pfield_computed _ -> layout_value_field
   | Punboxed_product_field (field, layouts) -> (Array.of_list layouts).(field)

lambda/lambda.mli

@@ -104,6 +104,10 @@ type region_close =
     tail call because the outer region needs to end there.)
 *)
 
+type lazy_block_tag =
+  | Lazy_tag
+  | Forward_tag
+
 (* CR layouts v5: When we add more blocks of non-scannable values, consider
    whether some of the primitives specific to ufloat records
    ([Pmakeufloatblock], [Pufloatfield], and [Psetufloatfield]) can/should be
@@ -121,6 +125,7 @@ type primitive =
   | Pmakefloatblock of mutable_flag * locality_mode
   | Pmakeufloatblock of mutable_flag * locality_mode
   | Pmakemixedblock of int * mutable_flag * mixed_block_shape * locality_mode
+  | Pmakelazyblock of lazy_block_tag * locality_mode
   | Pfield of int * immediate_or_pointer * field_read_semantics
   | Pfield_computed of field_read_semantics
   | Psetfield of int * immediate_or_pointer * initialization_or_assignment

lambda/matching.ml

@@ -2124,7 +2124,7 @@ let inline_lazy_force_cond arg pos loc =
     ( Strict,
       Lambda.layout_lazy,
       idarg,
-      arg,
+      Lprim (Popaque Lambda.layout_lazy, [arg], loc),
       Llet
         ( Alias,
           Lambda.layout_int,

lambda/printlambda.ml

@@ -469,6 +469,10 @@ let primitive ppf = function
   | Pmakemixedblock (tag, Mutable, abs, mode) ->
      fprintf ppf "make%amixedblock %i Mutable%a"
         locality_mode mode tag (mixed_block_shape (fun _ _ -> ())) abs
+  | Pmakelazyblock (Lazy_tag, mode) ->
+      fprintf ppf "make%alazyblock" locality_mode mode
+  | Pmakelazyblock (Forward_tag, mode) ->
+      fprintf ppf "make%aforwardblock" locality_mode mode
   | Pfield (n, ptr, sem) ->
       let instr =
         match ptr, sem with
@@ -918,6 +922,7 @@ let name_of_primitive = function
   | Pmakefloatblock _ -> "Pmakefloatblock"
   | Pmakeufloatblock _ -> "Pmakeufloatblock"
   | Pmakemixedblock _ -> "Pmakemixedblock"
+  | Pmakelazyblock _ -> "Pmakelazyblock"
   | Pfield _ -> "Pfield"
   | Pfield_computed _ -> "Pfield_computed"
   | Psetfield _ -> "Psetfield"

lambda/tmc.ml

@@ -910,6 +910,9 @@ let rec choice ctx t =
     | Punbox_int _ | Pbox_int _
     | Punbox_vector _ | Pbox_vector (_, _)
 
+    (* it doesn't seem worth it to support lazy blocks for tmc *)
+    | Pmakelazyblock _
+
     (* we don't handle array indices as destinations yet *)
     | (Pmakearray _ | Pduparray _ | Pmakearray_dynamic _)
 

lambda/translcore.ml

@@ -1012,8 +1012,7 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
           (* We don't need to wrap with Popaque: this forward
              block will never be shortcutted since it points to a float
              and Config.flat_float_array is true. *)
-         Lprim(Pmakeblock(Obj.forward_tag, Immutable, None,
-                          alloc_heap),
+         Lprim(Pmakelazyblock(Forward_tag, alloc_heap),
                 [transl_exp ~scopes Jkind.Sort.Const.for_lazy_body e],
                of_location ~scopes e.exp_loc)
       | `Identifier `Forward_value ->
@@ -1023,11 +1022,8 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
             optimisation in Flambda, but the concept of a mutable
             block doesn't really match what is going on here.  This
             value may subsequently turn into an immediate... *)
-         Lprim (Popaque Lambda.layout_lazy,
-                [Lprim(Pmakeblock(Obj.forward_tag, Immutable, None,
-                                  alloc_heap),
-                       [transl_exp ~scopes Jkind.Sort.Const.for_lazy_body e],
-                       of_location ~scopes e.exp_loc)],
+         Lprim(Pmakelazyblock(Forward_tag, alloc_heap),
+                [transl_exp ~scopes Jkind.Sort.Const.for_lazy_body e],
                 of_location ~scopes e.exp_loc)
       | `Identifier `Other ->
          transl_exp ~scopes Jkind.Sort.Const.for_lazy_body e
@@ -1053,7 +1049,7 @@ and transl_exp0 ~in_new_scope ~scopes sort e =
                                      Lambda.layout_lazy_contents
                                      (transl_exp ~scopes Jkind.Sort.Const.for_lazy_body e))
          in
-          Lprim(Pmakeblock(Config.lazy_tag, Mutable, None, alloc_heap), [fn],
+          Lprim(Pmakelazyblock(Lazy_tag, alloc_heap), [fn],
                 of_location ~scopes e.exp_loc)
       end
   | Texp_object (cs, meths) ->

lambda/translprim.ml

@@ -1929,7 +1929,7 @@ let lambda_primitive_needs_event_after = function
   | Parray_to_iarray | Parray_of_iarray
   | Pignore | Psetglobal _
   | Pgetglobal _ | Pgetpredef _ | Pmakeblock _ | Pmakefloatblock _
-  | Pmakeufloatblock _ | Pmakemixedblock _
+  | Pmakeufloatblock _ | Pmakemixedblock _ | Pmakelazyblock _
   | Pmake_unboxed_product _ | Punboxed_product_field _
   | Parray_element_size_in_bytes _
   | Pfield _ | Pfield_computed _ | Psetfield _

lambda/value_rec_compiler.ml

@@ -245,7 +245,7 @@ let compute_static_size lam =
         | Record_inlined (_, _, (Variant_unboxed | Variant_with_null)) ->
             Misc.fatal_error "size_of_primitive"
         end
-    | Pmakeblock _ ->
+    | Pmakeblock _ | Pmakelazyblock _ ->
         (* The block shape is unfortunately an option, so we rely on the
            number of arguments instead.
            Note that flat float arrays/records use Pmakearray, so we don't need

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -1053,7 +1053,8 @@ let close_primitive acc env ~let_bound_ids_with_kinds named
       | Patomic_fetch_add | Patomic_add | Patomic_sub | Patomic_land
       | Patomic_lor | Patomic_lxor | Pdls_get | Ppoll | Patomic_load _
       | Patomic_set _ | Preinterpret_tagged_int63_as_unboxed_int64
-      | Preinterpret_unboxed_int64_as_tagged_int63 | Ppeek _ | Ppoke _ ->
+      | Preinterpret_unboxed_int64_as_tagged_int63 | Ppeek _ | Ppoke _
+      | Pmakelazyblock _ ->
         (* Inconsistent with outer match *)
         assert false
     in

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -1382,6 +1382,9 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
     let shape = convert_block_shape shape ~num_fields:(List.length args) in
     let mutability = Mutability.from_lambda mutability in
     [Variadic (Make_block (Values (tag, shape), mutability, mode), args)]
+  | Pmakelazyblock (lazy_tag, mode), [[arg]] ->
+    let mode = Alloc_mode.For_allocations.from_lambda mode ~current_region in
+    [Unary (Make_lazy (lazy_tag, mode), arg)]
   | Pmake_unboxed_product layouts, _ ->
     (* CR mshinwell: this should check the unarized lengths of [layouts] and
        [args] (like [Parray_element_size_in_bytes] below) *)
@@ -2492,7 +2495,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
       | Pufloatfield _ | Patomic_load _ | Pmixedfield _
       | Preinterpret_unboxed_int64_as_tagged_int63
       | Preinterpret_tagged_int63_as_unboxed_int64
-      | Parray_element_size_in_bytes _ | Ppeek _ ),
+      | Parray_element_size_in_bytes _ | Ppeek _ | Pmakelazyblock _ ),
       ([] | _ :: _ :: _ | [([] | _ :: _ :: _)]) ) ->
     Misc.fatal_errorf
       "Closure_conversion.convert_primitive: Wrong arity for unary primitive \

middle_end/flambda2/parser/flambda_to_fexpr.ml

@@ -582,7 +582,7 @@ let unop env (op : Flambda_primitive.unary_primitive) : Fexpr.unop =
   | Boolean_not -> Boolean_not
   | Int_as_pointer _ | Duplicate_block _ | Duplicate_array _ | Bigarray_length _
   | Float_arith _ | Reinterpret_64_bit_word _ | Is_boxed_float | Obj_dup
-  | Get_header | Atomic_load _ | Peek _ ->
+  | Get_header | Atomic_load _ | Peek _ | Make_lazy _ ->
     Misc.fatal_errorf "TODO: Unary primitive: %a"
       Flambda_primitive.Without_args.print
       (Flambda_primitive.Without_args.Unary op)

middle_end/flambda2/simplify/simplify_unary_primitive.ml

@@ -888,6 +888,12 @@ let simplify_mutable_block_load _access_kind ~field:_ ~original_prim dacc
       (P.result_kind' original_prim)
       ~original_term
 
+let simplify_lazy ~original_prim dacc ~original_term ~arg:_ ~arg_ty:_
+    ~result_var =
+  SPR.create_unknown dacc ~result_var
+    (P.result_kind' original_prim)
+    ~original_term
+
 (* CR layouts v3: implement a real simplifier. *)
 let simplify_is_null dacc ~original_term ~arg:scrutinee ~arg_ty:scrutinee_ty
     ~result_var =
@@ -962,5 +968,6 @@ let simplify_unary_primitive dacc original_prim (prim : P.unary_primitive) ~arg
     | Atomic_load block_access_field_kind ->
       simplify_atomic_load block_access_field_kind ~original_prim
     | Peek _ -> simplify_peek ~original_prim
+    | Make_lazy _ -> simplify_lazy ~original_prim
   in
   simplifier dacc ~original_term ~arg ~arg_ty ~result_var

middle_end/flambda2/terms/code_size.ml

@@ -380,6 +380,7 @@ let unary_prim_size prim =
   | Obj_dup -> needs_caml_c_call_extcall_size + 1
   | Get_header -> 2
   | Atomic_load _ | Peek _ -> 1
+  | Make_lazy (_, _alloc_mode) -> alloc_size
 
 let binary_prim_size prim =
   match (prim : Flambda_primitive.binary_primitive) with