Remove subkinds from value slots (oxcaml#3981)

Author: Ekdohibs

middle_end/flambda2/compare/compare.ml

@@ -775,7 +775,7 @@ let sets_of_closures env set1 set2 : Set_of_closures.t Comparison.t =
   let ok = ref true in
   let () =
     let compare (kind1, value1, _var1) (kind2, value2, _var2) =
-      let c = Flambda_kind.With_subkind.compare kind1 kind2 in
+      let c = Flambda_kind.compare kind1 kind2 in
       if c = 0 then Simple.compare value1 value2 else c
     in
     iter2_merged (value_slots_by_value set1) (value_slots_by_value set2)

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -2627,7 +2627,17 @@ 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 kind) map)
+          let is_always_immediate =
+            match[@ocaml.warning "-4"]
+              Flambda_kind.With_subkind.non_null_value_subkind kind
+            with
+            | Tagged_immediate -> true
+            | _ -> false
+          in
+          Ident.Map.add id
+            (Value_slot.create compilation_unit ~name ~is_always_immediate
+               (Flambda_kind.With_subkind.kind kind))
+            map)
       (Function_decls.all_free_idents function_declarations)
       Ident.Map.empty
   in
@@ -2774,10 +2784,11 @@ let close_functions acc external_env ~current_region function_declarations =
         let external_simple, kind' =
           find_simple_from_id_with_kind external_env id
         in
-        if not (K.With_subkind.equal kind kind')
+        if not (K.equal kind (K.With_subkind.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'
+            K.print kind K.print
+            (K.With_subkind.kind 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
@@ -2946,7 +2957,7 @@ let wrap_partial_application acc env apply_continuation (apply : IR.apply)
   let function_slot =
     Function_slot.create
       (Compilation_unit.get_current_exn ())
-      ~name:(Ident.name wrapper_id) K.With_subkind.any_value
+      ~name:(Ident.name wrapper_id) ~is_always_immediate:false K.value
   in
   let num_provided = Flambda_arity.num_params provided_arity in
   let missing_arity_and_param_modes =

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -1353,7 +1353,7 @@ and cps_function env ~fid ~(recursive : Recursive.t) ?precomputed_free_idents
         Function_slot.create
           (Compilation_unit.get_current_exn ())
           ~name:(Ident.name fid ^ "_unboxed")
-          Flambda_kind.With_subkind.any_value
+          ~is_always_immediate:false Flambda_kind.value
       in
       let unboxed_return =
         if attr.unbox_return then unboxing_kind return else None
@@ -1424,7 +1424,7 @@ and cps_function env ~fid ~(recursive : Recursive.t) ?precomputed_free_idents
   let function_slot =
     Function_slot.create
       (Compilation_unit.get_current_exn ())
-      ~name:(Ident.name fid) Flambda_kind.With_subkind.any_value
+      ~name:(Ident.name fid) ~is_always_immediate:false Flambda_kind.value
   in
   let unboxed_products = ref Ident.Map.empty in
   let params =

middle_end/flambda2/identifiers/slot.ml

@@ -20,7 +20,11 @@ module type S = sig
   module Lmap : Lmap.S with type key = t
 
   val create :
-    Compilation_unit.t -> name:string -> Flambda_kind.With_subkind.t -> t
+    Compilation_unit.t ->
+    name:string ->
+    is_always_immediate:bool ->
+    Flambda_kind.t ->
+    t
 
   val get_compilation_unit : t -> Compilation_unit.t
 
@@ -32,7 +36,9 @@ module type S = sig
 
   val name : t -> string
 
-  val kind : t -> Flambda_kind.With_subkind.t
+  val kind : t -> Flambda_kind.t
+
+  val is_always_immediate : t -> bool
 
   val rename : t -> t
 end
@@ -45,7 +51,8 @@ end) : S = struct
       name : string;
       name_stamp : int;
           (** [name_stamp]s are unique within any given compilation unit. *)
-      kind : Flambda_kind.With_subkind.t
+      kind : Flambda_kind.t;
+      is_always_immediate : bool
     }
 
   module Self = Container_types.Make (struct
@@ -55,12 +62,14 @@ end) : S = struct
         ({ compilation_unit = compilation_unit1;
            name = _;
            name_stamp = name_stamp1;
-           kind = _
+           kind = _;
+           is_always_immediate = _
          } as t1)
         ({ compilation_unit = compilation_unit2;
            name = _;
            name_stamp = name_stamp2;
-           kind = _
+           kind = _;
+           is_always_immediate = _
          } as t2) =
       if t1 == t2
       then 0
@@ -83,8 +92,8 @@ end) : S = struct
       else
         Format.fprintf ppf "%a.%s/%d" Compilation_unit.print t.compilation_unit
           t.name t.name_stamp;
-      Format.fprintf ppf " @<1>\u{2237} %a" Flambda_kind.With_subkind.print
-        t.kind;
+      Format.fprintf ppf " @<1>\u{2237} %a%s" Flambda_kind.print t.kind
+        (if t.is_always_immediate then "(immediate)" else "");
       Format.fprintf ppf ")%t@]" Flambda_colours.pop
   end)
 
@@ -103,8 +112,13 @@ end) : S = struct
     incr next_stamp;
     stamp
 
-  let create compilation_unit ~name kind =
-    { compilation_unit; name; name_stamp = get_next_stamp (); kind }
+  let create compilation_unit ~name ~is_always_immediate kind =
+    { compilation_unit;
+      name;
+      name_stamp = get_next_stamp ();
+      kind;
+      is_always_immediate
+    }
 
   let get_compilation_unit t = t.compilation_unit
 
@@ -119,5 +133,7 @@ end) : S = struct
 
   let kind t = t.kind
 
+  let is_always_immediate t = t.is_always_immediate
+
   let rename t = { t with name_stamp = get_next_stamp () }
 end

middle_end/flambda2/identifiers/slot.mli

@@ -20,7 +20,11 @@ module type S = sig
   module Lmap : Lmap.S with type key = t
 
   val create :
-    Compilation_unit.t -> name:string -> Flambda_kind.With_subkind.t -> t
+    Compilation_unit.t ->
+    name:string ->
+    is_always_immediate:bool ->
+    Flambda_kind.t ->
+    t
 
   val get_compilation_unit : t -> Compilation_unit.t
 
@@ -32,7 +36,9 @@ module type S = sig
 
   val name : t -> string
 
-  val kind : t -> Flambda_kind.With_subkind.t
+  val kind : t -> Flambda_kind.t
+
+  val is_always_immediate : t -> bool
 
   val rename : t -> t
 end

middle_end/flambda2/parser/fexpr_to_flambda.ml

@@ -144,7 +144,7 @@ let fresh_function_slot env { Fexpr.txt = name; loc = _ } =
   let c =
     Function_slot.create
       (Compilation_unit.get_current_exn ())
-      ~name Flambda_kind.With_subkind.any_value
+      ~name ~is_always_immediate:false Flambda_kind.value
   in
   UT.add env.function_slots name c;
   c
@@ -155,7 +155,11 @@ let fresh_or_existing_function_slot env ({ Fexpr.txt = name; loc = _ } as id) =
   | Some function_slot -> function_slot
 
 let fresh_value_slot env { Fexpr.txt = name; loc = _ } kind =
-  let c = Value_slot.create (Compilation_unit.get_current_exn ()) ~name kind in
+  let c =
+    Value_slot.create
+      (Compilation_unit.get_current_exn ())
+      ~name ~is_always_immediate:false kind
+  in
   WT.add env.vars_within_closures name c;
   c
 
@@ -436,7 +440,7 @@ let unop env (unop : Fexpr.unop) : Flambda_primitive.unary_primitive =
     Opaque_identity { middle_end_only = false; kind = Flambda_kind.value }
   | Project_value_slot { project_from; value_slot } ->
     (* CR mshinwell: support non-value kinds *)
-    let kind = Flambda_kind.With_subkind.any_value in
+    let kind = Flambda_kind.value in
     let value_slot = fresh_or_existing_value_slot env value_slot kind in
     let project_from = fresh_or_existing_function_slot env project_from in
     Project_value_slot { project_from; value_slot }
@@ -562,8 +566,7 @@ let set_of_closures env fun_decls value_slots alloc =
   let value_slots : Simple.t Value_slot.Map.t =
     let convert ({ var; value } : Fexpr.one_value_slot) =
       (* CR mshinwell: support non-value kinds *)
-      ( fresh_or_existing_value_slot env var Flambda_kind.With_subkind.any_value,
-        simple env value )
+      fresh_or_existing_value_slot env var Flambda_kind.value, simple env value
     in
     List.map convert value_slots |> Value_slot.Map.of_list
   in

middle_end/flambda2/parser/flambda_to_fexpr.ml

@@ -676,10 +676,7 @@ let value_slots env map =
   List.map
     (fun (var, value) ->
       let kind = Value_slot.kind var in
-      if not
-           (Flambda_kind.equal
-              (Flambda_kind.With_subkind.kind kind)
-              Flambda_kind.value)
+      if not (Flambda_kind.equal kind Flambda_kind.value)
       then
         Misc.fatal_errorf "Value slot %a not of kind Value" Simple.print value;
       let var = Env.translate_value_slot env var in

middle_end/flambda2/simplify/simplify_apply_expr.ml

@@ -425,7 +425,7 @@ let simplify_direct_partial_application ~simplify_expr dacc apply
   let compilation_unit = Compilation_unit.get_current_exn () in
   let wrapper_function_slot =
     Function_slot.create compilation_unit ~name:"partial_app_closure"
-      K.With_subkind.any_value
+      ~is_always_immediate:false K.value
   in
   (* The allocation mode of the closure is directly determined by the alloc_mode
      of the application. We check here that it is consistent with
@@ -507,7 +507,15 @@ let simplify_direct_partial_application ~simplify_expr dacc apply
                 }
         end in
         let mk_value_slot kind =
-          Value_slot.create compilation_unit ~name:"arg" kind
+          let is_always_immediate =
+            match[@ocaml.warning "-4"]
+              K.With_subkind.non_null_value_subkind kind
+            with
+            | Tagged_immediate -> true
+            | _ -> false
+          in
+          Value_slot.create compilation_unit ~name:"arg" ~is_always_immediate
+            (K.With_subkind.kind kind)
         in
         let applied_value (value, kind) =
           Simple.pattern_match' value

middle_end/flambda2/simplify/simplify_set_of_closures.ml

@@ -681,7 +681,7 @@ let simplify_and_lift_set_of_closures dacc ~closure_bound_vars_inverse
   let value_slot_types =
     Value_slot.Map.mapi
       (fun value_slot in_slot ->
-        let kind = K.With_subkind.kind (Value_slot.kind value_slot) in
+        let kind = Value_slot.kind value_slot in
         Simple.pattern_match in_slot
           ~const:(fun _ -> T.alias_type_of kind in_slot)
           ~name:(fun name ~coercion ->

middle_end/flambda2/simplify/simplify_unary_primitive.ml

@@ -69,8 +69,7 @@ let simplify_project_value_slot function_slot value_slot ~min_name_mode dacc
           simple
       in
       let dacc =
-        DA.add_variable dacc result_var
-          (T.alias_type_of (K.With_subkind.kind kind) simple)
+        DA.add_variable dacc result_var (T.alias_type_of kind simple)
       in
       SPR.create (Named.create_simple simple) ~try_reify:true dacc
     | Need_meet ->
@@ -81,15 +80,16 @@ let simplify_project_value_slot function_slot value_slot ~min_name_mode dacc
             (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_kind:kind)
-          ~result_var ~result_kind:(K.With_subkind.kind kind)
+          ~result_var ~result_kind:kind
       in
       let dacc = DA.add_use_of_value_slot result.dacc value_slot in
       SPR.with_dacc result dacc
   in
   let dacc =
     Simplify_common.add_symbol_projection result.dacc ~projected_from:closure
       (Symbol_projection.Projection.project_value_slot function_slot value_slot)
-      ~projection_bound_to:result_var ~kind
+      ~projection_bound_to:result_var
+      ~kind:(Flambda_kind.With_subkind.anything kind)
   in
   SPR.with_dacc result dacc