flambda-backend: Undo sort changes during snapshot backtrack (#1885)

* record and undo sort changes

* add test and cr

* format and comment

* remove test that's no longer relevant

Author: alanechang

typing/jkind.ml

@@ -31,6 +31,15 @@ module Sort = struct
 
   and var = t option ref
 
+  (* To record changes to sorts, for use with `Types.{snapshot, backtrack}` *)
+  type change = var * t option
+
+  let change_log : (change -> unit) ref = ref (fun _ -> ())
+
+  let log_change change = !change_log change
+
+  let undo_change (v, t_op) = v := t_op
+
   let var_name : var -> string =
     let next_id = ref 1 in
     let named = ref [] in
@@ -58,6 +67,11 @@ module Sort = struct
 
   let new_var () = Var (ref None)
 
+  let set : var -> t option -> unit =
+   fun v t_op ->
+    log_change (v, t_op);
+    v := t_op
+
   (* Post-condition: If the result is a [Var v], then [!v] is [None]. *)
   let rec get : t -> t = function
     | Const _ as t -> t
@@ -66,7 +80,7 @@ module Sort = struct
       | None -> t
       | Some s ->
         let result = get s in
-        if result != s then r := Some result;
+        if result != s then set r (Some result);
         (* path compression *)
         result)
 
@@ -86,11 +100,11 @@ module Sort = struct
     | Var r -> (
       match !r with
       | None ->
-        r := default_value;
+        set r default_value;
         Value
       | Some s ->
         let result = get_default_value s in
-        r := default result;
+        set r (default result);
         (* path compression *)
         result)
 
@@ -119,7 +133,7 @@ module Sort = struct
     match !v1 with
     | Some s1 -> equate_sort_const s1 c2
     | None ->
-      v1 := Some (of_const c2);
+      set v1 (Some (of_const c2));
       Equal_mutated_first
 
   and equate_var v1 s2 =
@@ -135,7 +149,7 @@ module Sort = struct
       | Some s1, _ -> swap_equate_result (equate_var v2 s1)
       | _, Some s2 -> equate_var v1 s2
       | None, None ->
-        v1 := Some (of_var v2);
+        set v1 (Some (of_var v2));
         Equal_mutated_first
 
   and equate_sort_const s1 c2 =
@@ -167,7 +181,7 @@ module Sort = struct
       match !v with
       (* CR layouts v5: this should probably default to void now *)
       | None ->
-        v := some_value;
+        set v some_value;
         false
       | Some s -> is_void_defaulting s)
     | Const Value -> false

typing/jkind.mli

@@ -82,6 +82,13 @@ module Sort : sig
       it is set to [value] first. *)
   val get_default_value : t -> const
 
+  (** To record changes to sorts, for use with `Types.{snapshot, backtrack}` *)
+  type change
+
+  val change_log : (change -> unit) ref
+
+  val undo_change : change -> unit
+
   module Debug_printers : sig
     val t : Format.formatter -> t -> unit
 

typing/types.ml

@@ -609,6 +609,7 @@ type change =
   | Ccommu : [`var] commutable_gen -> change
   | Cuniv : type_expr option ref * type_expr option -> change
   | Cmodes : Mode.changes -> change
+  | Csort : Jkind.Sort.change -> change
 
 type changes =
     Change of change * changes ref
@@ -623,7 +624,8 @@ let log_change ch =
   trail := r'
 
 let () =
-  Mode.change_log := (fun changes -> log_change (Cmodes changes))
+  Mode.change_log := (fun changes -> log_change (Cmodes changes));
+  Jkind.Sort.change_log := (fun change -> log_change (Csort change))
 
 (* constructor and accessors for [field_kind] *)
 
@@ -873,6 +875,7 @@ let undo_change = function
   | Ccommu (Cvar r)  -> r.commu <- Cunknown
   | Cuniv  (r, v)    -> r := v
   | Cmodes ms -> Mode.undo_changes ms
+  | Csort change -> Jkind.Sort.undo_change change
 
 type snapshot = changes ref * int
 let last_snapshot = Local_store.s_ref 0
@@ -1030,5 +1033,3 @@ let undo_compress (changes, _old) =
             Transient_expr.set_desc ty desc; r := !next
         | _ -> ())
         log
-
-