Fix storing unboxed_vec128 in closures

Author: Ekdohibs

middle_end/clambda_layout.ml

@@ -17,6 +17,7 @@ type atom =
   | Value_int
   | Unboxed_float
   | Unboxed_int of Lambda.boxed_integer
+  | Unboxed_vector of Lambda.boxed_vector
 
 let rec fold_left_layout (f : 'acc -> 'e -> atom -> 'acc) (acc : 'acc)
     (expr : Clambda.ulambda) (layout : Clambda_primitives.layout) : 'acc =
@@ -28,6 +29,7 @@ let rec fold_left_layout (f : 'acc -> 'e -> atom -> 'acc) (acc : 'acc)
        closure."
   | Punboxed_float -> f acc expr Unboxed_float
   | Punboxed_int bi -> f acc expr (Unboxed_int bi)
+  | Punboxed_vector bv -> f acc expr (Unboxed_vector bv)
   | Pvalue Pintval -> f acc expr Value_int
   | Pvalue _ -> f acc expr Value
   | Punboxed_product layouts ->
@@ -59,6 +61,7 @@ let print_atom ppf = function
   | Unboxed_int Pint32 -> Format.fprintf ppf "unboxed_int32"
   | Unboxed_int Pint64 -> Format.fprintf ppf "unboxed_int64"
   | Unboxed_int Pnativeint -> Format.fprintf ppf "unboxed_nativeint"
+  | Unboxed_vector Pvec128 -> Format.fprintf ppf "unboxed_vec128"
 
 let equal_decomposition = ( = )
 
@@ -80,6 +83,7 @@ let rec decompose (layout : Lambda.layout) : _ decomposition' =
        closure."
   | Punboxed_float -> Gc_invisible ((), Unboxed_float)
   | Punboxed_int bi -> Gc_invisible ((), Unboxed_int bi)
+  | Punboxed_vector bv -> Gc_invisible ((), Unboxed_vector bv)
   | Pvalue Pintval -> Gc_invisible ((), Value_int)
   | Pvalue _ -> Gc_visible ((), Value)
   | Punboxed_product l -> Product (Array.of_list (List.map decompose l))
@@ -105,7 +109,9 @@ let rec fold_decompose (f1 : 'acc -> 'a -> atom -> 'acc * 'b)
     acc, Product elts
 
 let atom_size (layout : atom) =
-  match layout with Value | Value_int | Unboxed_float | Unboxed_int _ -> 1
+  match layout with
+  | Value | Value_int | Unboxed_float | Unboxed_int _ -> 1
+  | Unboxed_vector Pvec128 -> 2
 
 let assign_invisible_offsets init_pos (var, dec) =
   let f_visible acc () _layout = acc, () in

middle_end/clambda_layout.mli

@@ -17,6 +17,7 @@ type atom =
   | Value_int
   | Unboxed_float
   | Unboxed_int of Lambda.boxed_integer
+  | Unboxed_vector of Lambda.boxed_vector
 
 val fold_left_layout :
   ('acc -> Clambda.ulambda -> atom -> 'acc) ->

middle_end/closure/closure.ml

@@ -970,7 +970,7 @@ let split_closure_fv env kinds fv =
         let f (not_scanned_fv, scanned_fv) expr (atom : Clambda_layout.atom) =
           match atom with
           | Value -> not_scanned_fv, ((expr, atom) :: scanned_fv)
-          | Value_int | Unboxed_float | Unboxed_int _ ->
+          | Value_int | Unboxed_float | Unboxed_int _ | Unboxed_vector _ ->
             ((expr, atom) :: not_scanned_fv, scanned_fv)
         in
         Clambda_layout.fold_left_layout f acc (close_var env id) kind)
@@ -984,6 +984,7 @@ let layout_of_atom (atom : Closure_offsets.layout_atom) : Lambda.layout =
   | Value_int -> Pvalue Pintval
   | Unboxed_float -> Punboxed_float
   | Unboxed_int bi -> Punboxed_int bi
+  | Unboxed_vector bv -> Punboxed_vector bv
 
 let load_env_field ~base_offset
     ~closure (parts : Clambda_layout.decomposition) : Clambda.ulambda =

middle_end/flambda/flambda_to_clambda.ml

@@ -250,6 +250,7 @@ let layout_of_atom (atom : Closure_offsets.layout_atom) : Lambda.layout =
   | Value_int -> Pvalue Pintval
   | Unboxed_float -> Punboxed_float
   | Unboxed_int bi -> Punboxed_int bi
+  | Unboxed_vector bv -> Punboxed_vector bv
 
 let load_env_field ~fun_offset
     ~closure_using_field (parts : Closure_offsets.parts) : Clambda.ulambda =
@@ -735,7 +736,7 @@ and to_clambda_set_of_closures t env
             (expr: Clambda.ulambda) (atom : Closure_offsets.layout_atom) =
           match atom with
           | Value -> not_scanned_fv, (expr :: scanned_fv)
-          | Value_int | Unboxed_float | Unboxed_int _ ->
+          | Value_int | Unboxed_float | Unboxed_int _ | Unboxed_vector _ ->
             (expr :: not_scanned_fv, scanned_fv)
         in
         let closure, var_layout = subst_var env free_var.var in

middle_end/flambda2/simplify_shared/slot_offsets.ml

@@ -873,10 +873,12 @@ end = struct
           | Region | Rec_info ->
             Misc.fatal_errorf "Value slot %a has Region or Rec_info kind"
               Value_slot.print value_slot
-          | Naked_number _ ->
+          | Naked_number (Naked_immediate | Naked_float | Naked_int32 | Naked_int64 | Naked_nativeint) ->
             1, true
-            (* flambda only supports 64-bits for now, so naked numbers can only
-               be of size 1 *)
+          (* flambda only supports 64-bits for now, so naked numbers can only
+             be of size 1 *)
+          | Naked_number Naked_vec128 ->
+            2, true
           | Value -> (
             match[@ocaml.warning "-4"]
               Flambda_kind.With_subkind.subkind kind