Unboxed tuples fixes

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -450,8 +450,8 @@ let let_cont_nonrecursive_with_extra_params acc env ccenv ~is_exn_handler
               (Flambda_arity.unarize arity)
           in
           let handler_env =
-            Env.register_unboxed_product handler_env ~unboxed_product:id
-              ~before_unarization:arity_component ~fields
+            Env.register_unboxed_product_with_kinds handler_env
+              ~unboxed_product:id ~before_unarization:arity_component ~fields
           in
           let new_params_rev =
             List.map (fun (id, kind) -> id, IR.Not_user_visible, kind) fields
@@ -912,7 +912,7 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
               let arity = Flambda_arity.create [arity_component] in
               let fields = Flambda_arity.fresh_idents_unarized ~id arity in
               let env =
-                Env.register_unboxed_product env ~unboxed_product:id
+                Env.register_unboxed_product_with_kinds env ~unboxed_product:id
                   ~before_unarization:arity_component ~fields
               in
               env, fields
@@ -1085,7 +1085,7 @@ let rec cps acc env ccenv (lam : L.lambda) (k : cps_continuation)
                   let before_unarization =
                     Flambda_arity.Component_for_creation.from_lambda layout
                   in
-                  ( Env.register_unboxed_product handler_env
+                  ( Env.register_unboxed_product_with_kinds handler_env
                       ~unboxed_product:arg ~before_unarization ~fields,
                     fields ))
               handler_env
@@ -1612,6 +1612,17 @@ and cps_function env ~fid ~(recursive : Recursive.t) ?precomputed_free_idents
     Env.create ~current_unit:(Env.current_unit env)
       ~return_continuation:body_cont ~exn_continuation:body_exn_cont ~my_region
   in
+  let new_env, free_idents_of_body =
+    Ident.Set.fold
+      (fun id (new_env, free_idents_of_body) ->
+        match Env.get_unboxed_product_fields env id with
+        | None -> new_env, Ident.Set.add id free_idents_of_body
+        | Some (before_unarization, fields) ->
+          ( Env.register_unboxed_product new_env ~unboxed_product:id
+              ~before_unarization ~fields,
+            Ident.Set.union free_idents_of_body (Ident.Set.of_list fields) ))
+      free_idents_of_body (new_env, Ident.Set.empty)
+  in
   let exn_continuation : IR.exn_continuation =
     { exn_handler = body_exn_cont; extra_args = [] }
   in
@@ -1663,7 +1674,7 @@ and cps_function env ~fid ~(recursive : Recursive.t) ?precomputed_free_idents
     let new_env =
       Ident.Map.fold
         (fun unboxed_product (before_unarization, fields) new_env ->
-          Env.register_unboxed_product new_env ~unboxed_product
+          Env.register_unboxed_product_with_kinds new_env ~unboxed_product
             ~before_unarization ~fields)
         unboxed_products new_env
     in
@@ -1715,7 +1726,11 @@ and cps_switch acc env ccenv (switch : L.lambda_switch) ~condition_dbg
           in
           let k = restore_continuation_context_for_switch_arm env k in
           let consts_rev =
-            (arm, k, Debuginfo.none, None, IR.Var var :: extra_args)
+            ( arm,
+              k,
+              Debuginfo.none,
+              None,
+              get_unarized_vars var env @ extra_args )
             :: consts_rev
           in
           consts_rev, wrappers

middle_end/flambda2/from_lambda/lambda_to_flambda_env.ml

@@ -29,11 +29,8 @@ type t =
       (Ident.t * Flambda_kind.With_subkind.t) Ident.Map.t;
     mutables_needed_by_continuations : Ident.Set.t Continuation.Map.t;
     unboxed_product_components_in_scope :
-      ([`Complex] Flambda_arity.Component_for_creation.t
-      * (Ident.t * Flambda_kind.With_subkind.t) array)
+      ([`Complex] Flambda_arity.Component_for_creation.t * Ident.t list)
       Ident.Map.t;
-    unboxed_product_components_needed_by_continuations :
-      Ident.Set.t Continuation.Map.t;
     try_stack : Continuation.t list;
     try_stack_at_handler : Continuation.t list Continuation.Map.t;
     static_exn_continuation : Continuation.t Numeric_types.Int.Map.t;
@@ -53,15 +50,10 @@ let create ~current_unit ~return_continuation ~exn_continuation ~my_region =
   in
   let id = Ident.create_local "unused" in
   let ident_stamp_upon_starting = Ident.stamp id in
-  let unboxed_product_components_needed_by_continuations =
-    Continuation.Map.of_list
-      [return_continuation, Ident.Set.empty; exn_continuation, Ident.Set.empty]
-  in
   { current_unit;
     current_values_of_mutables_in_scope = Ident.Map.empty;
     mutables_needed_by_continuations;
     unboxed_product_components_in_scope = Ident.Map.empty;
-    unboxed_product_components_needed_by_continuations;
     try_stack = [];
     try_stack_at_handler = Continuation.Map.empty;
     static_exn_continuation = Numeric_types.Int.Map.empty;
@@ -108,12 +100,14 @@ let register_unboxed_product t ~unboxed_product ~before_unarization ~fields =
   { t with
     unboxed_product_components_in_scope =
       Ident.Map.add unboxed_product
-        (before_unarization, Array.of_list fields)
+        (before_unarization, fields)
         t.unboxed_product_components_in_scope
   }
 
-let unboxed_product_components_in_scope t =
-  Ident.Map.keys t.unboxed_product_components_in_scope
+let register_unboxed_product_with_kinds t ~unboxed_product ~before_unarization
+    ~fields =
+  register_unboxed_product t ~unboxed_product ~before_unarization
+    ~fields:(List.map fst fields)
 
 type add_continuation_result =
   { body_env : t;
@@ -139,17 +133,11 @@ let add_continuation t cont ~push_to_try_stack ~pop_region
       Continuation.Map.add cont (mutables_in_scope t)
         t.mutables_needed_by_continuations
     in
-    let unboxed_product_components_needed_by_continuations =
-      Continuation.Map.add cont
-        (unboxed_product_components_in_scope t)
-        t.unboxed_product_components_needed_by_continuations
-    in
     let try_stack =
       if push_to_try_stack then cont :: t.try_stack else t.try_stack
     in
     { t with
       mutables_needed_by_continuations;
-      unboxed_product_components_needed_by_continuations;
       try_stack;
       region_stack_in_cont_scope
     }
@@ -159,15 +147,6 @@ let add_continuation t cont ~push_to_try_stack ~pop_region
       (fun mut_var (_outer_value, kind) -> Ident.rename mut_var, kind)
       t.current_values_of_mutables_in_scope
   in
-  let unboxed_product_components_in_scope =
-    Ident.Map.map
-      (fun (before_unarization, fields) ->
-        let fields =
-          Array.map (fun (field, layout) -> Ident.rename field, layout) fields
-        in
-        before_unarization, fields)
-      t.unboxed_product_components_in_scope
-  in
   let handler_env =
     let handler_env =
       match recursive with
@@ -179,18 +158,12 @@ let add_continuation t cont ~push_to_try_stack ~pop_region
     in
     { handler_env with
       current_values_of_mutables_in_scope;
-      unboxed_product_components_in_scope;
       region_stack_in_cont_scope;
       region_stack
     }
   in
-  let extra_params_for_unboxed_products =
-    Ident.Map.data handler_env.unboxed_product_components_in_scope
-    |> List.map snd |> List.map Array.to_list |> List.concat
-  in
   let extra_params =
     Ident.Map.data handler_env.current_values_of_mutables_in_scope
-    @ extra_params_for_unboxed_products
   in
   { body_env; handler_env; extra_params }
 
@@ -241,41 +214,18 @@ let get_try_stack_at_handler t continuation =
   | stack -> stack
 
 let extra_args_for_continuation_with_kinds t cont =
-  let for_mutables =
-    match Continuation.Map.find cont t.mutables_needed_by_continuations with
-    | exception Not_found ->
-      Misc.fatal_errorf "Unbound continuation %a" Continuation.print cont
-    | mutables ->
-      let mutables = Ident.Set.elements mutables in
-      List.map
-        (fun mut ->
-          match Ident.Map.find mut t.current_values_of_mutables_in_scope with
-          | exception Not_found ->
-            Misc.fatal_errorf "No current value for %a" Ident.print mut
-          | current_value, kind -> current_value, kind)
-        mutables
-  in
-  let for_unboxed_products =
-    match
-      Continuation.Map.find cont
-        t.unboxed_product_components_needed_by_continuations
-    with
-    | exception Not_found ->
-      Misc.fatal_errorf "Unbound continuation %a" Continuation.print cont
-    | unboxed_products_to_fields ->
-      let unboxed_products = Ident.Set.elements unboxed_products_to_fields in
-      List.concat_map
-        (fun unboxed_product ->
-          match
-            Ident.Map.find unboxed_product t.unboxed_product_components_in_scope
-          with
-          | exception Not_found ->
-            Misc.fatal_errorf "No field list registered for unboxed product %a"
-              Ident.print unboxed_product
-          | _, fields -> Array.to_list fields)
-        unboxed_products
-  in
-  for_mutables @ for_unboxed_products
+  match Continuation.Map.find cont t.mutables_needed_by_continuations with
+  | exception Not_found ->
+    Misc.fatal_errorf "Unbound continuation %a" Continuation.print cont
+  | mutables ->
+    let mutables = Ident.Set.elements mutables in
+    List.map
+      (fun mut ->
+        match Ident.Map.find mut t.current_values_of_mutables_in_scope with
+        | exception Not_found ->
+          Misc.fatal_errorf "No current value for %a" Ident.print mut
+        | current_value, kind -> current_value, kind)
+      mutables
 
 let extra_args_for_continuation t cont =
   List.map fst (extra_args_for_continuation_with_kinds t cont)
@@ -289,8 +239,7 @@ let get_mutable_variable_with_kind t id =
 let get_unboxed_product_fields t id =
   match Ident.Map.find id t.unboxed_product_components_in_scope with
   | exception Not_found -> None
-  | before_unarization, fields ->
-    Some (before_unarization, List.map fst (Array.to_list fields))
+  | before_unarization, fields -> Some (before_unarization, fields)
 
 let entering_region t id ~continuation_closing_region
     ~continuation_after_closing_region =

middle_end/flambda2/from_lambda/lambda_to_flambda_env.mli

@@ -39,6 +39,13 @@ val register_mutable_variable :
 val update_mutable_variable : t -> Ident.t -> t * Ident.t
 
 val register_unboxed_product :
+  t ->
+  unboxed_product:Ident.t ->
+  before_unarization:[`Complex] Flambda_arity.Component_for_creation.t ->
+  fields:Ident.t list ->
+  t
+
+val register_unboxed_product_with_kinds :
   t ->
   unboxed_product:Ident.t ->
   before_unarization:[`Complex] Flambda_arity.Component_for_creation.t ->