Apply coercions to both closure types and function declarations (#473)

When applying a coercion to the type of a set of closures, we need to
update both the function declarations *and* the types of the closures.

middle_end/flambda/types/structures/closures_entry.rec.ml

@@ -202,3 +202,15 @@ let map_function_decl_types
       }
     in
     Ok t
+
+let map_closure_types
+      { function_decls; closure_types; closure_var_types; }
+      ~(f : Type_grammar.t -> Type_grammar.t Or_bottom.t) : _ Or_bottom.t =
+  let closure_types =
+    Product.Closure_id_indexed.map_types closure_types ~f
+  in
+  Or_bottom.map closure_types ~f:(fun closure_types ->
+      { function_decls;
+        closure_types;
+        closure_var_types;
+      })

middle_end/flambda/types/structures/closures_entry.rec.mli

@@ -37,6 +37,11 @@ val find_function_declaration
 
 val closure_types : t -> Type_grammar.t Closure_id.Map.t
 
+val map_closure_types
+   : t
+  -> f:(Type_grammar.t -> Type_grammar.t Or_bottom.t)
+  -> t Or_bottom.t
+
 val function_decl_types : t -> Function_declaration_type.t Closure_id.Map.t
 
 val closure_var_types : t -> Type_grammar.t Var_within_closure.Map.t

middle_end/flambda/types/structures/row_like.rec.ml

@@ -658,6 +658,10 @@ struct
       map_maps_to t ~f:(fun closures_entry ->
           Closures_entry.map_function_decl_types closures_entry ~f)
 
+    let map_closure_types t ~f =
+      map_maps_to t ~f:(fun closures_entry ->
+          Closures_entry.map_closure_types closures_entry ~f)
+
     let create_exactly
         (closure_id : Closure_id.t)
         (contents : Set_of_closures_contents.t)

middle_end/flambda/types/structures/row_like.rec.mli

@@ -119,6 +119,11 @@ module For_closures_entry_by_set_of_closures_contents : sig
     -> f:(Function_declaration_type.t -> Function_declaration_type.t Or_bottom.t)
     -> t Or_bottom.t
 
+  val map_closure_types
+     : t
+    -> f:(Type_grammar.t -> Type_grammar.t Or_bottom.t)
+    -> t Or_bottom.t
+
   include Type_structure_intf.S
     with type t := t
     with type flambda_type := Type_grammar.t

middle_end/flambda/types/type_of_kind/type_of_kind_value0.rec.ml

@@ -159,13 +159,25 @@ let all_ids_for_export t =
 let apply_coercion t coercion : _ Or_bottom.t =
   match t with
   | Closures { by_closure_id; } ->
-    Or_bottom.map
-      (Row_like.For_closures_entry_by_set_of_closures_contents.
-       map_function_decl_types
-        by_closure_id
-        ~f:(fun (decl : Function_declaration_type.t) : _ Or_bottom.t ->
-          Function_declaration_type.apply_coercion decl coercion))
-      ~f:(fun by_closure_id -> Closures { by_closure_id; })
+    begin match
+      Row_like.For_closures_entry_by_set_of_closures_contents.
+        map_function_decl_types
+         by_closure_id
+         ~f:(fun (decl : Function_declaration_type.t) : _ Or_bottom.t ->
+           Function_declaration_type.apply_coercion decl coercion)
+    with
+    | Bottom -> Bottom
+    | Ok by_closure_id ->
+      match
+        Row_like.For_closures_entry_by_set_of_closures_contents.
+          map_closure_types
+          by_closure_id
+          ~f:(fun ty -> Type_grammar.apply_coercion ty coercion)
+      with
+      | Bottom -> Bottom
+      | Ok by_closure_id ->
+        Ok (Closures { by_closure_id; })
+    end
   | Variant _
   | Boxed_float _
   | Boxed_int32 _