Unboxed types fixes extracted from PR1166 (ocaml-flambda#1202)

backend/cmm_helpers.ml

@@ -735,11 +735,14 @@ let rec remove_unit = function
 let field_address ptr n dbg =
   if n = 0 then ptr else Cop (Cadda, [ptr; Cconst_int (n * size_addr, dbg)], dbg)
 
+let get_field_gen_given_memory_chunk memory_chunk mut ptr n dbg =
+  Cop (Cload (memory_chunk, mut), [field_address ptr n dbg], dbg)
+
 let get_field_gen mut ptr n dbg =
-  Cop (Cload (Word_val, mut), [field_address ptr n dbg], dbg)
+  get_field_gen_given_memory_chunk Word_val mut ptr n dbg
 
 let get_field_codepointer mut ptr n dbg =
-  Cop (Cload (Word_int, mut), [field_address ptr n dbg], dbg)
+  get_field_gen_given_memory_chunk Word_int mut ptr n dbg
 
 let set_field ptr n newval init dbg =
   Cop (Cstore (Word_val, init), [field_address ptr n dbg; newval], dbg)

backend/cmm_helpers.mli

@@ -230,10 +230,20 @@ val remove_unit : expression -> expression
 val field_address : expression -> int -> Debuginfo.t -> expression
 
 (** [get_field_gen mut ptr n dbg] returns an expression for the access to the
-    [n]th field of the block pointed to by [ptr] *)
+    [n]th field of the block pointed to by [ptr].  The [memory_chunk] used is
+    always [Word_val]. *)
 val get_field_gen :
   Asttypes.mutable_flag -> expression -> int -> Debuginfo.t -> expression
 
+(** Like [get_field_gen] but allows use of a different [memory_chunk]. *)
+val get_field_gen_given_memory_chunk :
+  Cmm.memory_chunk ->
+  Asttypes.mutable_flag ->
+  expression ->
+  int ->
+  Debuginfo.t ->
+  expression
+
 (** Get the field of the given [block] whose index is specified by the Cmm
     expresson [index] (in words). *)
 val get_field_computed :

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -1176,12 +1176,14 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
 
      Note that free variables corresponding to predefined exception identifiers
      have been filtered out by [close_functions], above. *)
-  let value_slots_to_bind, value_slots_for_idents =
+  let ( (value_slots_to_bind :
+          (Value_slot.t * Flambda_kind.With_subkind.t) Variable.Map.t),
+        vars_for_idents ) =
     Ident.Map.fold
-      (fun id value_slots_for_idents (to_bind, var_for_ident) ->
+      (fun id value_slot (value_slots_to_bind, vars_for_idents) ->
         let var = Variable.create_with_same_name_as_ident id in
-        ( Variable.Map.add var value_slots_for_idents to_bind,
-          Ident.Map.add id var var_for_ident ))
+        ( Variable.Map.add var value_slot value_slots_to_bind,
+          Ident.Map.add id var vars_for_idents ))
       value_slots_from_idents
       (Variable.Map.empty, Ident.Map.empty)
   in
@@ -1245,7 +1247,7 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
           (Env.add_var env id var kind)
           var
           (find_value_approximation acc env simple))
-      value_slots_for_idents closure_env
+      vars_for_idents closure_env
   in
   let closure_env =
     List.fold_right
@@ -1320,16 +1322,13 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
   in
   let acc, body =
     Variable.Map.fold
-      (fun var value_slot (acc, body) ->
+      (fun var (value_slot, kind) (acc, body) ->
         let var = VB.create var Name_mode.normal in
         let named =
           Named.create_prim
             (Unary
                ( Project_value_slot
-                   { project_from = function_slot;
-                     value_slot;
-                     kind = K.With_subkind.any_value
-                   },
+                   { project_from = function_slot; value_slot; kind },
                  my_closure' ))
             Debuginfo.none
         in
@@ -1444,16 +1443,18 @@ let close_functions acc external_env ~current_region function_declarations =
         (* Filter out predefined exception identifiers and simple substitutions.
            The former will be turned into symbols, and the latter substituted
            when we closure-convert the body *)
-        let has_non_var_subst, subst_var =
+        let has_non_var_subst, subst_var, kind =
           match Env.find_simple_to_substitute_exn external_env id with
-          | exception Not_found -> false, None
-          | simple, _kind ->
+          | exception Not_found ->
+            let _, kind = find_simple_from_id_with_kind external_env id in
+            false, None, kind
+          | simple, kind ->
             Simple.pattern_match simple
-              ~const:(fun _ -> true, None)
+              ~const:(fun _ -> true, None, kind)
               ~name:(fun name ~coercion:_ ->
                 Name.pattern_match name
-                  ~var:(fun var -> false, Some var)
-                  ~symbol:(fun _ -> true, None))
+                  ~var:(fun var -> false, Some var, kind)
+                  ~symbol:(fun _ -> true, None, kind))
         in
         if has_non_var_subst || Ident.is_predef id
         then map
@@ -1463,7 +1464,7 @@ let close_functions acc external_env ~current_region function_declarations =
             | None -> Ident.name id
             | Some var -> Variable.name var
           in
-          Ident.Map.add id (Value_slot.create compilation_unit ~name) map)
+          Ident.Map.add id (Value_slot.create compilation_unit ~name, kind) map)
       (Function_decls.all_free_idents function_declarations)
       Ident.Map.empty
   in
@@ -1615,10 +1616,15 @@ let close_functions acc external_env ~current_region function_declarations =
   let function_decls = Function_declarations.create funs in
   let value_slots =
     Ident.Map.fold
-      (fun id value_slot map ->
-        let external_simple, kind =
+      (fun id (value_slot, kind) map ->
+        let external_simple, kind' =
           find_simple_from_id_with_kind external_env id
         in
+        if not (K.With_subkind.equal kind kind')
+        then
+          Misc.fatal_errorf "Value slot kinds %a and %a don't match for slot %a"
+            K.With_subkind.print kind K.With_subkind.print kind'
+            Value_slot.print value_slot;
         (* We're sure [external_simple] is a variable since
            [value_slot_from_idents] has already filtered constants and symbols
            out. *)

middle_end/flambda2/simplify/simplify_apply_expr.ml

@@ -709,18 +709,6 @@ let simplify_direct_function_call ~simplify_expr dacc apply
       "[Apply] terms with a [probe_name] (i.e. that call a tracing probe) must \
        always be marked as [Never_inline]:@ %a"
       Apply.print apply);
-  let result_arity_of_application = Apply.return_arity apply in
-  if not
-       (Flambda_arity.equal
-          (Flambda_arity.With_subkinds.to_arity result_arity)
-          (Flambda_arity.With_subkinds.to_arity result_arity_of_application))
-  then
-    Misc.fatal_errorf
-      "Wrong return arity for direct OCaml function call (expected %a, found \
-       %a):@ %a"
-      Flambda_arity.With_subkinds.print result_arity
-      Flambda_arity.With_subkinds.print result_arity_of_application Apply.print
-      apply;
   let coming_from_indirect = Option.is_none callee's_code_id_from_call_kind in
   let callee's_code_id : _ Or_bottom.t =
     match callee's_code_id_from_call_kind with
@@ -765,24 +753,66 @@ let simplify_direct_function_call ~simplify_expr dacc apply
       let args = Apply.args apply in
       let provided_num_args = List.length args in
       let num_params = Flambda_arity.With_subkinds.cardinal params_arity in
+      let result_arity_of_application = Apply.return_arity apply in
       if provided_num_args = num_params
-      then
+      then (
+        (* This check can only be performed for exact applications:
+
+           - In the partial application case, the type checker should have
+           specified kind Value as the return kind of the application
+           (propagated through Lambda to this point), and it would be wrong to
+           compare against the return arity of the fully-applied function.
+
+           - In the overapplication case, the correct return arity is only
+           present on the application expression, so all we can do is check that
+           the function being overapplied returns kind Value. *)
+        if not
+             (Flambda_arity.equal
+                (Flambda_arity.With_subkinds.to_arity result_arity)
+                (Flambda_arity.With_subkinds.to_arity
+                   result_arity_of_application))
+        then
+          Misc.fatal_errorf
+            "Wrong return arity for direct OCaml function call\n\
+            \     (expected %a, found  %a):@ %a"
+            Flambda_arity.With_subkinds.print result_arity
+            Flambda_arity.With_subkinds.print result_arity_of_application
+            Apply.print apply;
         simplify_direct_full_application ~simplify_expr dacc apply function_decl
           ~params_arity ~result_arity ~result_types ~down_to_up
-          ~coming_from_indirect ~callee's_code_metadata
+          ~coming_from_indirect ~callee's_code_metadata)
       else if provided_num_args > num_params
-      then
+      then (
+        (* See comment above. *)
+        if not
+             (Flambda_arity.is_singleton_value
+                (Flambda_arity.With_subkinds.to_arity result_arity))
+        then
+          Misc.fatal_errorf
+            "Non-singleton-value return arity for overapplied OCaml function:@ \
+             %a"
+            Apply.print apply;
         simplify_direct_over_application ~simplify_expr dacc apply ~down_to_up
           ~coming_from_indirect ~apply_alloc_mode ~current_region
-          ~callee's_code_id ~callee's_code_metadata
+          ~callee's_code_id ~callee's_code_metadata)
       else if provided_num_args > 0 && provided_num_args < num_params
-      then
+      then (
+        (* See comment above. *)
+        if not
+             (Flambda_arity.is_singleton_value
+                (Flambda_arity.With_subkinds.to_arity
+                   result_arity_of_application))
+        then
+          Misc.fatal_errorf
+            "Non-singleton-value return arity for partially-applied OCaml \
+             function:@ %a"
+            Apply.print apply;
         simplify_direct_partial_application ~simplify_expr dacc apply
           ~callee's_code_id ~callee's_code_metadata ~callee's_function_slot
           ~param_arity:params_arity ~result_arity ~recursive ~down_to_up
           ~coming_from_indirect ~closure_alloc_mode_from_type ~current_region
           ~num_trailing_local_params:
-            (Code_metadata.num_trailing_local_params callee's_code_metadata)
+            (Code_metadata.num_trailing_local_params callee's_code_metadata))
       else
         Misc.fatal_errorf
           "Function with %d params when simplifying direct OCaml function call \

middle_end/flambda2/simplify/simplify_set_of_closures.ml

@@ -795,7 +795,7 @@ let type_value_slots_and_make_lifting_decision_for_one_set dacc
      the case where we are considering lifting a set that has not been lifted
      before, there are never any other mutually-recursive sets ([Named.t] does
      not allow them). *)
-  let[@inline] variable_permits_lifting var =
+  let[@inline] variable_permits_lifting var kind =
     (* Variables (excluding ones bound to symbol projections; see below) in the
        definition of the set of closures will currently prevent lifting if the
        allocation mode is [Local] and we cannot show that such variables never
@@ -817,18 +817,18 @@ let type_value_slots_and_make_lifting_decision_for_one_set dacc
        | Local _ -> (
          match
            T.never_holds_locally_allocated_values (DA.typing_env dacc) var
-             K.value
+             (K.With_subkind.kind kind)
          with
          | Proved () -> true
          | Unknown -> false)
        | Heap -> true
   in
-  let value_slot_permits_lifting _value_slot (simple, _kind) =
+  let value_slot_permits_lifting _value_slot (simple, kind) =
     can_lift_coercion (Simple.coercion simple)
     && Simple.pattern_match' simple
          ~const:(fun _ -> true)
          ~symbol:(fun _ ~coercion:_ -> true)
-         ~var:(fun var ~coercion:_ -> variable_permits_lifting var)
+         ~var:(fun var ~coercion:_ -> variable_permits_lifting var kind)
   in
   let can_lift =
     Name_mode.is_normal name_mode_of_bound_vars

middle_end/flambda2/simplify/simplify_unary_primitive.ml

@@ -65,7 +65,8 @@ let simplify_project_value_slot function_slot value_slot kind ~min_name_mode
         else T.get_alias_exn (S.simplify_simple dacc simple ~min_name_mode)
       in
       let dacc =
-        DA.add_variable dacc result_var (T.alias_type_of K.value simple)
+        DA.add_variable dacc result_var
+          (T.alias_type_of (K.With_subkind.kind kind) simple)
       in
       SPR.create (Named.create_simple simple) ~try_reify:true dacc
     | Need_meet ->
@@ -75,7 +76,7 @@ let simplify_project_value_slot function_slot value_slot kind ~min_name_mode
           ~shape:
             (T.closure_with_at_least_this_value_slot
                ~this_function_slot:function_slot value_slot
-               ~value_slot_var:(Bound_var.var result_var))
+               ~value_slot_var:(Bound_var.var result_var) ~value_slot_kind:kind)
           ~result_var ~result_kind:(K.With_subkind.kind kind)
       in
       let dacc = DA.add_use_of_value_slot result.dacc value_slot in

middle_end/flambda2/simplify/unboxing/build_unboxing_denv.ml

@@ -61,14 +61,15 @@ let rec denv_of_decision denv ~param_var (decision : U.decision) : DE.t =
   | Unbox (Closure_single_entry { function_slot; vars_within_closure }) ->
     let denv =
       Value_slot.Map.fold
-        (fun _ ({ epa = { param = var; _ }; _ } : U.field_decision) denv ->
+        (fun _ ({ epa = { param = var; _ }; kind; _ } : U.field_decision) denv ->
           let v = VB.create var Name_mode.normal in
-          DE.define_variable denv v K.value)
+          DE.define_variable denv v (K.With_subkind.kind kind))
         vars_within_closure denv
     in
     let map =
       Value_slot.Map.map
-        (fun ({ epa = { param = var; _ }; _ } : U.field_decision) -> var)
+        (fun ({ epa = { param = var; _ }; kind; _ } : U.field_decision) ->
+          var, kind)
         vars_within_closure
     in
     let shape =

middle_end/flambda2/terms/flambda_primitive.ml

@@ -934,9 +934,8 @@ let result_kind_of_unary_primitive p : result_kind =
   | Bigarray_length _ -> Singleton K.naked_immediate
   | Unbox_number kind -> Singleton (K.Boxable_number.unboxed_kind kind)
   | Untag_immediate -> Singleton K.naked_immediate
-  | Box_number _ | Tag_immediate | Project_function_slot _
-  | Project_value_slot _ ->
-    Singleton K.value
+  | Box_number _ | Tag_immediate | Project_function_slot _ -> Singleton K.value
+  | Project_value_slot { kind; _ } -> Singleton (K.With_subkind.kind kind)
   | Is_boxed_float | Is_flat_float_array -> Singleton K.naked_immediate
   | Begin_try_region -> Singleton K.region
   | End_region -> Singleton K.value

middle_end/flambda2/to_cmm/to_cmm_primitive.ml

@@ -592,12 +592,17 @@ let unary_primitive env res dbg f arg =
       let message = dead_slots_msg dbg [c1; c2] [] in
       let expr, res = C.invalid res ~message in
       None, res, expr)
-  | Project_value_slot { project_from; value_slot; kind = _ } -> (
+  | Project_value_slot { project_from; value_slot; kind } -> (
     match
       value_slot_offset env value_slot, function_slot_offset env project_from
     with
     | Live_value_slot { offset; _ }, Live_function_slot { offset = base; _ } ->
-      None, res, C.get_field_gen Asttypes.Immutable arg (offset - base) dbg
+      let memory_chunk = To_cmm_shared.memory_chunk_of_kind kind in
+      let expr =
+        C.get_field_gen_given_memory_chunk memory_chunk Asttypes.Immutable arg
+          (offset - base) dbg
+      in
+      None, res, expr
     | Dead_value_slot, Live_function_slot _ ->
       let message = dead_slots_msg dbg [] [value_slot] in
       let expr, res = C.invalid res ~message in

middle_end/flambda2/to_cmm/to_cmm_shared.ml

@@ -48,6 +48,24 @@ let machtype_of_kind (k : Flambda_kind.t) =
     Cmm.typ_int
   | Region | Rec_info -> assert false
 
+let memory_chunk_of_kind (kind : Flambda_kind.With_subkind.t) : Cmm.memory_chunk
+    =
+  match Flambda_kind.With_subkind.kind kind with
+  | Value -> (
+    match Flambda_kind.With_subkind.subkind kind with
+    | Tagged_immediate -> Word_int
+    | Anything | Boxed_float | Boxed_int32 | Boxed_int64 | Boxed_nativeint
+    | Variant _ | Float_block _ | Float_array | Immediate_array | Value_array
+    | Generic_array ->
+      Word_val)
+  | Naked_number (Naked_int32 | Naked_int64 | Naked_nativeint | Naked_immediate)
+    ->
+    Word_int
+  | Naked_number Naked_float -> Double
+  | Region | Rec_info ->
+    Misc.fatal_errorf "Bad kind %a for [memory_chunk_of_kind]"
+      Flambda_kind.With_subkind.print kind
+
 let machtype_of_kinded_parameter p =
   Bound_parameter.kind p |> Flambda_kind.With_subkind.kind |> machtype_of_kind
 

middle_end/flambda2/to_cmm/to_cmm_shared.mli

@@ -31,6 +31,8 @@ val machtype_of_kind : Flambda_kind.t -> Cmm.machtype_component array
 val machtype_of_kinded_parameter :
   Bound_parameter.t -> Cmm.machtype_component array
 
+val memory_chunk_of_kind : Flambda_kind.With_subkind.t -> Cmm.memory_chunk
+
 (** Create a constant int expression from a targetint. *)
 val targetint : dbg:Debuginfo.t -> Targetint_32_64.t -> Cmm.expression
 

middle_end/flambda2/types/flambda2_types.mli

@@ -479,11 +479,12 @@ val closure_with_at_least_this_value_slot :
   this_function_slot:Function_slot.t ->
   Value_slot.t ->
   value_slot_var:Variable.t ->
+  value_slot_kind:Flambda_kind.With_subkind.t ->
   flambda_type
 
 val closure_with_at_least_these_value_slots :
   this_function_slot:Function_slot.t ->
-  Variable.t Value_slot.Map.t ->
+  (Variable.t * Flambda_kind.With_subkind.t) Value_slot.Map.t ->
   flambda_type
 
 val array_of_length :