Add kind to Function_slot.t and Value_slot.t + fix join (ocaml-flambd…

…a#1214)

middle_end/flambda2/compare/compare.ml

@@ -270,8 +270,8 @@ let subst_set_of_closures env set =
   let value_slots =
     Set_of_closures.value_slots set
     |> Value_slot.Map.bindings
-    |> List.map (fun (var, (simple, kind)) ->
-           subst_value_slot env var, (subst_simple env simple, kind))
+    |> List.map (fun (var, simple) ->
+           subst_value_slot env var, subst_simple env simple)
     |> Value_slot.Map.of_list
   in
   let alloc = Set_of_closures.alloc_mode set in
@@ -756,7 +756,8 @@ let sets_of_closures env set1 set2 : Set_of_closures.t Comparison.t =
    * similar (and less worrisome) with function slots. *)
   let value_slots_by_value set =
     Value_slot.Map.bindings (Set_of_closures.value_slots set)
-    |> List.map (fun (var, (value, kind)) -> kind, subst_simple env value, var)
+    |> List.map (fun (var, value) ->
+           Value_slot.kind var, subst_simple env value, var)
   in
   (* We want to process the whole map to find new correspondences between
    * value slots, so we need to remember whether we've found any mismatches *)

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -1176,9 +1176,7 @@ 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_slot.t * Flambda_kind.With_subkind.t) Variable.Map.t),
-        vars_for_idents ) =
+  let (value_slots_to_bind : Value_slot.t Variable.Map.t), vars_for_idents =
     Ident.Map.fold
       (fun id value_slot (value_slots_to_bind, vars_for_idents) ->
         let var = Variable.create_with_same_name_as_ident id in
@@ -1322,8 +1320,9 @@ let close_one_function acc ~code_id ~external_env ~by_function_slot decl
   in
   let acc, body =
     Variable.Map.fold
-      (fun var (value_slot, kind) (acc, body) ->
+      (fun var value_slot (acc, body) ->
         let var = VB.create var Name_mode.normal in
+        let kind = Value_slot.kind value_slot in
         let named =
           Named.create_prim
             (Unary
@@ -1464,7 +1463,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, kind) map)
+          Ident.Map.add id (Value_slot.create compilation_unit ~name kind) map)
       (Function_decls.all_free_idents function_declarations)
       Ident.Map.empty
   in
@@ -1616,7 +1615,8 @@ 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, kind) map ->
+      (fun id value_slot map ->
+        let kind = Value_slot.kind value_slot in
         let external_simple, kind' =
           find_simple_from_id_with_kind external_env id
         in
@@ -1628,7 +1628,7 @@ let close_functions acc external_env ~current_region function_declarations =
         (* We're sure [external_simple] is a variable since
            [value_slot_from_idents] has already filtered constants and symbols
            out. *)
-        Value_slot.Map.add value_slot (external_simple, kind) map)
+        Value_slot.Map.add value_slot external_simple map)
       value_slots_from_idents Value_slot.Map.empty
   in
   let set_of_closures =
@@ -1779,7 +1779,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)
+      ~name:(Ident.name wrapper_id) K.With_subkind.any_value
   in
   let num_provided = List.length provided in
   let params =

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -1644,7 +1644,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)
+      ~name:(Ident.name fid) Flambda_kind.With_subkind.any_value
   in
   let body acc ccenv =
     let ccenv = CCenv.set_path_to_root ccenv loc in

middle_end/flambda2/identifiers/slot.ml

@@ -19,7 +19,8 @@ module type S = sig
 
   module Lmap : Lmap.S with type key = t
 
-  val create : Compilation_unit.t -> name:string -> t
+  val create :
+    Compilation_unit.t -> name:string -> Flambda_kind.With_subkind.t -> t
 
   val get_compilation_unit : t -> Compilation_unit.t
 
@@ -31,6 +32,8 @@ module type S = sig
 
   val name : t -> string
 
+  val kind : t -> Flambda_kind.With_subkind.t
+
   val rename : t -> t
 end
 
@@ -40,36 +43,56 @@ end) : S = struct
   type t =
     { compilation_unit : Compilation_unit.t;
       name : string;
-      name_stamp : int
+      name_stamp : int;
           (** [name_stamp]s are unique within any given compilation unit. *)
+      kind : Flambda_kind.With_subkind.t
     }
 
   module Self = Container_types.Make (struct
     type nonrec t = t
 
-    let compare t1 t2 =
+    let compare
+        ({ compilation_unit = compilation_unit1;
+           name = _;
+           name_stamp = name_stamp1;
+           kind = kind1
+         } as t1)
+        ({ compilation_unit = compilation_unit2;
+           name = _;
+           name_stamp = name_stamp2;
+           kind = kind2
+         } as t2) =
       if t1 == t2
       then 0
       else
-        let c = t1.name_stamp - t2.name_stamp in
+        let c = name_stamp1 - name_stamp2 in
         if c <> 0
         then c
-        else Compilation_unit.compare t1.compilation_unit t2.compilation_unit
+        else
+          let c =
+            Compilation_unit.compare compilation_unit1 compilation_unit2
+          in
+          if c <> 0 then c else Flambda_kind.With_subkind.compare kind1 kind2
 
     let equal t1 t2 = compare t1 t2 = 0
 
     let hash t =
-      Hashtbl.hash (t.name_stamp, Compilation_unit.hash t.compilation_unit)
+      Hashtbl.hash
+        ( t.name_stamp,
+          Compilation_unit.hash t.compilation_unit,
+          Flambda_kind.With_subkind.hash t.kind )
 
     let print ppf t =
-      Format.fprintf ppf "%t" P.colour;
+      Format.fprintf ppf "@[%t(" P.colour;
       if Compilation_unit.equal t.compilation_unit
            (Compilation_unit.get_current_exn ())
       then Format.fprintf ppf "%s/%d" t.name t.name_stamp
       else
         Format.fprintf ppf "%a.%s/%d" Compilation_unit.print t.compilation_unit
           t.name t.name_stamp;
-      Format.fprintf ppf "%t" Flambda_colours.pop
+      Format.fprintf ppf " @<1>\u{2237} %a" Flambda_kind.With_subkind.print
+        t.kind;
+      Format.fprintf ppf ")%t@]" Flambda_colours.pop
   end)
 
   include Self
@@ -87,8 +110,8 @@ end) : S = struct
     incr next_stamp;
     stamp
 
-  let create compilation_unit ~name =
-    { compilation_unit; name; name_stamp = get_next_stamp () }
+  let create compilation_unit ~name kind =
+    { compilation_unit; name; name_stamp = get_next_stamp (); kind }
 
   let get_compilation_unit t = t.compilation_unit
 
@@ -101,5 +124,7 @@ end) : S = struct
 
   let name t = t.name
 
+  let kind t = t.kind
+
   let rename t = { t with name_stamp = get_next_stamp () }
 end

middle_end/flambda2/identifiers/slot.mli

@@ -19,7 +19,8 @@ module type S = sig
 
   module Lmap : Lmap.S with type key = t
 
-  val create : Compilation_unit.t -> name:string -> t
+  val create :
+    Compilation_unit.t -> name:string -> Flambda_kind.With_subkind.t -> t
 
   val get_compilation_unit : t -> Compilation_unit.t
 
@@ -31,6 +32,8 @@ module type S = sig
 
   val name : t -> string
 
+  val kind : t -> Flambda_kind.With_subkind.t
+
   val rename : t -> t
 end
 

middle_end/flambda2/parser/fexpr_to_flambda.ml

@@ -133,7 +133,11 @@ let fresh_code_id env { Fexpr.txt = name; loc = _ } =
   c, { env with code_ids = DM.add name c env.code_ids }
 
 let fresh_function_slot env { Fexpr.txt = name; loc = _ } =
-  let c = Function_slot.create (Compilation_unit.get_current_exn ()) ~name in
+  let c =
+    Function_slot.create
+      (Compilation_unit.get_current_exn ())
+      ~name Flambda_kind.With_subkind.any_value
+  in
   UT.add env.function_slots name c;
   c
 
@@ -142,14 +146,14 @@ let fresh_or_existing_function_slot env ({ Fexpr.txt = name; loc = _ } as id) =
   | None -> fresh_function_slot env id
   | Some function_slot -> function_slot
 
-let fresh_value_slot env { Fexpr.txt = name; loc = _ } =
-  let c = Value_slot.create (Compilation_unit.get_current_exn ()) ~name in
+let fresh_value_slot env { Fexpr.txt = name; loc = _ } kind =
+  let c = Value_slot.create (Compilation_unit.get_current_exn ()) ~name kind in
   WT.add env.vars_within_closures name c;
   c
 
-let fresh_or_existing_value_slot env ({ Fexpr.txt = name; _ } as id) =
+let fresh_or_existing_value_slot env ({ Fexpr.txt = name; _ } as id) kind =
   match WT.find_opt env.vars_within_closures name with
-  | None -> fresh_value_slot env id
+  | None -> fresh_value_slot env id kind
   | Some value_slot -> value_slot
 
 let print_scoped_location ppf loc =
@@ -385,10 +389,11 @@ let unop env (unop : Fexpr.unop) : Flambda_primitive.unary_primitive =
   | Opaque_identity ->
     Opaque_identity { middle_end_only = false; kind = Flambda_kind.value }
   | Project_value_slot { project_from; value_slot } ->
-    let value_slot = fresh_or_existing_value_slot env value_slot in
+    (* CR mshinwell: support non-value kinds *)
+    let kind = Flambda_kind.With_subkind.any_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; kind = Flambda_kind.With_subkind.any_value }
+    Project_value_slot { project_from; value_slot; kind }
   | Project_function_slot { move_from; move_to } ->
     let move_from = fresh_or_existing_function_slot env move_from in
     let move_to = fresh_or_existing_function_slot env move_to in
@@ -490,10 +495,11 @@ let set_of_closures env fun_decls value_slots alloc =
     |> Function_slot.Lmap.of_list |> Function_declarations.create
   in
   let value_slots = Option.value value_slots ~default:[] in
-  let value_slots : (Simple.t * Flambda_kind.With_subkind.t) Value_slot.Map.t =
+  let value_slots : Simple.t Value_slot.Map.t =
     let convert ({ var; value } : Fexpr.one_value_slot) =
-      ( fresh_or_existing_value_slot env var,
-        (simple env value, Flambda_kind.With_subkind.any_value) )
+      (* CR mshinwell: support non-value kinds *)
+      ( fresh_or_existing_value_slot env var Flambda_kind.With_subkind.any_value,
+        simple env value )
     in
     List.map convert value_slots |> Value_slot.Map.of_list
   in

middle_end/flambda2/parser/flambda_to_fexpr.ml

@@ -601,7 +601,8 @@ let prim env (p : Flambda_primitive.t) : Fexpr.prim =
 
 let value_slots env map =
   List.map
-    (fun (var, (value, kind)) ->
+    (fun (var, value) ->
+      let kind = Value_slot.kind var in
       if not
            (Flambda_kind.equal
               (Flambda_kind.With_subkind.kind kind)

middle_end/flambda2/simplify/simplify_apply_expr.ml

@@ -398,6 +398,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
   in
   let new_closure_alloc_mode, num_trailing_local_params =
     (* If the closure has a local suffix, and we've supplied enough args to hit
@@ -473,12 +474,13 @@ let simplify_direct_partial_application ~simplify_expr dacc apply
             { var : Variable.t;
               (* name to bind to projected variable *)
               value : Simple.t;
-              kind : K.With_subkind.t;
-              (* value to store in closure, with kind *)
+              (* value to store in closure *)
               value_slot : Value_slot.t
             }
     end in
-    let mk_value_slot () = Value_slot.create compilation_unit ~name:"arg" in
+    let mk_value_slot kind =
+      Value_slot.create compilation_unit ~name:"arg" kind
+    in
     let applied_value (value, kind) =
       Simple.pattern_match' value
         ~const:(fun const -> Const const)
@@ -492,14 +494,9 @@ let simplify_direct_partial_application ~simplify_expr dacc apply
               Misc.fatal_errorf
                 "Simple %a which is a symbol should be of kind Value"
                 Simple.print value;
-            In_closure
-              { var;
-                value;
-                kind = K.With_subkind.any_value;
-                value_slot = mk_value_slot ()
-              })
+            In_closure { var; value; value_slot = mk_value_slot kind })
         ~var:(fun var ~coercion:_ ->
-          In_closure { var; value; kind; value_slot = mk_value_slot () })
+          In_closure { var; value; value_slot = mk_value_slot kind })
     in
     let applied_callee =
       applied_value (Apply.callee apply, K.With_subkind.any_value)
@@ -539,8 +536,9 @@ let simplify_direct_partial_application ~simplify_expr dacc apply
         (fun (expr, cost_metrics, free_names) applied_value ->
           match applied_value with
           | Const _ | Symbol _ -> expr, cost_metrics, free_names
-          | In_closure { var; value_slot; value = _; kind } ->
+          | In_closure { var; value_slot; value = _ } ->
             let arg = VB.create var Name_mode.normal in
+            let kind = Value_slot.kind value_slot in
             let prim =
               P.Unary
                 ( Project_value_slot
@@ -617,8 +615,7 @@ let simplify_direct_partial_application ~simplify_expr dacc apply
         (fun value ->
           match value with
           | Const _ | Symbol _ -> None
-          | In_closure { value_slot; value; kind; var = _ } ->
-            Some (value_slot, (value, kind)))
+          | In_closure { value_slot; value; var = _ } -> Some (value_slot, value))
         applied_values
       |> Value_slot.Map.of_list
     in

middle_end/flambda2/simplify/simplify_let_expr.ml

@@ -151,8 +151,7 @@ let rebuild_let simplify_named_result removed_operations ~rewrite_id
     in
     after_rebuild body uacc
 
-let record_one_value_slot_for_data_flow symbol value_slot (simple, _kind)
-    data_flow =
+let record_one_value_slot_for_data_flow symbol value_slot simple data_flow =
   Flow.Acc.record_value_slot (Name.symbol symbol) value_slot
     (Simple.free_names simple) data_flow