Fix bug in lazy pattern matching (#3647)

* 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).

* cleanup

* Update code size of Make_lazy to match Make_block

* Remove allocation mode from lazy primitives

* Formatting

---------

Co-authored-by: Vincent Laviron <[email protected]>

Author: Gbury

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
   | 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 _ -> Some alloc_heap
   | 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
   | 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/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 ->
+      fprintf ppf "makelazyblock"
+  | Pmakelazyblock Forward_tag ->
+      fprintf ppf "makeforwardblock"
   | 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,
                 [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,
+                [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, [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

@@ -1375,6 +1375,7 @@ 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, [[arg]] -> [Unary (Make_lazy lazy_tag, arg)]
   | Pmake_unboxed_product layouts, _ ->
     (* CR mshinwell: this should check the unarized lengths of [layouts] and
        [args] (like [Parray_element_size_in_bytes] below) *)
@@ -2485,7 +2486,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_size + 1
 
 let binary_prim_size prim =
   match (prim : Flambda_primitive.binary_primitive) with

middle_end/flambda2/terms/flambda_primitive.ml

@@ -25,6 +25,26 @@ type classification_for_printing =
   | Destructive
   | Neither
 
+module Lazy_block_tag = struct
+  type t = Lambda.lazy_block_tag =
+    | Lazy_tag
+    | Forward_tag
+
+  let print ppf t =
+    match t with
+    | Lazy_tag -> Format.fprintf ppf "Lazy_block"
+    | Forward_tag -> Format.fprintf ppf "Forward_block"
+
+  let compare t1 t2 =
+    match t1, t2 with
+    | Lazy_tag, Lazy_tag | Forward_tag, Forward_tag -> 0
+    | Lazy_tag, Forward_tag -> -1
+    | Forward_tag, Lazy_tag -> 1
+
+  let to_tag t =
+    match t with Lazy_tag -> Tag.lazy_tag | Forward_tag -> Tag.forward_tag
+end
+
 module Block_kind = struct
   type t =
     | Values of Tag.Scannable.t * K.With_subkind.t list
@@ -57,7 +77,7 @@ module Block_kind = struct
          @[<hov 1>(shape@ @[<hov 1>(%a)@])@])@]"
        Tag.Scannable.print tag
        (Format.pp_print_list ~pp_sep:Format.pp_print_space
-      K.print) (Array.to_list (K.Mixed_block_shape.field_kinds shape))
+          K.print) (Array.to_list (K.Mixed_block_shape.field_kinds shape))
 
   let compare t1 t2 =
     match t1, t2 with
@@ -1060,6 +1080,7 @@ type unary_primitive =
   | Get_header
   | Atomic_load of Block_access_field_kind.t
   | Peek of Flambda_kind.Standard_int_or_float.t
+  | Make_lazy of Lazy_block_tag.t
 
 (* Here and below, operations that are genuine projections shouldn't be eligible
    for CSE, since we deal with projections through types. *)
@@ -1091,7 +1112,9 @@ let unary_primitive_eligible_for_cse p ~arg =
     Simple.is_var arg
   | Project_function_slot _ | Project_value_slot _ -> false
   | Is_boxed_float | Is_flat_float_array -> true
-  | End_region _ | End_try_region _ | Obj_dup | Atomic_load _ | Peek _ -> false
+  | End_region _ | End_try_region _ | Obj_dup | Atomic_load _ | Peek _
+  | Make_lazy _ ->
+    false
 
 let compare_unary_primitive p1 p2 =
   let unary_primitive_numbering p =
@@ -1126,6 +1149,7 @@ let compare_unary_primitive p1 p2 =
     | Atomic_load _ -> 27
     | Is_null -> 28
     | Peek _ -> 29
+    | Make_lazy _ -> 30
   in
   match p1, p2 with
   | ( Block_load { kind = kind1; mut = mut1; field = field1 },
@@ -1212,14 +1236,16 @@ let compare_unary_primitive p1 p2 =
     Bool.compare ghost1 ghost2
   | Peek kind1, Peek kind2 ->
     Flambda_kind.Standard_int_or_float.compare kind1 kind2
+  | Make_lazy lazy_tag1, Make_lazy lazy_tag2 ->
+    Lazy_block_tag.compare lazy_tag1 lazy_tag2
   | ( ( Block_load _ | Duplicate_array _ | Duplicate_block _ | Is_int _
       | Is_null | Get_tag | String_length _ | Int_as_pointer _
       | Opaque_identity _ | Int_arith _ | Num_conv _ | Boolean_not
       | Reinterpret_64_bit_word _ | Float_arith _ | Array_length _
       | Bigarray_length _ | Unbox_number _ | Box_number _ | Untag_immediate
       | Tag_immediate | Project_function_slot _ | Project_value_slot _
       | Is_boxed_float | Is_flat_float_array | End_region _ | End_try_region _
-      | Obj_dup | Get_header | Atomic_load _ | Peek _ ),
+      | Obj_dup | Get_header | Atomic_load _ | Peek _ | Make_lazy _ ),
       _ ) ->
     Stdlib.compare (unary_primitive_numbering p1) (unary_primitive_numbering p2)
 
@@ -1289,6 +1315,8 @@ let print_unary_primitive ppf p =
   | Peek kind ->
     fprintf ppf "@[(Peek@ %a)@]"
       Flambda_kind.Standard_int_or_float.print_lowercase kind
+  | Make_lazy lazy_tag ->
+    fprintf ppf "@[<hov 1>(Make_lazy@ %a)@]" Lazy_block_tag.print lazy_tag
 
 let arg_kind_of_unary_primitive p =
   match p with
@@ -1324,6 +1352,7 @@ let arg_kind_of_unary_primitive p =
   | Get_header -> K.value
   | Atomic_load _ -> K.value
   | Peek _ -> K.naked_nativeint
+  | Make_lazy _ -> K.value
 
 let result_kind_of_unary_primitive p : result_kind =
   match p with
@@ -1362,6 +1391,7 @@ let result_kind_of_unary_primitive p : result_kind =
   | Get_header -> Singleton K.naked_nativeint
   | Atomic_load _ -> Singleton K.value
   | Peek kind -> Singleton (K.Standard_int_or_float.to_kind kind)
+  | Make_lazy _ -> Singleton K.value
 
 let effects_and_coeffects_of_unary_primitive p : Effects_and_coeffects.t =
   match p with
@@ -1453,6 +1483,7 @@ let effects_and_coeffects_of_unary_primitive p : Effects_and_coeffects.t =
   | Atomic_load _ | Peek _ ->
     (* For the moment, prevent [Peek] from being moved. *)
     Arbitrary_effects, Has_coeffects, Strict
+  | Make_lazy _ -> Only_generative_effects Mutable, No_coeffects, Strict
 
 let unary_classify_for_printing p =
   match p with
@@ -1471,6 +1502,7 @@ let unary_classify_for_printing p =
   | End_region _ | End_try_region _ -> Neither
   | Get_header -> Neither
   | Peek _ -> Neither
+  | Make_lazy _ -> Constructive
 
 let free_names_unary_primitive p =
   match p with
@@ -1493,7 +1525,8 @@ let free_names_unary_primitive p =
   | Is_boxed_float | Is_flat_float_array | End_region _ | End_try_region _
   | Obj_dup | Get_header
   | Atomic_load (_ : Block_access_field_kind.t)
-  | Peek (_ : Flambda_kind.Standard_int_or_float.t) ->
+  | Peek (_ : Flambda_kind.Standard_int_or_float.t)
+  | Make_lazy _ ->
     Name_occurrences.empty
 
 let apply_renaming_unary_primitive p renaming =
@@ -1515,7 +1548,8 @@ let apply_renaming_unary_primitive p renaming =
   | Is_boxed_float | Is_flat_float_array | End_region _ | End_try_region _
   | Project_function_slot _ | Project_value_slot _ | Obj_dup | Get_header
   | Atomic_load (_ : Block_access_field_kind.t)
-  | Peek (_ : Flambda_kind.Standard_int_or_float.t) ->
+  | Peek (_ : Flambda_kind.Standard_int_or_float.t)
+  | Make_lazy _ ->
     p
 
 let ids_for_export_unary_primitive p =
@@ -1529,7 +1563,8 @@ let ids_for_export_unary_primitive p =
   | Is_boxed_float | Is_flat_float_array | End_region _ | End_try_region _
   | Project_function_slot _ | Project_value_slot _ | Obj_dup | Get_header
   | Atomic_load (_ : Block_access_field_kind.t)
-  | Peek (_ : Flambda_kind.Standard_int_or_float.t) ->
+  | Peek (_ : Flambda_kind.Standard_int_or_float.t)
+  | Make_lazy _ ->
     Ids_for_export.empty
 
 type binary_int_arith_op =