diff --git a/middle_end/flambda/types/structures/closures_entry.rec.ml b/middle_end/flambda/types/structures/closures_entry.rec.ml index 709c43772c59..4a6282bb13ec 100644 --- a/middle_end/flambda/types/structures/closures_entry.rec.ml +++ b/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; + }) diff --git a/middle_end/flambda/types/structures/closures_entry.rec.mli b/middle_end/flambda/types/structures/closures_entry.rec.mli index 1404b3605528..6babf555fac5 100644 --- a/middle_end/flambda/types/structures/closures_entry.rec.mli +++ b/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 diff --git a/middle_end/flambda/types/structures/row_like.rec.ml b/middle_end/flambda/types/structures/row_like.rec.ml index 676e1e3149dd..fc7589e4d6bb 100644 --- a/middle_end/flambda/types/structures/row_like.rec.ml +++ b/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) diff --git a/middle_end/flambda/types/structures/row_like.rec.mli b/middle_end/flambda/types/structures/row_like.rec.mli index d0b87cb54ae9..5854eb8f1916 100644 --- a/middle_end/flambda/types/structures/row_like.rec.mli +++ b/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 diff --git a/middle_end/flambda/types/type_of_kind/type_of_kind_value0.rec.ml b/middle_end/flambda/types/type_of_kind/type_of_kind_value0.rec.ml index a63b6ce25f86..43bb69a97663 100644 --- a/middle_end/flambda/types/type_of_kind/type_of_kind_value0.rec.ml +++ b/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 _