Allow overlap of invalid rewrite id and extra_args

As the comment states, it can happen that, when adding a extra param and
args, there is an overlap between the domain of the extra_args map, and
the set of invalids (because of the way these are computed by the
unboxing code). This case is reasonable, so in such a case, we can allow
the invalid set to take precedenhce.

Author: Gbury

middle_end/flambda2/simplify/continuation_extra_params_and_args.ml

@@ -65,18 +65,12 @@ let empty = Empty
 let is_empty = function Empty -> true | Non_empty _ -> false
 
 let add t ~invalids ~extra_param ~extra_args =
-  if not
-       (Apply_cont_rewrite_id.Set.is_empty
-          (Apply_cont_rewrite_id.Set.inter invalids
-             (Apply_cont_rewrite_id.Map.keys extra_args)))
-  then
-    Misc.fatal_errorf
-      "Broken invariants: when adding an extra param to a continuation, every \
-       Apply_cont_rewrite_id should either have a valid extra arg, or be \
-       invalid, but not both:@ %a@ %a"
-      Apply_cont_rewrite_id.Set.print invalids
-      (Apply_cont_rewrite_id.Map.print Extra_arg.print)
-      extra_args;
+  (* Note: there can be some overlap between the invalid ids and the keys of the
+     [extra_args] map. This is notably used by the unboxing code which may
+     compute some extra args and only later (when computing extra args for
+     another parameter) realize that some rewrite ids are invalids, and then
+     call this function with this new invalid set and the extra_args computed
+     before this invalid set was known. *)
   match t with
   | Empty ->
     let extra_params = Bound_parameters.create [extra_param] in
@@ -96,6 +90,10 @@ let add t ~invalids ~extra_param ~extra_args =
     let extra_args =
       Apply_cont_rewrite_id.Map.merge
         (fun id already_extra_args extra_arg ->
+          (* The [invalids] set is expected to be small (actually, empty most of
+             the time), so the lookups in each case of the merge should be
+             reasonable, compared to merging (and allocating) the [invalids] set
+             and the [extra_args] map. *)
           match already_extra_args, extra_arg with
           | None, None -> None
           | None, Some _ ->
@@ -125,7 +123,9 @@ let add t ~invalids ~extra_param ~extra_args =
                 extra_args Apply_cont_rewrite_id.Set.print invalids print t
           | Some Or_invalid.Invalid, Some _ -> Some Or_invalid.Invalid
           | Some (Or_invalid.Ok already_extra_args), Some extra_arg ->
-            Some (Or_invalid.Ok (extra_arg :: already_extra_args)))
+            if Apply_cont_rewrite_id.Set.mem id invalids
+            then Some Or_invalid.Invalid
+            else Some (Or_invalid.Ok (extra_arg :: already_extra_args)))
         already_extra_args extra_args
     in
     Non_empty { extra_params; extra_args }