Rec_info: Add coercions

compilerlibs/Makefile.compilerlibs

@@ -189,6 +189,7 @@ MIDDLE_END_FLAMBDA_COMPILENV_DEPS=\
   middle_end/flambda/compilenv_deps/flambda_colours.cmo \
   middle_end/flambda/compilenv_deps/compilation_unit.cmo \
   middle_end/flambda/compilenv_deps/rec_info.cmo \
+  middle_end/flambda/compilenv_deps/coercion.cmo \
   middle_end/flambda/compilenv_deps/reg_width_things.cmo \
   middle_end/flambda/compilenv_deps/symbol.cmo \
   middle_end/flambda/compilenv_deps/variable.cmo \

middle_end/flambda/basic/kinded_parameter.ml

@@ -111,8 +111,6 @@ module List = struct
 
   let name_set t = Name.Set.of_list (List.map Name.var (vars t))
 
-  let simple_set t = Simple.Set.of_list (simples t)
-
   let rename t = List.map (fun t -> rename t) t
 
   let arity t = List.map (fun t -> Flambda_kind.With_subkind.kind (kind t)) t

middle_end/flambda/basic/kinded_parameter.mli

@@ -68,8 +68,6 @@ module List : sig
   (** As for [var_set] but returns a set of [Name]s. *)
   val name_set : t -> Name.Set.t
 
-  val simple_set : t -> Simple.Set.t
-
   val equal_vars : t -> Variable.t list -> bool
 
   val rename : t -> t

middle_end/flambda/basic/name.ml

@@ -104,3 +104,11 @@ let must_be_symbol_opt t =
   pattern_match t
     ~var:(fun _ -> None)
     ~symbol:(fun sym -> Some sym)
+
+module Pair = struct
+  include Identifiable.Make_pair
+    (Reg_width_things.Name)
+    (Reg_width_things.Name)
+
+  type nonrec t = t * t
+end

middle_end/flambda/basic/name.mli

@@ -55,3 +55,10 @@ val must_be_var_opt : t -> Variable.t option
 val must_be_symbol_opt : t -> Symbol.t option
 
 val rename : t -> t
+
+module Pair : sig
+  type nonrec t = t * t
+
+  include Identifiable.S with type t := t
+end
+

middle_end/flambda/basic/simple.ml

@@ -53,25 +53,30 @@ let pattern_match' t ~var ~symbol ~const =
 
 let const_from_descr descr = const (RWC.of_descr descr)
 
-let without_rec_info t = pattern_match t ~name ~const
+let without_coercion t = pattern_match t ~name ~const
 
-let merge_rec_info t ~newer_rec_info =
-  if is_const t then None
+let apply_coercion t applied_coercion =
+  if Coercion.is_id applied_coercion then Some t
   else
-    match newer_rec_info with
-    | None -> Some t
-    | Some newer_rec_info ->
-      let rec_info =
-        match rec_info t with
-        | None -> newer_rec_info
-        | Some older_rec_info ->
-          Rec_info.merge older_rec_info ~newer:newer_rec_info
-      in
-      Some (with_rec_info (without_rec_info t) rec_info)
+    let coercion = 
+      let existing_coercion = coercion t in
+      if Coercion.is_id existing_coercion then Some applied_coercion
+      else Coercion.compose existing_coercion ~then_:applied_coercion
+    in
+    coercion
+    |> Option.map (fun coercion -> with_coercion (without_coercion t) coercion)
+
+let apply_coercion_exn t applied_coercion =
+  match apply_coercion t applied_coercion with
+  | Some t -> t
+  | None ->
+    Misc.fatal_errorf "Cannot apply coercion %a to %a"
+      print t
+      Coercion.print applied_coercion
 
 (* CR mshinwell: Make naming consistent with [Name] re. the option type *)
 
-(* CR mshinwell: Careful that Rec_info doesn't get dropped using the
+(* CR mshinwell: Careful that coercions don't get dropped using the
    following *)
 
 let [@inline always] must_be_var t =
@@ -86,7 +91,7 @@ let [@inline always] must_be_name t =
 let to_name t =
   match must_be_name t with
   | None -> None
-  | Some name -> Some (rec_info t, name)
+  | Some name -> Some (coercion t, name)
 
 let map_name t ~f =
   match must_be_name t with
@@ -158,14 +163,6 @@ module List = struct
     else result
 end
 
-module Pair = struct
-  include Identifiable.Make_pair
-    (Reg_width_things.Simple)
-    (Reg_width_things.Simple)
-
-  type nonrec t = t * t
-end
-
 module With_kind = struct
   type nonrec t = t * Flambda_kind.t
 

middle_end/flambda/basic/simple.mli

@@ -23,11 +23,13 @@ include module type of struct include Reg_width_things.Simple end
 
 include Contains_names.S with type t := t
 
-val has_rec_info : t -> bool
+val has_coercion : t -> bool
 
-val merge_rec_info : t -> newer_rec_info:Rec_info.t option -> t option
+val apply_coercion : t -> Coercion.t -> t option
 
-val without_rec_info : t -> t
+val apply_coercion_exn : t -> Coercion.t -> t
+
+val without_coercion : t -> t
 
 val must_be_var : t -> Variable.t option
 
@@ -66,7 +68,7 @@ val const_from_descr : Reg_width_const.Descr.t -> t
 
 val map_name : t -> f:(Name.t -> Name.t) -> t
 
-val to_name : t -> (Rec_info.t option * Name.t) option
+val to_name : t -> (Coercion.t * Name.t) option
 
 (* CR mshinwell: remove these next two? *)
 val map_var : t -> f:(Variable.t -> Variable.t) -> t
@@ -94,12 +96,6 @@ module List : sig
   include Identifiable.S with type t := t
 end
 
-module Pair : sig
-  type nonrec t = t * t
-
-  include Identifiable.S with type t := t
-end
-
 module With_kind : sig
   type nonrec t = t * Flambda_kind.t
 

middle_end/flambda/cmx/flambda_cmx_format.ml

@@ -62,7 +62,7 @@ let create ~final_typing_env ~all_code ~exported_offsets ~used_closure_vars =
       Variable.Map.empty
   in
   let simples =
-    Simple.Set.fold (fun simple simples ->
+    Reg_width_things.Simple.Set.fold (fun simple simples ->
         Simple.Map.add simple (Simple.export simple) simples)
       exported_ids.simples
       Simple.Map.empty

middle_end/flambda/cmx/ids_for_export.ml

@@ -75,9 +75,9 @@ let add_name t name =
 
 let add_simple t simple =
   let simples =
-    match Simple.rec_info simple with
-    | None -> t.simples
-    | Some _ -> Simple.Set.add simple t.simples
+    match Simple.coercion simple with
+    | Id -> t.simples
+    | _ -> Simple.Set.add simple t.simples
   in
   let t = { t with simples; } in
   Simple.pattern_match simple
@@ -90,14 +90,13 @@ let add_code_id t code_id =
 let add_continuation t continuation =
   { t with continuations = Continuation.Set.add continuation t.continuations }
 
-
 let from_simple simple =
   let simples =
-    match Simple.rec_info simple with
-    | None ->
+    match Simple.coercion simple with
+    | Id ->
       (* This simple will not be in the grand_table_of_simples *)
       Simple.Set.empty
-    | Some _ -> Simple.Set.singleton simple
+    | _ -> Simple.Set.singleton simple
   in
   Simple.pattern_match simple
     ~const:(fun const ->

middle_end/flambda/compare/compare.ml

@@ -766,6 +766,25 @@ let function_decls env decl1 decl2 : unit Comparison.t =
   else Different { approximant = () }
 ;;
 
+(** Match up equal elements in two lists and iterate through both of them,
+    using [f] analogously to [Map.S.merge] *)
+let iter2_merged l1 l2 ~compare ~f =
+  let l1 = List.sort compare l1 in
+  let l2 = List.sort compare l2 in
+  let rec go l1 l2 =
+    match l1, l2 with
+    | []      , []       -> ()
+    | a1 :: l1, []       -> f (Some a1) None     ; go l1         []
+    | []      , a2 :: l2 -> f None      (Some a2); go []         l2 
+    | a1 :: l1, a2 :: l2 ->
+      begin match compare a1 a2 with
+      | 0                -> f (Some a1) (Some a2); go l1         l2
+      | c when c < 0     -> f (Some a1) None     ; go l1         (a2 :: l2)
+      | _                -> f None      (Some a2); go (a1 :: l1) l2
+      end
+  in
+  go l1 l2
+
 let sets_of_closures env set1 set2 : Set_of_closures.t Comparison.t =
   (* Need to do unification on closure vars and closure ids, we we're going to
    * invert both maps, figuring the closure vars with the same value should be
@@ -777,28 +796,28 @@ let sets_of_closures env set1 set2 : Set_of_closures.t Comparison.t =
     |> List.map (fun (var, value) ->
          subst_simple env value, var
        )
-    |> Simple.Map.of_list
   in
   (* We want to process the whole map to find new correspondences between
    * closure vars, so we need to remember whether we've found any mismatches *)
   let ok = ref true in
-  (* Using merge here as a map version of [List.iter2]; always returning None
-   * means the returned map is always empty, so this shouldn't waste much *)
-  let _ : unit Simple.Map.t =
-    Simple.Map.merge (fun _value var1 var2 ->
-      begin
-        match var1, var2 with
-        | None, None -> ()
-        | Some _, None | None, Some _ -> ok := false
-        | Some var1, Some var2 ->
-          begin
-            match closure_vars env var1 var2 with
-            | Equivalent -> ()
-            | Different { approximant = _ } -> ok := false
-          end
-      end;
-      None
-    ) (closure_vars_by_value set1) (closure_vars_by_value set2)
+  let () =
+    let compare (value1, _var1) (value2, _var2) =
+      Simple.compare value1 value2
+    in
+    iter2_merged (closure_vars_by_value set1) (closure_vars_by_value set2)
+      ~compare
+      ~f:(fun elt1 elt2 ->
+        begin
+          match elt1, elt2 with
+          | None, None -> ()
+          | Some _, None | None, Some _ -> ok := false
+          | Some (_value1, var1), Some (_value2, var2) ->
+            begin
+              match closure_vars env var1 var2 with
+              | Equivalent -> ()
+              | Different { approximant = _ } -> ok := false
+            end
+        end)
   in
   let closure_ids_and_fun_decls_by_code_id set =
     let map = Function_declarations.funs (Set_of_closures.function_decls set) in
@@ -809,6 +828,8 @@ let sets_of_closures env set1 set2 : Set_of_closures.t Comparison.t =
        )
     |> Code_id.Map.of_list
   in
+  (* Using merge here as a map version of [List.iter2]; always returning None
+   * means the returned map is always empty, so this shouldn't waste much *)
   let _ : unit Code_id.Map.t =
     Code_id.Map.merge (fun _code_id value1 value2 ->
       begin

middle_end/flambda/compilenv_deps/coercion.ml

@@ -0,0 +1,74 @@
+(**************************************************************************)
+(*                                                                        *)
+(*                                 OCaml                                  *)
+(*                                                                        *)
+(*                   Mark Shinwell, Jane Street Europe                    *)
+(*                                                                        *)
+(*   Copyright 2019 Jane Street Group LLC                                 *)
+(*                                                                        *)
+(*   All rights reserved.  This file is distributed under the terms of    *)
+(*   the GNU Lesser General Public License version 2.1, with the          *)
+(*   special exception on linking described in the file LICENSE.          *)
+(*                                                                        *)
+(**************************************************************************)
+
+type t =
+  | Id
+  | Non_id of {
+      from_depth : int;
+      to_depth : int;
+    }
+
+let id = Id
+
+let change_depth ~from:from_depth ~to_:to_depth =
+  if from_depth = to_depth then Id else Non_id { from_depth; to_depth }
+
+let is_id = function
+  | Id -> true
+  | Non_id _ -> false
+
+let inverse = function
+  | Id -> Id
+  | Non_id { from_depth; to_depth } ->
+    Non_id { from_depth = to_depth; to_depth = from_depth }
+
+let print ppf = function
+  | Id ->
+    Format.fprintf ppf "@<0>%sid@<0>%s"
+      (Flambda_colours.elide ())
+      (Flambda_colours.normal ())
+  | Non_id { from_depth; to_depth; } ->
+    Format.fprintf ppf "@<0>%s@[<hov 1>(depth@ %d ->@ %d)@]@<0>%s"
+      (Flambda_colours.coercion ())
+      from_depth to_depth
+      (Flambda_colours.normal ())
+
+let compose t1 ~then_:t2 =
+  match t1, t2 with
+  | Id, _ -> Some t2
+  | _, Id -> Some t1
+  | Non_id { from_depth = from_depth1; to_depth = to_depth1 },
+    Non_id { from_depth = from_depth2; to_depth = to_depth2 } ->
+    if to_depth1 = from_depth2 then
+      Some (change_depth ~from:from_depth1 ~to_:to_depth2)
+    else
+      None
+
+let compose_exn t1 ~then_:t2 =
+  match compose t1 ~then_:t2 with
+  | Some t -> t
+  | None ->
+    Misc.fatal_errorf "Invalid composition: %a@ >>@ %a" print t1 print t2
+
+let equal t1 t2 =
+  match t1, t2 with
+  | Id, Id -> true
+  | Non_id { from_depth = from_depth1; to_depth = to_depth1 },
+    Non_id { from_depth = from_depth2; to_depth = to_depth2 } ->
+    from_depth1 = from_depth2 && to_depth1 = to_depth2
+  | _, _ -> false
+
+let hash = Hashtbl.hash
+
+let apply_to_rec_info _ rec_info = rec_info

middle_end/flambda/compilenv_deps/coercion.mli

@@ -0,0 +1,42 @@
+(**************************************************************************)
+(*                                                                        *)
+(*                                 OCaml                                  *)
+(*                                                                        *)
+(*                   Mark Shinwell, Jane Street Europe                    *)
+(*                                                                        *)
+(*   Copyright 2019 Jane Street Group LLC                                 *)
+(*                                                                        *)
+(*   All rights reserved.  This file is distributed under the terms of    *)
+(*   the GNU Lesser General Public License version 2.1, with the          *)
+(*   special exception on linking described in the file LICENSE.          *)
+(*                                                                        *)
+(**************************************************************************)
+
+[@@@ocaml.warning "+a-4-30-40-41-42"]
+
+type t = private
+  | Id
+  | Non_id of {
+      from_depth : int;
+      to_depth : int;
+    }
+
+val change_depth : from:int -> to_:int -> t
+
+val id : t
+
+val is_id : t -> bool
+
+val inverse : t -> t
+
+val compose : t -> then_:t -> t option
+
+val compose_exn : t -> then_:t -> t
+
+val print : Format.formatter -> t -> unit
+
+val equal : t -> t -> bool
+
+val hash : t -> int
+
+val apply_to_rec_info : t -> Rec_info.t -> Rec_info.t