flambda-backend: Fix issue with gadts and layouts (#2618)

* Fix issue with gadts and layouts

* Improved fix

* Guard fix

* when clauses should be pure

Author: ccasin

testsuite/tests/typing-layouts/basics.ml

@@ -2493,3 +2493,80 @@ Line 1, characters 13-47:
 Error: The primitive [%sendcache] is used in an invalid declaration.
        The declaration contains argument/return types with the wrong layout.
 |}]
+
+(*********************************************************)
+(* Test 43: GADT refinement works on layouts as expected *)
+
+type ('a : any) t_gadt_simple =
+  | Float64 : ('a : float64) t_gadt_simple
+
+let f_match_allowed (type a : any) (x : a t_gadt_simple) : int =
+  match x with
+  | Float64 -> 1;;
+[%%expect{|
+type ('a : any) t_gadt_simple = Float64 : ('a : float64). 'a t_gadt_simple
+val f_match_allowed : ('a : any). 'a t_gadt_simple -> int = <fun>
+|}]
+
+let not_magic  (type a : any) (x : a t_gadt_simple) : 'b =
+  match x with
+  | _ -> .
+[%%expect{|
+Line 3, characters 4-5:
+3 |   | _ -> .
+        ^
+Error: This match case could not be refuted.
+       Here is an example of a value that would reach it: Float64
+|}]
+
+type ('a : any) t =
+  | UFloat : float# t
+  | Float : float t
+  | Int : int t
+
+let make_pi (type a : any) (x : a t) : unit -> a =
+  match x with
+  | UFloat -> fun () -> #3.14
+  | Float -> fun () -> 3.14
+  | Int -> fun () -> 3;;
+[%%expect{|
+type ('a : any) t = UFloat : float# t | Float : float t | Int : int t
+val make_pi : ('a : any). 'a t -> unit -> 'a = <fun>
+|}]
+
+type ('a : any) repr =
+  | Float64 : ('a : float64) repr
+  | Value : ('a : value) repr
+
+let lpoly_id (type a : any) (x : a repr) : a -> a =
+  match x with
+  | Float64 -> fun x -> x
+  | Value -> fun x -> x
+[%%expect{|
+type ('a : any) repr = Float64 : ('a : float64). 'a repr | Value : 'a repr
+val lpoly_id : ('a : any). 'a repr -> 'a -> 'a = <fun>
+|}]
+
+type 'a s = 'a
+
+module M = struct
+  type t : immediate
+end
+
+module N = struct
+  type ('a,'b) eq =
+    | Refl : ('a, 'a) eq
+
+  let f (x : (M.t, 'a s) eq) : int =
+    match x with
+    | Refl -> 42
+end
+[%%expect{|
+type 'a s = 'a
+module M : sig type t : immediate end
+module N :
+  sig
+    type ('a, 'b) eq = Refl : ('a, 'a) eq
+    val f : (M.t, M.t s) eq -> int
+  end
+|}]

typing/ctype.ml

@@ -3333,7 +3333,10 @@ let unify1_var env t1 t2 =
     | _ -> assert false
   in
   occur_for Unify env t1 t2;
-  match occur_univar_for Unify env t2 with
+  match
+    occur_univar_for Unify env t2;
+    unification_jkind_check env t2 jkind
+  with
   | () ->
       begin
         try
@@ -3342,7 +3345,6 @@ let unify1_var env t1 t2 =
         with Escape e ->
           raise_for Unify (Escape e)
       end;
-      unification_jkind_check env t2 jkind;
       link_type t1 t2;
       true
   | exception Unify_trace _ when in_pattern_mode () ->
@@ -3364,6 +3366,18 @@ let unify3_var env jkind1 t1' t2 t2' =
         record_equation t1' t2';
       end
 
+(* This is used to check whether we should add a gadt equation refining a
+   Tconstr's jkind during pattern unification. *)
+let constr_jkind_refinable env t jkind =
+  let snap = Btype.snapshot () in
+  let refinable =
+    match unification_jkind_check env t jkind with
+    | () -> false
+    | exception Unify_trace _ -> true
+  in
+  Btype.backtrack snap;
+  refinable
+
 (*
    1. When unifying two non-abbreviated types, one type is made a link
       to the other. When unifying an abbreviated type with a
@@ -3491,6 +3505,23 @@ and unify3 env t1 t1' t2 t2' =
     (Tunivar { jkind = k1 }, Tunivar { jkind = k2 }) ->
       unify_univar_for Unify t1' t2' k1 k2 !univar_pairs;
       link_type t1' t2'
+  (* Before layouts, the following two cases were unnecessary because unifying a
+     [Tconstr] and a [Tvar] couldn't refine the [Tconstr] in any interesting
+     way. *)
+  | (Tconstr (path,[],_), Tvar { jkind })
+    when is_instantiable !env ~for_jkind_eqn:false path
+      && can_generate_equations ()
+      && constr_jkind_refinable !env t1' jkind ->
+      reify env t2';
+      record_equation t1' t2';
+      add_gadt_equation env path t2'
+  | (Tvar { jkind }, Tconstr (path,[],_))
+    when is_instantiable !env ~for_jkind_eqn:false path
+      && can_generate_equations ()
+      && constr_jkind_refinable !env t2' jkind ->
+      reify env t1';
+      record_equation t1' t2';
+      add_gadt_equation env path t1'
   | (Tvar { jkind }, _) ->
       unify3_var env jkind t1' t2 t2'
   | (_, Tvar { jkind }) ->