Reroute legacy immediate support for jkinds

ocaml/parsing/builtin_attributes.mli

@@ -175,6 +175,8 @@ val has_unique: Parsetree.attributes -> (bool,unit) result
 
 val has_once : Parsetree.attributes -> (bool, unit) result
 
+(* CR layouts v1.5: Remove everything except for [Immediate64] and [Immediate]
+   after rerouting [@@immediate]. *)
 type jkind_attribute =
   | Immediate64
   | Immediate

ocaml/parsing/jane_asttypes.ml

@@ -2,7 +2,7 @@
 (*                                                                        *)
 (*                                 OCaml                                  *)
 (*                                                                        *)
-(*               Antal Spector-Zabusky, Jane Street, New York             *)
+(*                   Nick Roberts, Jane Street, New York                  *)
 (*                                                                        *)
 (*   Copyright 2023 Jane Street Group LLC                                 *)
 (*                                                                        *)

ocaml/parsing/jane_asttypes.mli

@@ -34,7 +34,7 @@ type global_flag =
 (** [const_jkind] is private to limit confusion with type variables, which
     are also strings in the parser.
 *)
-type const_jkind = private string
+type const_jkind
 
 val jkind_of_string : string -> const_jkind
 

ocaml/typing/jkind.ml

@@ -438,7 +438,16 @@ module Const : sig
     | Immediate
     | Float64
 
+<<<<<<< HEAD
   val of_attribute : Builtin_attributes.jkind_attribute -> t
+=======
+  (** The function names are suffixed with "unchecked" to suggest that
+      they don't check whether the layouts extension is enabled.
+  *)
+
+  val of_user_written_attribute_unchecked :
+    Builtin_attributes.jkind_attribute -> t
+>>>>>>> nroberts/layouts-are-strings-in-parsetree
 
   val of_user_written_annotation_unchecked :
     Jane_asttypes.const_jkind -> t option
@@ -486,13 +495,10 @@ type const = Const.t =
   | Immediate
   | Float64
 
-let string_of_const : const -> _ = function
-  | Any -> "any"
-  | Value -> "value"
-  | Void -> "void"
-  | Immediate64 -> "immediate64"
-  | Immediate -> "immediate"
-  | Float64 -> "float64"
+type annotation = const * Jane_asttypes.jkind_annotation
+
+let string_of_const const =
+  Jane_asttypes.jkind_to_string (Const.to_user_written_annotation const)
 
 let equal_const (c1 : const) (c2 : const) =
   match c1, c2 with
@@ -567,6 +573,7 @@ let of_const ~why : const -> t = function
   | Void -> fresh_jkind (Sort Sort.void) ~why
   | Float64 -> fresh_jkind (Sort Sort.float64) ~why
 
+<<<<<<< HEAD
 let const_of_user_written_annotation ~context ~is_type_decl
     Location.{ loc; txt = annot } =
   match Const.of_user_written_annotation_unchecked annot with
@@ -582,13 +589,44 @@ let const_of_user_written_annotation ~context ~is_type_decl
 
 let const_to_user_written_annotation = Const.to_user_written_annotation
 
+=======
+let check_extension_for_const ?(legacy_immediate = false) ~context ~loc annot =
+  match annot with
+  | (Immediate | Immediate64 | Value) as const when legacy_immediate -> const
+  | const ->
+    let required_layouts_level = get_required_layouts_level context const in
+    if not (Language_extension.is_at_least Layouts required_layouts_level)
+    then raise ~loc (Insufficient_level (context, const));
+    const
+
+let const_of_user_written_annotation ?legacy_immediate ~context
+    Location.{ loc; txt = annot } =
+  match Const.of_user_written_annotation_unchecked annot with
+  | None -> raise ~loc (Unknown_jkind annot)
+  | Some unchecked ->
+    check_extension_for_const ?legacy_immediate ~context ~loc unchecked
+
+let const_of_user_written_attribute ?legacy_immediate ~context
+    Location.{ loc; txt = attribute } =
+  let unchecked = Const.of_user_written_attribute_unchecked attribute in
+  let checked =
+    check_extension_for_const ?legacy_immediate ~context ~loc unchecked
+  in
+  Location.{ loc; txt = checked }
+
+let const_of_attributes ~legacy_immediate ~context attrs =
+  Builtin_attributes.jkind ~legacy_immediate attrs
+  |> Result.map
+       (Option.map (const_of_user_written_attribute ~legacy_immediate ~context))
+
+>>>>>>> nroberts/layouts-are-strings-in-parsetree
 let of_annotated_const ~context Location.{ txt = const; loc = const_loc } =
   of_const ~why:(Annotated (context, const_loc)) const
 
 let of_annotation ~context ~is_type_decl (annot : _ Location.loc) =
   let const = const_of_user_written_annotation ~is_type_decl ~context annot in
   let jkind = of_annotated_const { txt = const; loc = annot.loc } ~context in
-  jkind, const
+  jkind, (const, annot)
 
 let of_annotation_option_default ~default ~context ~is_type_decl =
   Option.fold ~none:(default, None) ~some:(fun annot ->
@@ -1345,7 +1383,7 @@ end
 let report_error ~loc = function
   | Unknown_jkind jkind ->
     Location.errorf ~loc
-      (* CR layouts errors: use the context to produce a better error message.
+      (* CR layouts v2.9: use the context to produce a better error message.
          When RAE tried this, some types got printed like [t/2], but the
          [/2] shouldn't be there. Investigate and fix. *)
       "@[<v>Unknown layout %a@]" Jane_syntax.Layouts.Pprint.const_jkind jkind

ocaml/typing/jkind.mli

@@ -327,12 +327,15 @@ type const =
   | Immediate
   | Float64
 
-type 'a loc := 'a Location.loc
-
 val const_of_user_written_annotation :
+<<<<<<< HEAD
   context:annotation_context -> Jane_asttypes.const_jkind loc -> const
-
-val const_to_user_written_annotation : const -> Jane_asttypes.const_jkind
+=======
+  ?legacy_immediate:bool ->
+  context:annotation_context ->
+  Jane_asttypes.jkind_annotation ->
+  const
+>>>>>>> nroberts/layouts-are-strings-in-parsetree
 
 val string_of_const : const -> string
 
@@ -376,14 +379,29 @@ val of_sort_for_error : why:concrete_jkind_reason -> sort -> t
 
 val of_const : why:creation_reason -> const -> t
 
+<<<<<<< HEAD
 val of_annotation :
   context:annotation_context -> Jane_asttypes.jkind_annotation -> t * const
+=======
+(** The typed jkind together with its user-written annotation. *)
+type annotation = const * Jane_asttypes.jkind_annotation
+
+(* CR layouts v1.5: remove legacy_immediate when the old attributes mechanism
+   is rerouted away from the new annotations mechanism *)
+val of_annotation :
+  ?legacy_immediate:bool ->
+  context:annotation_context ->
+  Jane_asttypes.jkind_annotation ->
+  t * annotation
+>>>>>>> nroberts/layouts-are-strings-in-parsetree
 
 val of_annotation_option_default :
   default:t ->
   context:annotation_context ->
   Jane_asttypes.jkind_annotation option ->
-  t * const option
+  t * annotation option
+
+(* CR layouts v1.5: remove [of_attributes] when we reroute [@@immediate]. *)
 
 (** Find a jkind from a type declaration. Type declarations are special because
     the jkind may have been provided via [: jkind] syntax (which goes through

ocaml/typing/oprint.ml

@@ -272,7 +272,7 @@ let pr_present =
 let pr_var = Printast.tyvar
 
 let print_out_jkind ppf = function
-  | Olay_const jkind -> fprintf ppf "%s" (Jane_asttypes.jkind_to_string jkind)
+  | Olay_const jkind -> fprintf ppf "%s" jkind
   | Olay_var v     -> fprintf ppf "%s" v
 
 let print_out_jkind_annot ppf = function
@@ -377,7 +377,7 @@ let mode_agree expected real =
   linearity_agree expected.oam_linearity real.oam_linearity
 
 let print_out_jkind ppf = function
-  | Olay_const jkind -> fprintf ppf "%s" (Jane_asttypes.jkind_to_string jkind)
+  | Olay_const jkind -> fprintf ppf "%s" jkind
   | Olay_var v     -> fprintf ppf "%s" v
 
 let is_local mode =

ocaml/typing/outcometree.mli

@@ -57,7 +57,11 @@ type out_value =
   | Oval_variant of string * out_value option
 
 type out_jkind =
-  | Olay_const of Jane_asttypes.const_jkind
+  (* CR layouts v1.5: It would be a bit better to store [Jkind.const] here, and
+     indeed I will do that in the stacked PR after deleting the one site where
+     we need to populate a string.
+  *)
+  | Olay_const of string
   | Olay_var of string
 
 type out_type_param =

ocaml/typing/printtyp.ml

@@ -1184,7 +1184,7 @@ let print_labels = ref true
 let out_jkind_option_of_jkind jkind =
   match Jkind.get jkind with
   | Const Value -> None
-  | Const jkind -> Some (Olay_const (Jkind.const_to_user_written_annotation jkind))
+  | Const jkind -> Some (Olay_const (Jkind.string_of_const jkind))
   | Var v -> (* This handles (X1). *)
     if !Clflags.verbose_types
     then Some (Olay_var (Jkind.Sort.var_name v))
@@ -1650,8 +1650,20 @@ let rec tree_of_type_decl id decl =
       otype_private = priv;
       otype_jkind =
         Option.map
+<<<<<<< HEAD
           (fun x -> Olay_const (Jkind.const_to_user_written_annotation x))
           jkind_annotation;
+=======
+          (fun { txt } ->
+             let jkind_attribute =
+               Builtin_attributes.jkind_attribute_to_string txt
+             in
+             (* CR layouts 1.5: This is a bit of a lie: we're interpreting the
+                jkind attribute as a jkind *annotation*. This will go away in a
+                child PR when we move jkind annotations into Jane Syntax. *)
+             Olay_const jkind_attribute)
+          lay;
+>>>>>>> nroberts/layouts-are-strings-in-parsetree
       otype_unboxed = unboxed;
       otype_cstrs = constraints }
 
@@ -2267,7 +2279,7 @@ let trees_of_type_expansion'
       match get_desc ty with
       | Tvar { jkind; _ } | Tunivar { jkind; _ } ->
           let olay = match Jkind.get jkind with
-            | Const clay -> Olay_const (Jkind.const_to_user_written_annotation clay)
+            | Const clay -> Olay_const (Jkind.string_of_const clay)
             | Var v      -> Olay_var (Jkind.Sort.var_name v)
           in
           Otyp_jkind_annot (out, olay)

ocaml/typing/printtyped.ml

@@ -163,11 +163,10 @@ let typevar_jkind ~print_quote ppf (v, l) =
   in
   match l with
   | None -> fprintf ppf " %a" pptv v
-  | Some lay ->
-      let lay = Jkind.const_to_user_written_annotation lay in
+  | Some (_, lay) ->
       fprintf ppf " (%a : %a)"
         pptv v
-        Jane_syntax.Layouts.Pprint.const_jkind lay
+        Jane_syntax.Layouts.Pprint.const_jkind lay.txt
 
 let typevars ppf vs =
   List.iter (typevar_jkind ~print_quote:true ppf) vs
@@ -211,7 +210,7 @@ let attributes i ppf l =
     Printast.payload (i + 1) ppf a.Parsetree.attr_payload
   ) l
 
-let jkind_annotation i ppf jkind =
+let jkind_annotation i ppf (jkind, _) =
   line i ppf "%s" (Jkind.string_of_const jkind)
 
 let rec core_type i ppf x =

ocaml/typing/tast_iterator.ml

@@ -460,7 +460,7 @@ let typ sub {ctyp_desc; ctyp_env; _} =
   sub.env sub ctyp_env;
   match ctyp_desc with
   | Ttyp_var (_, jkind) ->
-      Option.iter (sub.jkind_annotation sub) jkind
+      Option.iter (fun (jkind, _) -> sub.jkind_annotation sub jkind) jkind
   | Ttyp_arrow (_, ct1, ct2) ->
       sub.typ sub ct1;
       sub.typ sub ct2
@@ -470,10 +470,10 @@ let typ sub {ctyp_desc; ctyp_env; _} =
   | Ttyp_class (_, _, list) -> List.iter (sub.typ sub) list
   | Ttyp_alias (ct, _, jkind) ->
     sub.typ sub ct;
-    Option.iter (sub.jkind_annotation sub) jkind
+    Option.iter (fun (jkind, _) -> sub.jkind_annotation sub jkind) jkind
   | Ttyp_variant (list, _, _) -> List.iter (sub.row_field sub) list
   | Ttyp_poly (vars, ct) ->
-      List.iter (fun (_, l) -> Option.iter (sub.jkind_annotation sub) l) vars;
+      List.iter (fun (_, l) -> Option.iter (fun (j, _) -> sub.jkind_annotation sub j) l) vars;
       sub.typ sub ct
   | Ttyp_package pack -> sub.package_type sub pack
 

ocaml/typing/tast_mapper.ml

@@ -37,7 +37,7 @@ type mapper =
     expr: mapper -> expression -> expression;
     extension_constructor: mapper -> extension_constructor ->
       extension_constructor;
-    jkind_annotation: mapper -> Jkind.const -> Jkind.const;
+    jkind_annotation: mapper -> Jkind.annotation -> Jkind.annotation;
     module_binding: mapper -> module_binding -> module_binding;
     module_coercion: mapper -> module_coercion -> module_coercion;
     module_declaration: mapper -> module_declaration -> module_declaration;

ocaml/typing/tast_mapper.mli

@@ -36,7 +36,7 @@ type mapper =
     expr: mapper -> expression -> expression;
     extension_constructor: mapper -> extension_constructor ->
       extension_constructor;
-    jkind_annotation: mapper -> Jkind.const -> Jkind.const;
+    jkind_annotation: mapper -> Jkind.annotation -> Jkind.annotation;
     module_binding: mapper -> module_binding -> module_binding;
     module_coercion: mapper -> module_coercion -> module_coercion;
     module_declaration: mapper -> module_declaration -> module_declaration;

ocaml/typing/typedecl.ml

@@ -402,19 +402,6 @@ let transl_constructor_arguments env univars closed = function
       Types.Cstr_record lbls',
       Cstr_record lbls
 
-let transl_constructor_type_parameters (svars : _ Either.t) ~cstr_path =
-  match svars with
-  | Left vars_only -> List.map (fun v -> v.txt, None) vars_only
-  | Right vars_jkinds ->
-      List.map
-        (fun (v, l) ->
-           v.txt,
-           Option.map
-             (Jkind.const_of_user_written_annotation
-                ~context:(Constructor_type_parameter (cstr_path, v.txt)))
-             l)
-        vars_jkinds
-
 (* Note that [make_constructor] does not fill in the [ld_jkind] field of any
    computed record types, because it's called too early in the translation of a
    type declaration to compute accurate jkinds in the presence of recursively
@@ -423,12 +410,28 @@ let transl_constructor_type_parameters (svars : _ Either.t) ~cstr_path =
 let make_constructor
       env loc ~cstr_path ~type_path type_params (svars : _ Either.t)
       sargs sret_type =
+  let tvars = match svars with
+    | Left vars_only -> List.map (fun v -> v.txt, None) vars_only
+    | Right vars_jkinds ->
+        List.map
+          (fun (v, l) ->
+            v.txt,
+            Option.map
+              (fun annot ->
+                 let const =
+                    Jkind.const_of_user_written_annotation
+                      ~context:(Constructor_type_parameter (cstr_path, v.txt))
+                      annot
+                 in
+                 const, annot)
+              l)
+          vars_jkinds
+  in
   match sret_type with
   | None ->
       let args, targs =
         transl_constructor_arguments env None true sargs
       in
-      let tvars = transl_constructor_type_parameters svars ~cstr_path in
         tvars, targs, None, args, None
   | Some sret_type -> TyVarEnv.with_local_scope begin fun () ->
       (* if it's a generalized constructor we must work in a narrowed
@@ -446,11 +449,6 @@ let make_constructor
                    ~context:(fun v -> Constructor_type_parameter (cstr_path, v))
                    vars_jkinds), true
       in
-      let tvars =
-        match univars with
-        | None -> transl_constructor_type_parameters svars ~cstr_path
-        | Some univars -> TyVarEnv.ttyp_poly_arg univars
-      in
       let args, targs =
         transl_constructor_arguments env univars closed sargs
       in