Remove region from Flambda_primitive.Init_or_assign.t (ocaml-flambd…

…a#1326)

The region in the `Alloc_mode.For_allocations.t` in an `Init_or_assign.t` was
always ignored (and meaningless). Replaced it with the new
`Alloc_mode.For_assignments.t`, which is simply either `Heap` or `Local`.

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -90,12 +90,10 @@ let standard_int_or_float_of_boxed_integer (bint : L.boxed_integer) :
 let convert_block_access_field_kind i_or_p : P.Block_access_field_kind.t =
   match i_or_p with L.Immediate -> Immediate | L.Pointer -> Any_value
 
-let convert_init_or_assign (i_or_a : L.initialization_or_assignment)
-    ~current_region : P.Init_or_assign.t =
+let convert_init_or_assign (i_or_a : L.initialization_or_assignment) :
+    P.Init_or_assign.t =
   match i_or_a with
-  | Assignment mode ->
-    Assignment
-      (Alloc_mode.For_allocations.from_lambda_modify mode ~current_region)
+  | Assignment mode -> Assignment (Alloc_mode.For_assignments.from_lambda mode)
   | Heap_initialization -> Initialization
   | Root_initialization ->
     Misc.fatal_error "[Root_initialization] should not appear in Flambda input"
@@ -503,21 +501,19 @@ let array_load_unsafe ~array ~index (array_kind : P.Array_kind.t)
       (Binary (Array_load (Naked_floats, Mutable), array, index))
       ~current_region
 
-let array_set_unsafe ~array ~index ~new_value (array_kind : P.Array_kind.t)
-    ~current_region : H.expr_primitive =
+let array_set_unsafe ~array ~index ~new_value (array_kind : P.Array_kind.t) :
+    H.expr_primitive =
   match array_kind with
   | Immediates | Values ->
     Ternary
-      ( Array_set (array_kind, Assignment Alloc_mode.For_allocations.heap),
+      ( Array_set (array_kind, Assignment Alloc_mode.For_assignments.heap),
         array,
         index,
         new_value )
   | Naked_floats ->
     Ternary
       ( Array_set
-          ( Naked_floats,
-            Assignment (Alloc_mode.For_allocations.local ~region:current_region)
-          ),
+          (Naked_floats, Assignment (Alloc_mode.For_assignments.local ())),
         array,
         index,
         unbox_float new_value )
@@ -754,8 +750,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list)
       Values { tag = Unknown; size = Unknown; field_kind }
     in
     Ternary
-      ( Block_set
-          (block_access, convert_init_or_assign init_or_assign ~current_region),
+      ( Block_set (block_access, convert_init_or_assign init_or_assign),
         obj,
         field,
         value )
@@ -931,9 +926,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list)
     let imm = Targetint_31_63.of_int index in
     check_non_negative_imm imm "Psetfield";
     let field = Simple.const (Reg_width_const.tagged_immediate imm) in
-    let init_or_assign =
-      convert_init_or_assign initialization_or_assignment ~current_region
-    in
+    let init_or_assign = convert_init_or_assign initialization_or_assignment in
     let block_access : P.Block_access_kind.t =
       Values { tag = Unknown; size = Unknown; field_kind }
     in
@@ -946,9 +939,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list)
     let block_access : P.Block_access_kind.t =
       Naked_floats { size = Unknown }
     in
-    let init_or_assign =
-      convert_init_or_assign initialization_or_assignment ~current_region
-    in
+    let init_or_assign = convert_init_or_assign initialization_or_assignment in
     Ternary
       ( Block_set (block_access, init_or_assign),
         block,
@@ -990,11 +981,11 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list)
          (array_load_unsafe ~array ~index ~current_region))
   | Parraysetu array_kind, [array; index; new_value] ->
     match_on_array_kind ~array array_kind
-      (array_set_unsafe ~array ~index ~new_value ~current_region)
+      (array_set_unsafe ~array ~index ~new_value)
   | Parraysets array_kind, [array; index; new_value] ->
     check_array_access ~dbg ~array ~index
       (match_on_array_kind ~array array_kind
-         (array_set_unsafe ~array ~index ~new_value ~current_region))
+         (array_set_unsafe ~array ~index ~new_value))
   | Pbytessetu (* unsafe *), [bytes; index; new_value] ->
     bytes_like_set_unsafe ~access_size:Eight Bytes bytes index new_value
   | Pbytessets, [bytes; index; new_value] ->
@@ -1022,9 +1013,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list)
     in
     Ternary
       ( Block_set
-          ( block_access,
-            Assignment (Alloc_mode.For_allocations.local ~region:current_region)
-          ),
+          (block_access, Assignment (Alloc_mode.For_assignments.local ())),
         block,
         Simple Simple.const_zero,
         new_ref_value )

middle_end/flambda2/parser/fexpr.ml

@@ -225,9 +225,13 @@ type alloc_mode_for_types =
   | Heap_or_local
   | Local
 
+type alloc_mode_for_assignments =
+  | Heap
+  | Local
+
 type init_or_assign =
   | Initialization
-  | Assignment of alloc_mode_for_allocations
+  | Assignment of alloc_mode_for_assignments
 
 type 'signed_or_unsigned comparison =
       'signed_or_unsigned Flambda_primitive.comparison =

middle_end/flambda2/parser/fexpr_to_flambda.ml

@@ -364,11 +364,16 @@ let alloc_mode_for_types (alloc : Fexpr.alloc_mode_for_types) =
   | Heap_or_local -> Alloc_mode.For_types.unknown ()
   | Local -> Alloc_mode.For_types.local ()
 
-let init_or_assign env (ia : Fexpr.init_or_assign) :
+let alloc_mode_for_assignments (alloc : Fexpr.alloc_mode_for_assignments) =
+  match alloc with
+  | Heap -> Alloc_mode.For_assignments.heap
+  | Local -> Alloc_mode.For_assignments.local ()
+
+let init_or_assign _env (ia : Fexpr.init_or_assign) :
     Flambda_primitive.Init_or_assign.t =
   match ia with
   | Initialization -> Initialization
-  | Assignment alloc -> Assignment (alloc_mode_for_allocations env alloc)
+  | Assignment alloc -> Assignment (alloc_mode_for_assignments alloc)
 
 let nullop (nullop : Fexpr.nullop) : Flambda_primitive.nullary_primitive =
   match nullop with Begin_region -> Begin_region

middle_end/flambda2/parser/flambda_parser.mly

@@ -496,7 +496,7 @@ convertible_type:
 init_or_assign:
   | EQUAL { Initialization }
   | LESSMINUS { Assignment Heap }
-  | LESSMINUS AMP; region = region { Assignment (Local { region }) }
+  | LESSMINUS AMP { Assignment Local }
 
 alloc_mode_for_types_opt:
   | { Heap }

middle_end/flambda2/parser/flambda_to_fexpr.ml

@@ -483,6 +483,10 @@ let alloc_mode_for_allocations env (alloc : Alloc_mode.For_allocations.t) :
     let r = Env.find_region_exn env r in
     Local { region = r }
 
+let alloc_mode_for_assignments _env (alloc : Alloc_mode.For_assignments.t) :
+    Fexpr.alloc_mode_for_assignments =
+  match alloc with Heap -> Heap | Local -> Local
+
 let alloc_mode_for_types (alloc : Alloc_mode.For_types.t) :
     Fexpr.alloc_mode_for_types =
   match alloc with
@@ -494,7 +498,7 @@ let init_or_assign env (ia : Flambda_primitive.Init_or_assign.t) :
     Fexpr.init_or_assign =
   match ia with
   | Initialization -> Initialization
-  | Assignment alloc -> Assignment (alloc_mode_for_allocations env alloc)
+  | Assignment alloc -> Assignment (alloc_mode_for_assignments env alloc)
 
 let nullop _env (op : Flambda_primitive.nullary_primitive) : Fexpr.nullop =
   match op with

middle_end/flambda2/parser/print_fexpr.ml

@@ -303,10 +303,8 @@ let alloc_mode_for_types_opt ppf (alloc : alloc_mode_for_types) ~space =
 let init_or_assign ppf ia =
   match ia with
   | Initialization -> Format.pp_print_string ppf "="
-  | Assignment alloc ->
-    Format.fprintf ppf "@[<h><-%a@]"
-      (alloc_mode_for_allocations_opt ~space:Before)
-      alloc
+  | Assignment Heap -> Format.pp_print_string ppf "<-"
+  | Assignment Local -> Format.pp_print_string ppf "<-&"
 
 let boxed_variable ppf var ~kind =
   Format.fprintf ppf "%a : %s boxed" variable var kind

middle_end/flambda2/term_basics/alloc_mode.ml

@@ -98,21 +98,6 @@ module For_allocations = struct
       assert (Flambda_features.stack_allocation_enabled ());
       Lambda.alloc_local
 
-  let from_lambda_modify (mode : Lambda.modify_mode) ~current_region =
-    if not (Flambda_features.stack_allocation_enabled ())
-    then Heap
-    else
-      match mode with
-      | Modify_heap -> Heap
-      | Modify_maybe_stack -> Local { region = current_region }
-
-  let to_lambda_modify t =
-    match t with
-    | Heap -> Lambda.modify_heap
-    | Local _ ->
-      assert (Flambda_features.stack_allocation_enabled ());
-      Lambda.modify_maybe_stack
-
   let free_names t =
     match t with
     | Heap -> Name_occurrences.empty
@@ -131,3 +116,38 @@ module For_allocations = struct
     | Heap -> Ids_for_export.empty
     | Local { region } -> Ids_for_export.singleton_variable region
 end
+
+module For_assignments = struct
+  type t =
+    | Heap
+    | Local
+
+  let print ppf t =
+    match t with
+    | Heap -> Format.pp_print_string ppf "Heap"
+    | Local -> Format.pp_print_string ppf "Local"
+
+  let compare t1 t2 =
+    match t1, t2 with
+    | Heap, Heap -> 0
+    | Local, Local -> 0
+    | Heap, Local -> -1
+    | Local, Heap -> 1
+
+  let heap = Heap
+
+  let local () =
+    if Flambda_features.stack_allocation_enabled () then Local else Heap
+
+  let from_lambda (mode : Lambda.modify_mode) =
+    if not (Flambda_features.stack_allocation_enabled ())
+    then Heap
+    else match mode with Modify_heap -> Heap | Modify_maybe_stack -> Local
+
+  let to_lambda t =
+    match t with
+    | Heap -> Lambda.modify_heap
+    | Local ->
+      assert (Flambda_features.stack_allocation_enabled ());
+      Lambda.modify_maybe_stack
+end

middle_end/flambda2/term_basics/alloc_mode.mli

@@ -61,11 +61,27 @@ module For_allocations : sig
 
   val to_lambda : t -> Lambda.alloc_mode
 
-  val from_lambda_modify : Lambda.modify_mode -> current_region:Variable.t -> t
-
-  val to_lambda_modify : t -> Lambda.modify_mode
-
   include Contains_names.S with type t := t
 
   include Contains_ids.S with type t := t
 end
+
+module For_assignments : sig
+  (** Decisions on assignment locations *)
+  type t = private
+    | Heap
+    | Local
+
+  val print : Format.formatter -> t -> unit
+
+  val compare : t -> t -> int
+
+  val heap : t
+
+  (* Returns [Heap] if stack allocation is disabled! *)
+  val local : unit -> t
+
+  val from_lambda : Lambda.modify_mode -> t
+
+  val to_lambda : t -> Lambda.modify_mode
+end

middle_end/flambda2/terms/code_size.ml

@@ -145,14 +145,14 @@ let block_set (kind : Flambda_primitive.Block_access_kind.t)
     (init : Flambda_primitive.Init_or_assign.t) =
   match kind, init with
   | Values _, Assignment Heap -> nonalloc_extcall_size (* caml_modify *)
-  | Values _, (Assignment (Local _) | Initialization) -> 1 (* cadda + store *)
+  | Values _, (Assignment Local | Initialization) -> 1 (* cadda + store *)
   | Naked_floats _, (Assignment _ | Initialization) -> 1
 
 let array_set (kind : Flambda_primitive.Array_kind.t)
     (init : Flambda_primitive.Init_or_assign.t) =
   match kind, init with
   | Values, Assignment Heap -> nonalloc_extcall_size
-  | Values, (Assignment (Local _) | Initialization) -> 1
+  | Values, (Assignment Local | Initialization) -> 1
   | (Immediates | Naked_floats), (Assignment _ | Initialization) -> 1
 
 let string_or_bigstring_load kind width =

middle_end/flambda2/terms/flambda_primitive.ml

@@ -230,41 +230,21 @@ type string_or_bytes =
 module Init_or_assign = struct
   type t =
     | Initialization
-    | Assignment of Alloc_mode.For_allocations.t
+    | Assignment of Alloc_mode.For_assignments.t
 
   let [@ocamlformat "disable"] print ppf t =
     let fprintf = Format.fprintf in
     match t with
     | Initialization -> fprintf ppf "Init"
-    | Assignment mode -> fprintf ppf "Assign %a" Alloc_mode.For_allocations.print mode
+    | Assignment Heap -> fprintf ppf "Assign Heap"
+    | Assignment Local -> fprintf ppf "Assign Local"
 
   let compare = Stdlib.compare
 
   let to_lambda t : Lambda.initialization_or_assignment =
     match t with
     | Initialization -> Heap_initialization
-    | Assignment mode ->
-      Assignment (Alloc_mode.For_allocations.to_lambda_modify mode)
-
-  let free_names t =
-    match t with
-    | Initialization -> Name_occurrences.empty
-    | Assignment alloc_mode -> Alloc_mode.For_allocations.free_names alloc_mode
-
-  let apply_renaming t renaming =
-    match t with
-    | Initialization -> Initialization
-    | Assignment alloc_mode ->
-      let alloc_mode' =
-        Alloc_mode.For_allocations.apply_renaming alloc_mode renaming
-      in
-      if alloc_mode == alloc_mode' then t else Assignment alloc_mode'
-
-  let ids_for_export t =
-    match t with
-    | Initialization -> Ids_for_export.empty
-    | Assignment alloc_mode ->
-      Alloc_mode.For_allocations.ids_for_export alloc_mode
+    | Assignment mode -> Assignment (Alloc_mode.For_assignments.to_lambda mode)
 end
 
 type array_like_operation =
@@ -1444,37 +1424,17 @@ let ternary_classify_for_printing p =
 
 let free_names_ternary_primitive p =
   match p with
-  | Block_set (_kind, init_or_assign) ->
-    Init_or_assign.free_names init_or_assign
-  | Array_set (_kind, init_or_assign) ->
-    Init_or_assign.free_names init_or_assign
-  | Bytes_or_bigstring_set _ | Bigarray_set _ -> Name_occurrences.empty
+  | Block_set _ | Array_set _ | Bytes_or_bigstring_set _ | Bigarray_set _ ->
+    Name_occurrences.empty
 
-let apply_renaming_ternary_primitive p renaming =
+let apply_renaming_ternary_primitive p _ =
   match p with
-  | Block_set (kind, init_or_assign) ->
-    let init_or_assign' =
-      Init_or_assign.apply_renaming init_or_assign renaming
-    in
-    if init_or_assign == init_or_assign'
-    then p
-    else Block_set (kind, init_or_assign')
-  | Array_set (kind, init_or_assign) ->
-    let init_or_assign' =
-      Init_or_assign.apply_renaming init_or_assign renaming
-    in
-    if init_or_assign == init_or_assign'
-    then p
-    else Array_set (kind, init_or_assign')
-  | Bytes_or_bigstring_set _ | Bigarray_set _ -> p
+  | Block_set _ | Array_set _ | Bytes_or_bigstring_set _ | Bigarray_set _ -> p
 
 let ids_for_export_ternary_primitive p =
   match p with
-  | Block_set (_kind, init_or_assign) ->
-    Init_or_assign.ids_for_export init_or_assign
-  | Array_set (_kind, init_or_assign) ->
-    Init_or_assign.ids_for_export init_or_assign
-  | Bytes_or_bigstring_set _ | Bigarray_set _ -> Ids_for_export.empty
+  | Block_set _ | Array_set _ | Bytes_or_bigstring_set _ | Bigarray_set _ ->
+    Ids_for_export.empty
 
 type variadic_primitive =
   | Make_block of Block_kind.t * Mutability.t * Alloc_mode.For_allocations.t

middle_end/flambda2/terms/flambda_primitive.mli

@@ -119,7 +119,7 @@ type string_or_bytes =
 module Init_or_assign : sig
   type t =
     | Initialization
-    | Assignment of Alloc_mode.For_allocations.t
+    | Assignment of Alloc_mode.For_assignments.t
 
   val to_lambda : t -> Lambda.initialization_or_assignment
 end