flambda-backend: Add Jkind.sort to Tpat_array (#1971)

* add Jkind.sort to Tpat_array

* fix chamelon

* update cr comment

* print sort

* move case to minimize merge conflict

* ocaml 5 minus fix

* add assert

Author: alanechang

lambda/matching.ml

@@ -317,7 +317,7 @@ end = struct
       | `Record (fields, closed) ->
           let alpha_field env (lid, l, p) = (lid, l, alpha_pat env p) in
           `Record (List.map (alpha_field env) fields, closed)
-      | `Array (am, ps) -> `Array (am, List.map (alpha_pat env) ps)
+      | `Array (am, arg_sort, ps) -> `Array (am, arg_sort, List.map (alpha_pat env) ps)
       | `Lazy p -> `Lazy (alpha_pat env p)
     in
     { p with pat_desc }
@@ -468,7 +468,7 @@ let matcher discr (p : Simple.pattern) rem =
   | Variant _, (Constant _ | Construct _ | Lazy | Array _ | Record _ | Tuple _)
     ->
       no ()
-  | Array (am1, n1), Array (am2, n2) -> yesif (am1 = am2 && n1 = n2)
+  | Array (am1, _, n1), Array (am2, _, n2) -> yesif (am1 = am2 && n1 = n2)
   | Array _, (Constant _ | Construct _ | Variant _ | Lazy | Record _ | Tuple _)
     ->
       no ()
@@ -2322,19 +2322,19 @@ let divide_record all_labels ~scopes head ctx pm =
 (* Matching against an array pattern *)
 
 let get_key_array = function
-  | { pat_desc = Tpat_array (_, patl) } -> List.length patl
+  | { pat_desc = Tpat_array (_, _, patl) } -> List.length patl
   | _ -> assert false
 
 let get_pat_args_array p rem =
   match p with
-  | { pat_desc = Tpat_array (_, patl) } -> patl @ rem
+  | { pat_desc = Tpat_array (_, _, patl) } -> patl @ rem
   | _ -> assert false
 
 let get_expr_args_array ~scopes kind head (arg, _mut, _sort, _layout) rem =
-  let am, len =
+  let am, arg_sort, len =
     let open Patterns.Head in
     match head.pat_desc with
-    | Array (am, len) -> am, len
+    | Array (am, arg_sort, len) -> am, arg_sort, len
     | _ -> assert false
   in
   let loc = head_loc ~scopes head in
@@ -2353,7 +2353,7 @@ let get_expr_args_array ~scopes kind head (arg, _mut, _sort, _layout) rem =
         (match am with
         | Mutable   -> StrictOpt
         | Immutable -> Alias),
-        Jkind.Sort.for_array_get_result,
+        arg_sort,
         result_layout)
       :: make_args (pos + 1)
   in

typing/parmatch.ml

@@ -143,7 +143,7 @@ let all_coherent column =
     | Tuple l1, Tuple l2 -> l1 = l2
     | Record (lbl1 :: _), Record (lbl2 :: _) ->
       Array.length lbl1.lbl_all = Array.length lbl2.lbl_all
-    | Array (am1, _), Array (am2, _) -> am1 = am2
+    | Array (am1, _, _), Array (am2, _, _) -> am1 = am2
     | Any, _
     | _, Any
     | Record [], Record []
@@ -302,7 +302,7 @@ module Compat
   | Tpat_record (l1,_),Tpat_record (l2,_) ->
       let ps,qs = records_args l1 l2 in
       compats ps qs
-  | Tpat_array (am1, ps), Tpat_array (am2, qs) ->
+  | Tpat_array (am1, _, ps), Tpat_array (am2, _, qs) ->
       am1 = am2 &&
       List.length ps = List.length qs &&
       compats ps qs
@@ -365,7 +365,7 @@ let simple_match d h =
   | Lazy, Lazy -> true
   | Record _, Record _ -> true
   | Tuple len1, Tuple len2 -> len1 = len2
-  | Array (am1, len1), Array (am2, len2) -> am1 = am2 && len1 = len2
+  | Array (am1, _, len1), Array (am2, _, len2) -> am1 = am2 && len1 = len2
   | _, Any -> true
   | _, _ -> false
 
@@ -405,7 +405,7 @@ let simple_match_args discr head args =
       | Variant { has_arg = true }
       | Lazy -> [Patterns.omega]
       | Record lbls ->  omega_list lbls
-      | Array (_, len)
+      | Array (_, _, len)
       | Tuple len -> Patterns.omegas len
       | Variant { has_arg = false }
       | Any
@@ -529,10 +529,10 @@ let do_set_args ~erase_mutable q r = match q with
         make_pat (Tpat_lazy arg) q.pat_type q.pat_env::rest
     | _ -> fatal_error "Parmatch.do_set_args (lazy)"
     end
-| {pat_desc = Tpat_array (am, omegas)} ->
+| {pat_desc = Tpat_array (am, arg_sort, omegas)} ->
     let args,rest = read_args omegas r in
     make_pat
-      (Tpat_array (am, args)) q.pat_type q.pat_env::
+      (Tpat_array (am, arg_sort, args)) q.pat_type q.pat_env::
     rest
 | {pat_desc=Tpat_constant _|Tpat_any} ->
     q::r (* case any is used in matching.ml *)
@@ -1044,17 +1044,17 @@ let build_other ext env =
                     | _ -> assert false)
             (function f -> Tpat_constant(Const_float (string_of_float f)))
             0.0 (fun f -> f +. 1.0) d env
-      | Array (am, _) ->
+      | Array (am, arg_sort, _) ->
           let all_lengths =
             List.map
               (fun (p,_) -> match p.pat_desc with
-              | Array (am', len) when am = am' -> len
+              | Array (am', _, len) when am = am' -> len
               | _ -> assert false)
               env in
           let rec try_arrays l =
             if List.mem l all_lengths then try_arrays (l+1)
             else
-              make_pat (Tpat_array (am, omegas l))
+              make_pat (Tpat_array (am, arg_sort, omegas l))
                 d.pat_type d.pat_env in
           try_arrays 0
       | _ -> Patterns.omega
@@ -1064,7 +1064,7 @@ let rec has_instance p = match p.pat_desc with
   | Tpat_any | Tpat_var _ | Tpat_constant _ | Tpat_variant (_,None,_) -> true
   | Tpat_alias (p,_,_,_,_) | Tpat_variant (_,Some p,_) -> has_instance p
   | Tpat_or (p1,p2,_) -> has_instance p1 || has_instance p2
-  | Tpat_construct (_,_,ps, _) | Tpat_tuple ps | Tpat_array (_, ps) ->
+  | Tpat_construct (_,_,ps, _) | Tpat_tuple ps | Tpat_array (_, _, ps) ->
       has_instances ps
   | Tpat_record (lps,_) -> has_instances (List.map (fun (_,_,x) -> x) lps)
   | Tpat_lazy p
@@ -1712,7 +1712,7 @@ let rec le_pat p q =
   | Tpat_record (l1,_), Tpat_record (l2,_) ->
       let ps,qs = records_args l1 l2 in
       le_pats ps qs
-  | Tpat_array(am1, ps), Tpat_array(am2, qs) ->
+  | Tpat_array(am1, _, ps), Tpat_array(am2, _, qs) ->
       am1 = am2 && List.length ps = List.length qs && le_pats ps qs
 (* In all other cases, enumeration is performed *)
   | _,_  -> not (satisfiable [[p]] [q])
@@ -1766,10 +1766,10 @@ let rec lub p q = match p.pat_desc,q.pat_desc with
     let rs = record_lubs l1 l2 in
     make_pat (Tpat_record (rs, closed))
       p.pat_type p.pat_env
-| Tpat_array (am1, ps), Tpat_array (am2, qs)
+| Tpat_array (am1, arg_sort, ps), Tpat_array (am2, _, qs)
       when am1 = am2 && List.length ps = List.length qs ->
         let rs = lubs ps qs in
-        make_pat (Tpat_array (am1, rs))
+        make_pat (Tpat_array (am1, arg_sort, rs))
           p.pat_type p.pat_env
 | _,_  ->
     raise Empty
@@ -1939,7 +1939,7 @@ let rec collect_paths_from_pat r p = match p.pat_desc with
       (if extendable_path path then add_path path r else r)
       ps
 | Tpat_any|Tpat_var _|Tpat_constant _| Tpat_variant (_,None,_) -> r
-| Tpat_tuple ps | Tpat_array (_, ps)
+| Tpat_tuple ps | Tpat_array (_, _, ps)
 | Tpat_construct (_, {cstr_tag=Extension _}, ps, _)->
     List.fold_left collect_paths_from_pat r ps
 | Tpat_record (lps,_) ->
@@ -2067,7 +2067,7 @@ let inactive ~partial pat =
   | Total -> begin
       let rec loop pat =
         match pat.pat_desc with
-        | Tpat_lazy _ | Tpat_array (Mutable, _) ->
+        | Tpat_lazy _ | Tpat_array (Mutable, _, _) ->
           false
         | Tpat_any | Tpat_var _ | Tpat_variant (_, None, _) ->
             true
@@ -2078,7 +2078,7 @@ let inactive ~partial pat =
             | Const_int32 _ | Const_int64 _ | Const_nativeint _ -> true
           end
         | Tpat_tuple ps | Tpat_construct (_, _, ps, _)
-        | Tpat_array (Immutable, ps) ->
+        | Tpat_array (Immutable, _, ps) ->
             List.for_all (fun p -> loop p) ps
         | Tpat_alias (p,_,_,_,_) | Tpat_variant (_, Some p, _) ->
             loop p

typing/patterns.ml

@@ -58,7 +58,7 @@ module Simple = struct
     | `Variant of label * pattern option * row_desc ref
     | `Record of
         (Longident.t loc * label_description * pattern) list * closed_flag
-    | `Array of mutable_flag * pattern list
+    | `Array of mutable_flag * Jkind.sort * pattern list
     | `Lazy of pattern
   ]
 
@@ -101,7 +101,7 @@ module General = struct
        `Variant (cstr, arg, row_desc)
     | Tpat_record (fields, closed) ->
        `Record (fields, closed)
-    | Tpat_array (am,ps) -> `Array (am, ps)
+    | Tpat_array (am, arg_sort, ps) -> `Array (am, arg_sort, ps)
     | Tpat_or (p, q, row_desc) -> `Or (p, q, row_desc)
     | Tpat_lazy p -> `Lazy p
 
@@ -120,7 +120,7 @@ module General = struct
        Tpat_variant (cstr, arg, row_desc)
     | `Record (fields, closed) ->
        Tpat_record (fields, closed)
-    | `Array (am, ps) -> Tpat_array (am, ps)
+    | `Array (am, arg_sort, ps) -> Tpat_array (am, arg_sort, ps)
     | `Or (p, q, row_desc) -> Tpat_or (p, q, row_desc)
     | `Lazy p -> Tpat_lazy p
 
@@ -147,7 +147,7 @@ module Head : sig
         { tag: label; has_arg: bool;
           cstr_row: row_desc ref;
           type_row : unit -> row_desc; }
-    | Array of mutable_flag * int
+    | Array of mutable_flag * Jkind.sort * int
     | Lazy
 
   type t = desc pattern_data
@@ -174,7 +174,7 @@ end = struct
           type_row : unit -> row_desc; }
           (* the row of the type may evolve if [close_variant] is called,
              hence the (unit -> ...) delay *)
-    | Array of mutable_flag * int
+    | Array of mutable_flag * Jkind.sort * int
     | Lazy
 
   type t = desc pattern_data
@@ -199,8 +199,8 @@ end = struct
             | _ -> assert false
           in
           Variant {tag; has_arg; cstr_row; type_row}, pats
-      | `Array (am, args) ->
-          Array (am, List.length args), args
+      | `Array (am, arg_sort, args) ->
+          Array (am, arg_sort, List.length args), args
       | `Record (largs, _) ->
           let lbls = List.map (fun (_,lbl,_) -> lbl) largs in
           let pats = List.map (fun (_,_,pat) -> pat) largs in
@@ -216,7 +216,7 @@ end = struct
       | Any -> 0
       | Constant _ -> 0
       | Construct c -> c.cstr_arity
-      | Tuple n | Array (_, n) -> n
+      | Tuple n | Array (_, _, n) -> n
       | Record l -> List.length l
       | Variant { has_arg; _ } -> if has_arg then 1 else 0
       | Lazy -> 1
@@ -229,7 +229,7 @@ end = struct
       | Lazy -> Tpat_lazy omega
       | Constant c -> Tpat_constant c
       | Tuple n -> Tpat_tuple (omegas n)
-      | Array (am, n) -> Tpat_array (am, omegas n)
+      | Array (am, arg_sort, n) -> Tpat_array (am, arg_sort, omegas n)
       | Construct c ->
           let lid_loc = mkloc (Longident.Lident c.cstr_name) in
           Tpat_construct (lid_loc, c, omegas c.cstr_arity, None)

typing/patterns.mli

@@ -46,7 +46,7 @@ module Simple : sig
     | `Variant of label * pattern option * row_desc ref
     | `Record of
         (Longident.t loc * label_description * pattern) list * closed_flag
-    | `Array of mutable_flag * pattern list
+    | `Array of mutable_flag * Jkind.sort * pattern list
     | `Lazy of pattern
   ]
   type pattern = view pattern_data
@@ -89,7 +89,7 @@ module Head : sig
           type_row : unit -> row_desc; }
           (* the row of the type may evolve if [close_variant] is called,
              hence the (unit -> ...) delay *)
-    | Array of mutable_flag * int
+    | Array of mutable_flag * Jkind.sort * int
     | Lazy
 
   type t = desc pattern_data

typing/printpat.ml

@@ -94,7 +94,7 @@ let rec pretty_val : type k . _ -> k general_pattern -> _ = fun ppf v ->
           fprintf ppf "@[{%a%t}@]"
             pretty_lvals filtered_lvs elision_mark
       end
-  | Tpat_array (am, vs) ->
+  | Tpat_array (am, _arg_sort, vs) ->
       let punct = match am with
         | Mutable   -> '|'
         | Immutable -> ':'

typing/printtyped.ml

@@ -290,8 +290,9 @@ and pattern : type k . _ -> _ -> k general_pattern -> unit = fun i ppf x ->
   | Tpat_record (l, _c) ->
       line i ppf "Tpat_record\n";
       list i longident_x_pattern ppf l;
-  | Tpat_array (am, l) ->
+  | Tpat_array (am, arg_sort, l) ->
       line i ppf "Tpat_array %a\n" fmt_mutable_flag am;
+      line i ppf "%a\n" Jkind.Sort.format arg_sort;
       list i pattern ppf l;
   | Tpat_lazy p ->
       line i ppf "Tpat_lazy\n";

typing/tast_iterator.ml

@@ -250,7 +250,7 @@ let pat
   | Tpat_variant (_, po, _) -> Option.iter (sub.pat sub) po
   | Tpat_record (l, _) ->
       List.iter (fun (lid, _, i) -> iter_loc sub lid; sub.pat sub i) l
-  | Tpat_array (_, l) -> List.iter (sub.pat sub) l
+  | Tpat_array (_, _, l) -> List.iter (sub.pat sub) l
   | Tpat_alias (p, _, s, _, _) -> sub.pat sub p; iter_loc sub s
   | Tpat_lazy p -> sub.pat sub p
   | Tpat_value p -> sub.pat sub (p :> pattern)

typing/tast_mapper.ml

@@ -307,7 +307,7 @@ let pat
         Tpat_variant (l, Option.map (sub.pat sub) po, rd)
     | Tpat_record (l, closed) ->
         Tpat_record (List.map (tuple3 (map_loc sub) id (sub.pat sub)) l, closed)
-    | Tpat_array (am, l) -> Tpat_array (am, List.map (sub.pat sub) l)
+    | Tpat_array (am, arg_sort, l) -> Tpat_array (am, arg_sort, List.map (sub.pat sub) l)
     | Tpat_alias (p, id, s, uid, m) ->
         Tpat_alias (sub.pat sub p, id, map_loc sub s, uid, m)
     | Tpat_lazy p -> Tpat_lazy (sub.pat sub p)

typing/typecore.ml

@@ -1485,12 +1485,16 @@ let solve_Ppat_array ~refine loc env mutability expected_ty =
     | Immutable -> Predef.type_iarray
     | Mutable -> Predef.type_array
   in
-  (* CR layouts v4: in the future we'll have arrays of other jkinds *)
-  let ty_elt = newgenvar (Jkind.value ~why:Array_element) in
+  (* CR layouts v4: The code below is written this way to make it easier to update
+     when we generalize array to contain non-value jkinds. When that happens,
+     change the next two lines to use [Jkind.of_new_sort_var]. *)
+  let jkind = Jkind.value ~why:Array_element in
+  let arg_sort = Jkind.sort_of_jkind jkind in
+  let ty_elt = newgenvar jkind in
   let expected_ty = generic_instance expected_ty in
   unify_pat_types ~refine
     loc env (type_some_array ty_elt) expected_ty;
-  ty_elt
+  ty_elt, arg_sort
 
 let solve_Ppat_lazy  ~refine loc env expected_ty =
   let nv = newgenvar (Jkind.value ~why:Lazy_expression) in
@@ -2279,15 +2283,15 @@ and type_pat_aux
        keep them in sync, at the cost of a worse diff with upstream; it
        shouldn't be too bad.  We can inline this when we upstream this code and
        combine the two array pattern constructors. *)
-    let ty_elt = solve_Ppat_array ~refine loc env mutability expected_ty in
+    let ty_elt, arg_sort = solve_Ppat_array ~refine loc env mutability expected_ty in
     let alloc_mode =
       match mutability with
       | Mutable -> simple_pat_mode Value.legacy
       | Immutable -> alloc_mode
     in
     let pl = List.map (fun p -> type_pat ~alloc_mode tps Value p ty_elt) spl in
     rvp {
-      pat_desc = Tpat_array (mutability, pl);
+      pat_desc = Tpat_array (mutability, arg_sort, pl);
       pat_loc = loc; pat_extra=[];
       pat_type = instance expected_ty;
       pat_attributes;
@@ -3061,10 +3065,11 @@ let rec check_counter_example_pat
       in
       map_fold_cont type_label_pat fields
         (fun fields -> mkp k (Tpat_record (fields, closed)))
-  | Tpat_array (mut, tpl) ->
-      let ty_elt = solve_Ppat_array ~refine loc env mut expected_ty in
+  | Tpat_array (mut, original_arg_sort, tpl) ->
+      let ty_elt, arg_sort = solve_Ppat_array ~refine loc env mut expected_ty in
+      assert (Jkind.Sort.equate original_arg_sort arg_sort);
       map_fold_cont (fun p -> check_rec p ty_elt) tpl
-        (fun pl -> mkp k (Tpat_array (mut, pl)))
+        (fun pl -> mkp k (Tpat_array (mut, arg_sort, pl)))
   | Tpat_or(tp1, tp2, _) ->
       (* We are in counter-example mode, but try to avoid backtracking *)
       let must_split =

typing/typedtree.ml

@@ -79,7 +79,7 @@ and 'k pattern_desc =
         closed_flag ->
       value pattern_desc
   | Tpat_array :
-      mutable_flag * value general_pattern list -> value pattern_desc
+      mutable_flag * Jkind.sort * value general_pattern list -> value pattern_desc
   | Tpat_lazy : value general_pattern -> value pattern_desc
   (* computation patterns *)
   | Tpat_value : tpat_value_argument -> computation pattern_desc
@@ -782,7 +782,7 @@ let shallow_iter_pattern_desc
   | Tpat_variant(_, pat, _) -> Option.iter f.f pat
   | Tpat_record (lbl_pat_list, _) ->
       List.iter (fun (_, _, pat) -> f.f pat) lbl_pat_list
-  | Tpat_array (_, patl) -> List.iter f.f patl
+  | Tpat_array (_, _, patl) -> List.iter f.f patl
   | Tpat_lazy p -> f.f p
   | Tpat_any
   | Tpat_var _
@@ -804,8 +804,8 @@ let shallow_map_pattern_desc
       Tpat_record (List.map (fun (lid, l,p) -> lid, l, f.f p) lpats, closed)
   | Tpat_construct (lid, c, pats, ty) ->
       Tpat_construct (lid, c, List.map f.f pats, ty)
-  | Tpat_array (am, pats) ->
-      Tpat_array (am, List.map f.f pats)
+  | Tpat_array (am, arg_sort, pats) ->
+      Tpat_array (am, arg_sort, List.map f.f pats)
   | Tpat_lazy p1 -> Tpat_lazy (f.f p1)
   | Tpat_variant (x1, Some p1, x2) ->
       Tpat_variant (x1, Some (f.f p1), x2)
@@ -914,7 +914,7 @@ let iter_pattern_full ~both_sides_of_or f sort pat =
       | Tpat_tuple patl -> List.iter (loop f Jkind.Sort.value) patl
         (* CR layouts v5: tuple case to change when we allow non-values in
            tuples *)
-      | Tpat_array (_, patl) -> List.iter (loop f Jkind.Sort.value) patl
+      | Tpat_array (_, arg_sort, patl) -> List.iter (loop f arg_sort) patl
       | Tpat_lazy p | Tpat_exception p -> loop f Jkind.Sort.value p
       (* Cases without variables: *)
       | Tpat_any | Tpat_constant _ -> ()