Improve the results of meet functions

Rebase of ocaml-flambda/ocaml#316

Author: lthls

middle_end/flambda2/lifting/reification.ml

@@ -120,8 +120,7 @@ let try_to_reify dacc (term : Simplified_named.t) ~bound_to ~kind_of_bound_to
   match term with
   | Invalid _ ->
     let ty = T.bottom_like ty in
-    let denv = DE.add_equation_on_variable denv bound_to ty in
-    Simplified_named.invalid (), DA.with_denv dacc denv, ty
+    Simplified_named.invalid (), dacc, ty
   | Reachable _ ->
     let typing_env = DE.typing_env denv in
     let reify_result =

middle_end/flambda2/tests/meet_test.ml

@@ -40,9 +40,13 @@ let test_meet_chains_two_vars () =
     T.print new_type_for_var2;
   match T.meet env first_type_for_var2 new_type_for_var2 with
   | Bottom -> assert false
-  | Ok (meet_ty, env_extension) ->
+  | Ok (meet_result, env_extension) ->
     Format.eprintf "Env extension:@ %a\n%!" TEE.print env_extension;
     let env = TE.add_env_extension env env_extension in
+    let meet_ty =
+      Meet_result.extract_value meet_result
+        first_type_for_var2 new_type_for_var2
+    in
     let env = TE.add_equation env (Name.var var2) meet_ty in
     Format.eprintf "Final situation:@ %a\n%!" TE.print env
 
@@ -86,7 +90,11 @@ let test_meet_chains_three_vars () =
     T.print new_type_for_var3;
   match T.meet env first_type_for_var3 new_type_for_var3 with
   | Bottom -> assert false
-  | Ok (meet_ty, env_extension) ->
+  | Ok (meet_result, env_extension) ->
+    let meet_ty =
+      Meet_result.extract_value meet_result
+        first_type_for_var3 new_type_for_var3
+    in
     Format.eprintf "Env extension:@ %a\n%!" TEE.print env_extension;
     let env = TE.add_env_extension env env_extension in
     let env = TE.add_equation env (Name.var var3) meet_ty in
@@ -126,7 +134,10 @@ let meet_variants_don't_lose_aliases () =
     T.variant ~const_ctors ~non_const_ctors in
   match T.meet env ty1 ty2 with
   | Bottom -> assert false
-  | Ok (meet_ty, env_extension) ->
+  | Ok (meet_result, env_extension) ->
+    let meet_ty =
+      Meet_result.extract_value meet_result ty1 ty2
+    in
     Format.eprintf "@[<hov 2>Meet:@ %a@ /\\@ %a =>@ %a +@ %a@]@."
       T.print ty1 T.print ty2
       T.print meet_ty TEE.print env_extension;
@@ -136,7 +147,11 @@ let meet_variants_don't_lose_aliases () =
     let t_tag_1 = T.this_naked_immediate Targetint_31_63.one in
     match T.meet env t_get_tag t_tag_1 with
     | Bottom -> assert false
-    | Ok (tag_meet_ty, tag_meet_env_extension) ->
+    | Ok (tag_meet_result, tag_meet_env_extension) ->
+      let tag_meet_ty =
+        Meet_result.extract_value tag_meet_result
+          t_get_tag t_tag_1
+      in
       Format.eprintf "t_get_tag: %[email protected]_tag: %a@."
         T.print t_get_tag
         T.print t_tag_1;
@@ -179,13 +194,12 @@ let test_meet_two_blocks () =
    * test block2 block1 env; *)
 
   let f b1 b2 =
-    match
-    T.meet env
-      (T.alias_type_of K.value (Simple.var b1))
-      (T.alias_type_of K.value (Simple.var b2))
-    with
+    let ty1 = T.alias_type_of K.value (Simple.var b1) in
+    let ty2 = T.alias_type_of K.value (Simple.var b2) in
+    match T.meet env ty1 ty2 with
     | Bottom -> assert false
-    | Ok (t, tee) ->
+    | Ok (result, tee) ->
+        let t = Meet_result.extract_value result ty1 ty2 in
         Format.eprintf "Res:@ %a@.%a@."
           T.print t
           TEE.print tee;

middle_end/flambda2/types/basic/meet_result.ml

@@ -0,0 +1,61 @@
+(**************************************************************************)
+(*                                                                        *)
+(*                                 OCaml                                  *)
+(*                                                                        *)
+(*                       Vincent Laviron, OCamlPro                        *)
+(*                                                                        *)
+(*   Copyright 2021 OCamlPro SAS                                          *)
+(*                                                                        *)
+(*   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 'a return_value =
+  | Left_input
+  | Right_input
+  | Both_inputs
+  | New_result of 'a
+
+type ('a, 'ext) t =
+  | Bottom
+  | Ok of 'a return_value * 'ext
+
+let map_result ~f = function
+  | Bottom -> Bottom
+  | Ok (Left_input, ext) -> Ok (Left_input, ext)
+  | Ok (Right_input, ext) -> Ok (Right_input, ext)
+  | Ok (Both_inputs, ext) -> Ok (Both_inputs, ext)
+  | Ok (New_result x, ext) -> Ok (New_result (f x), ext)
+
+let extract_value res left right =
+  match res with
+  | Left_input -> left
+  | Right_input -> right
+  | Both_inputs -> left
+  | New_result value -> value
+
+let set_meet (type a) ~no_extension (module S : Set.S with type t = a)
+      (s1 : a) (s2 : a) =
+  match S.subset s1 s2, S.subset s2 s1 with
+  | true, true -> Ok (Both_inputs, no_extension)
+  | true, false -> Ok (Left_input, no_extension)
+  | false, true -> Ok (Right_input, no_extension)
+  | false, false ->
+    let s = S.inter s1 s2 in
+    if S.is_empty s then Bottom
+    else Ok (New_result s, no_extension)
+
+let combine rval1 left1 right1 rval2 left2 right2 f =
+  match rval1, rval2 with
+  | Both_inputs, Both_inputs -> Both_inputs
+  | (Both_inputs | Left_input), (Both_inputs | Left_input) -> Left_input
+  | (Both_inputs | Right_input), (Both_inputs | Right_input) -> Right_input
+  | New_result _, _ | _, New_result _
+  | Left_input, Right_input | Right_input, Left_input ->
+    let val1 = extract_value rval1 left1 right1 in
+    let val2 = extract_value rval2 left2 right2 in
+    New_result (f val1 val2)

middle_end/flambda2/types/basic/meet_result.mli

@@ -0,0 +1,58 @@
+(**************************************************************************)
+(*                                                                        *)
+(*                                 OCaml                                  *)
+(*                                                                        *)
+(*                       Vincent Laviron, OCamlPro                        *)
+(*                                                                        *)
+(*   Copyright 2021 OCamlPro SAS                                          *)
+(*                                                                        *)
+(*   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 for the result of meet operations *)
+
+[@@@ocaml.warning "+a-4-30-40-41-42"]
+
+type 'a return_value =
+  | Left_input
+  | Right_input
+  | Both_inputs
+  | New_result of 'a
+
+type ('a, 'ext) t =
+  | Bottom
+  | Ok of 'a return_value * 'ext
+
+val map_result
+   : f:('a -> 'b)
+  -> ('a, 'ext) t
+  -> ('b, 'ext) t
+
+(** Given a result and the inputs used to generate it,
+    returns the corresponding value.
+    In the Both_inputs case, the left input is returned. *)
+val extract_value
+   : 'a return_value
+  -> 'a
+  -> 'a
+  -> 'a
+
+val set_meet
+   : no_extension:'ext
+  -> (module Set.S with type t = 'a)
+  -> 'a
+  -> 'a
+  -> ('a, 'ext) t
+
+val combine
+   : 'a return_value
+  -> 'a
+  -> 'a
+  -> 'b return_value
+  -> 'b
+  -> 'b
+  -> ('a -> 'b -> 'c)
+  -> 'c return_value

middle_end/flambda2/types/env/typing_env.rec.ml

@@ -784,8 +784,17 @@ let mem ?min_name_mode t name =
     ~symbol:(fun sym ->
       (* CR mshinwell: This might not take account of symbols in missing
          .cmx files *)
-      Symbol.Set.mem sym t.defined_symbols
-        || Name.Set.mem name (t.get_imported_names ()))
+      let comp_unit = Symbol.compilation_unit sym in
+      if Compilation_unit.equal comp_unit (Compilation_unit.get_current_exn ())
+      then
+        Symbol.Set.mem sym t.defined_symbols
+      else
+        match (resolver t) comp_unit with
+        | None ->
+          (* The cmx is unavailable, but the symbol is valid *)
+          true
+        | Some _ ->
+          Name.Set.mem name (t.get_imported_names ()))
 
 let mem_simple ?min_name_mode t simple =
   Simple.pattern_match simple
@@ -947,18 +956,22 @@ let rec add_equation0 (t:t) name ty =
       | exception Not_found -> true
       | _ -> false
     in
-    if is_concrete then begin
-      let canonical =
-        Aliases.get_canonical_ignoring_name_mode (aliases t) name
-        |> Simple.without_coercion
-      in
-      if not (Simple.equal canonical (Simple.name name)) then begin
-        Misc.fatal_errorf "Trying to add equation giving concrete type on %a \
-            which is not canonical (its canonical is %a): %a"
-          Name.print name
-          Simple.print canonical
-          Type_grammar.print ty
-      end
+    let canonical =
+      Aliases.get_canonical_ignoring_name_mode (aliases t) name
+      |> Simple.without_coercion
+    in
+    if is_concrete && not (Simple.equal canonical (Simple.name name)) then begin
+      Misc.fatal_errorf "Trying to add equation giving concrete type on %a \
+          which is not canonical (its canonical is %a): %a"
+        Name.print name
+        Simple.print canonical
+        Type_grammar.print ty
+    end;
+    if not is_concrete && Simple.equal canonical (Simple.name name) then begin
+      Misc.fatal_errorf "Trying to add equation giving alias type on %a \
+          which is canonical: %a"
+        Name.print name
+        Type_grammar.print ty
     end
   end;
   invariant_for_new_equation t name ty;
@@ -1110,14 +1123,14 @@ and add_equation t name ty =
   let ty, t =
     let [@inline always] name eqn_name ~coercion =
       assert (Coercion.is_id coercion); (* true by definition *)
-      if Name.equal name eqn_name then ty, t
-      else
-        let env = Meet_env.create t in
-        let existing_ty = find t eqn_name (Some (Type_grammar.kind ty)) in
-        match Type_grammar.meet env ty existing_ty with
-        | Bottom -> Type_grammar.bottom_like ty, t
-        | Ok (meet_ty, env_extension) ->
-          meet_ty, add_env_extension t env_extension
+      let env = Meet_env.create t in
+      let existing_ty = find t eqn_name (Some (Type_grammar.kind ty)) in
+      match Type_grammar.meet env ty existing_ty with
+      | Bottom ->
+        Type_grammar.bottom_like ty, t
+      | Ok (meet_result, env_extension) ->
+        let meet_ty = Meet_result.extract_value meet_result ty existing_ty in
+        meet_ty, add_env_extension t env_extension
     in
     Simple.pattern_match bare_lhs ~name ~const:(fun _ -> ty, t)
   in
@@ -1207,8 +1220,24 @@ let meet_equations_on_params t ~params ~param_types =
       let existing_type = find t name (Some kind) in
       let env = Meet_env.create t in
       match Type_grammar.meet env existing_type param_type with
-      | Bottom -> add_equation t name (Type_grammar.bottom kind)
-      | Ok (meet_ty, env_extension) ->
+      | Bottom ->
+        (* Ideally we should adapt the type of this function so
+           that the caller can handle this case.
+           In practice, this would likely mean treating the
+           corresponding continuation handler as unreachable. *)
+        (* Misc.fatal_errorf
+         *   "Bottom equation added on parameter:@.\
+         *    Parameter:@ %a@.\
+         *    Parameter type:@ %a@.\
+         *    Typing env:@ %a"
+         *   Kinded_parameter.print param
+         *   Type_grammar.print param_type
+         *   print t *)
+        add_equation t name (Type_grammar.bottom kind)
+      | Ok (meet_result, env_extension) ->
+        let meet_ty =
+          Meet_result.extract_value meet_result existing_type param_type
+        in
         let t = add_equation t name meet_ty in
         add_env_extension t env_extension)
     t

middle_end/flambda2/types/env/typing_env_extension.rec.ml

@@ -115,9 +115,18 @@ let rec meet0 env (t1 : t) (t2 : t) extra_extensions =
         | Some ty0 ->
           begin match Type_grammar.meet env ty0 ty with
           | Bottom -> raise Bottom_meet
-          | Ok (ty, new_ext) ->
-            Type_grammar.check_equation name ty;
-            Name.Map.add (*replace*) name ty eqs, new_ext :: extra_extensions
+          | Ok (res, new_ext) ->
+            let extra_extensions = new_ext :: extra_extensions in
+            begin match res with
+            | Left_input | Both_inputs ->
+              (* The equation is already there *)
+              eqs, extra_extensions
+            | Right_input ->
+              Name.Map.add (*replace*) name ty eqs, extra_extensions
+            | New_result ty ->
+              Type_grammar.check_equation name ty;
+              Name.Map.add (*replace*) name ty eqs, extra_extensions
+            end
           end)
       t2.equations
       (t1.equations, extra_extensions)

middle_end/flambda2/types/env/typing_env_extension.rec.mli

@@ -45,6 +45,8 @@ val from_map : Type_grammar.t Name.Map.t -> t
 
 val add_or_replace_equation : t -> Name.t -> Type_grammar.t -> t
 
+(* Note: this doesn't use Meet_result.t, as the meet function here is used
+   to merge extensions together. *)
 val meet : Meet_env.t -> t -> t -> t Or_bottom.t
 
 val join : Join_env.t -> t -> t -> t

middle_end/flambda2/types/flambda_type.mli

@@ -228,7 +228,7 @@ module Typing_env : sig
   end
 end
 
-val meet : Typing_env.t -> t -> t -> (t * Typing_env_extension.t) Or_bottom.t
+val meet : Typing_env.t -> t -> t -> (t, Typing_env_extension.t) Meet_result.t
 
 val meet_shape
    : Typing_env.t