Ensure consistent variable binding times (#3217)

When joining environments, we exclude variables that are no longer
useful in the joined environment. This causes the remaining variables to
have new (lower) binding times.

If we later join the resulting environment again with another
environment that still has those variables, we end up having the
same variable with different binding times in the joined level, and we
define the variable twice in the resulting environment.

This can allow breaking the order between two variable's binding times,
which means that different canonical elements can be chosen in different
environments for the same pair of variables.

This patch fixes the issue by keeping track of `next_binding_time`
*globally* across all related environments and keeping the binding time
during join, ensuring that the binding time of a variable is globally
unique and preventing any such re-orderings.

Longer explanation
------------------

Consider 4 environments that define variables `a`, `b`, `c`, and `d`.

```
      _________________________________
     /            |            |       \
  Env 1         Env 2        Env 3    Env 4
   /              |            |         \
define a       define a     define a  define a
   |              |            |         |
define b       define b     define b  define b
   |              |            |         |
define c       define c     define c  define c
   |              |            |         |
define d       define d     define d  define d

```

Suppose that we first join environments 3 and 4 (filtering out `a` and
`b`) and environments 1 and 2 (filtering out `d`):

```
      ____________
     /            \
  Env 1 \/ 2      Env 3 \/ 4
   /                \
define a         define c
   |                |
define b         define d
   |
define c
```

Once we join the two resulting environments, if we filter out `a` and
`b` but keep `c` and `d`, the current implementation creates a level
which defines `c`, then `d`, then `c` again. Since we keep the last
definition, we end up with the environment:

```
define d
   |
define c
```

We have now swapped the binding times of `c` and `d` which, if they were
in the same alias set, can introduce subtle bugs.

Author: bclement-ocp

middle_end/flambda2/types/env/typing_env.ml

@@ -113,7 +113,10 @@ type t =
     (* [prev_levels] is sorted with the greatest scope at the head of the
        list *)
     current_level : One_level.t;
-    next_binding_time : Binding_time.t;
+    next_binding_time : Binding_time.t ref;
+    (* [next_binding_time] is a reference shared by all related environments to
+       the latest binding time ever used. Used to ensure consistent binding
+       times between variables defined in parallel branches. *)
     min_binding_time : Binding_time.t;
         (* Earlier variables have mode In_types *)
     is_bottom : bool
@@ -385,7 +388,7 @@ let create ~resolver ~get_imported_names =
        to allow an efficient implementation of [cut] (see below), we always
        increment the scope by one here. *)
     current_level = One_level.create_empty (Scope.next Scope.initial);
-    next_binding_time = Binding_time.earliest_var;
+    next_binding_time = ref Binding_time.earliest_var;
     defined_symbols = Symbol.Set.empty;
     code_age_relation = Code_age_relation.empty;
     min_binding_time = Binding_time.earliest_var;
@@ -620,17 +623,14 @@ let alias_is_bound_strictly_earlier t ~bound_name ~alias =
 
 let with_current_level t ~current_level = { t with current_level }
 
-let with_current_level_and_next_binding_time t ~current_level next_binding_time
-    =
-  { t with current_level; next_binding_time }
-
 let with_aliases t ~aliases =
   let current_level = One_level.with_aliases t.current_level ~aliases in
   with_current_level t ~current_level
 
 let cached t = One_level.just_after_level t.current_level
 
-let add_variable_definition t var kind name_mode =
+let add_variable_definition_with_binding_time ~binding_time t var kind name_mode
+    =
   (* We can add equations in our own compilation unit on variables and symbols
      defined in another compilation unit. However we can't define other
      compilation units' variables or symbols (except for predefined symbols such
@@ -644,24 +644,28 @@ let add_variable_definition t var kind name_mode =
        %a@ in environment:@ %a"
       Variable.print var print t;
   let name = Name.var var in
-  if Flambda_features.check_invariants () && mem t name
+  if Flambda_features.check_light_invariants () && mem t name
   then
     Misc.fatal_errorf "Cannot rebind %a in environment:@ %a" Name.print name
       print t;
   let level =
-    TEL.add_definition
-      (One_level.level t.current_level)
-      var kind t.next_binding_time
+    TEL.add_definition (One_level.level t.current_level) var kind binding_time
   in
   let just_after_level =
     Cached_level.add_or_replace_binding (cached t) name (MTC.unknown kind)
-      t.next_binding_time name_mode
+      binding_time name_mode
   in
   let current_level =
     One_level.create (current_scope t) level ~just_after_level
   in
-  with_current_level_and_next_binding_time t ~current_level
-    (Binding_time.succ t.next_binding_time)
+  with_current_level t ~current_level
+
+let get_next_binding_time t = !(t.next_binding_time)
+
+let add_variable_definition t var kind name_mode =
+  let binding_time = get_next_binding_time t in
+  t.next_binding_time := Binding_time.succ binding_time;
+  add_variable_definition_with_binding_time ~binding_time t var kind name_mode
 
 let add_symbol_definition t sym =
   (* CR-someday mshinwell: check for redefinition when invariants enabled? *)
@@ -958,7 +962,9 @@ and add_env_extension_with_extra_variables t
 let add_env_extension_from_level t level ~meet_type : t =
   let t =
     TEL.fold_on_defined_vars
-      (fun var kind t -> add_variable_definition t var kind Name_mode.in_types)
+      (fun ~binding_time var kind t ->
+        add_variable_definition_with_binding_time ~binding_time t var kind
+          Name_mode.in_types)
       level t
   in
   let t =
@@ -1194,7 +1200,7 @@ let compute_joined_aliases base_env alias_candidates envs_at_uses =
   with_aliases base_env ~aliases:new_aliases
 
 let closure_env t =
-  increment_scope { t with min_binding_time = t.next_binding_time }
+  increment_scope { t with min_binding_time = get_next_binding_time t }
 
 let rec free_names_transitive_of_type_of_name t name ~result =
   let result = Name_occurrences.add_name result name Name_mode.in_types in

middle_end/flambda2/types/env/typing_env_level.ml

@@ -69,11 +69,11 @@ let [@ocamlformat "disable"] print ppf
 
 let fold_on_defined_vars f t init =
   Binding_time.Map.fold
-    (fun _bt vars acc ->
+    (fun bt vars acc ->
       Variable.Set.fold
         (fun var acc ->
           let kind = Variable.Map.find var t.defined_vars in
-          f var kind acc)
+          f ~binding_time:bt var kind acc)
         vars acc)
     t.binding_times init
 
@@ -113,7 +113,8 @@ let add_symbol_projection t var proj =
   { t with symbol_projections }
 
 let add_definition t var kind binding_time =
-  if Flambda_features.check_invariants () && Variable.Map.mem var t.defined_vars
+  if Flambda_features.check_light_invariants ()
+     && Variable.Map.mem var t.defined_vars
   then
     Misc.fatal_errorf "Environment extension already binds variable %a:@ %a"
       Variable.print var print t;

middle_end/flambda2/types/env/typing_env_level.mli

@@ -58,6 +58,9 @@ val variables_by_binding_time : t -> Variable.Set.t Binding_time.Map.t
 val variable_is_defined : t -> Variable.t -> bool
 
 val fold_on_defined_vars :
-  (Variable.t -> Flambda_kind.t -> 'a -> 'a) -> t -> 'a -> 'a
+  (binding_time:Binding_time.t -> Variable.t -> Flambda_kind.t -> 'a -> 'a) ->
+  t ->
+  'a ->
+  'a
 
 val as_extension_without_bindings : t -> Type_grammar.Env_extension.t