OxCaml: Support for kind annotations

CHANGES.md

@@ -5,6 +5,7 @@
 - `markdown-generate` command now accepts multiple `.odocl` files in a single
   invocation, eliminating the need for shell scripting (@davesnx, #1387)
 - Support for OxCaml (@lukemaurer, @art-w, #1399)
+- Support for OxCaml kind annotations (@art-w, #1410)
 
 ### Fixed
 - Fix compile-time crashing bugs #930 and #1385 (@jonludlam, #1400)

src/document/generator.ml

@@ -306,6 +306,10 @@ module Make (Syntax : SYNTAX) = struct
 
     val format_type_path :
       delim:[ `parens | `brackets ] -> Lang.TypeExpr.t list -> text -> text
+
+    val kind_annotation : Odoc_model.Lang.KindAnnotation.t -> text
+
+    val with_kind_annotation : Odoc_model.Lang.KindAnnotation.t -> text -> text
   end = struct
     let rec te_variant (t : Odoc_model.Lang.TypeExpr.Polymorphic_variant.t) =
       let style_arguments ~constant arguments =
@@ -426,6 +430,30 @@ module Make (Syntax : SYNTAX) = struct
       then enclose ~l:lparen res ~r:")"
       else res
 
+    and kind_annotation (k : Odoc_model.Lang.KindAnnotation.t) =
+      match k with
+      | Default -> O.noop
+      | Abbreviation s -> O.txt s
+      | Mod (base, modes) ->
+          kind_annotation base ++ O.txt " " ++ O.keyword "mod"
+          ++ O.txt (" " ^ String.concat ~sep:" " modes)
+      | With (base, ty, modalities) -> (
+          kind_annotation base ++ O.txt " " ++ O.keyword "with" ++ O.txt " "
+          ++ type_expr ty
+          ++
+          match modalities with
+          | [] -> O.noop
+          | mods ->
+              O.txt " " ++ O.keyword "mod"
+              ++ O.txt (" " ^ String.concat ~sep:" " mods))
+      | Kind_of ty -> O.keyword "kind_of_" ++ O.txt " " ++ type_expr ty
+      | Product ks -> O.list ks ~sep:(O.txt " & ") ~f:kind_annotation
+
+    and with_kind_annotation kind base =
+      match kind with
+      | Odoc_model.Lang.KindAnnotation.Default -> base
+      | k -> O.txt "(" ++ base ++ O.txt " : " ++ kind_annotation k ++ O.txt ")"
+
     and type_expr ?(needs_parentheses = false) (t : Odoc_model.Lang.TypeExpr.t)
         =
       let enclose_parens_if_needed res =
@@ -478,9 +506,11 @@ module Make (Syntax : SYNTAX) = struct
           format_type_path ~delim:`brackets args
             (Link.from_path (path :> Paths.Path.t))
       | Poly (polyvars, t) ->
-          enclose_parens_if_needed
-          @@ O.txt ("'" ^ String.concat ~sep:" '" polyvars ^ ". ")
-             ++ type_expr t
+          let format_poly_var (name, kind) =
+            with_kind_annotation kind (O.txt ("'" ^ name))
+          in
+          let vars = O.list polyvars ~sep:(O.txt " ") ~f:format_poly_var in
+          enclose_parens_if_needed @@ (vars ++ O.txt ". " ++ type_expr t)
       | Quote t -> O.span (O.txt "<[ " ++ O.box_hv (type_expr t) ++ O.txt " ]>")
       | Splice t -> O.span (O.txt "$" ++ type_expr ~needs_parentheses:true t)
       | Package pkg ->
@@ -821,8 +851,8 @@ module Make (Syntax : SYNTAX) = struct
         Odoc_model.Lang.TypeDecl.param list ->
         text =
      fun ?(delim = `parens) params ->
-      let format_param { Odoc_model.Lang.TypeDecl.desc; variance; injectivity }
-          =
+      let format_param_str
+          { Odoc_model.Lang.TypeDecl.desc; variance; injectivity; kind = _ } =
         let desc =
           match desc with
           | Odoc_model.Lang.TypeDecl.Any -> [ "_" ]
@@ -838,15 +868,22 @@ module Make (Syntax : SYNTAX) = struct
         let final = if injectivity then "!" :: var_desc else var_desc in
         String.concat ~sep:"" final
       in
-      O.txt
-        (match params with
-        | [] -> ""
-        | [ x ] -> format_param x |> Syntax.Type.handle_format_params
-        | lst -> (
-            let params = String.concat ~sep:", " (List.map format_param lst) in
-            (match delim with `parens -> "(" | `brackets -> "[")
-            ^ params
-            ^ match delim with `parens -> ")" | `brackets -> "]"))
+      let format_param p =
+        Type_expression.with_kind_annotation p.Odoc_model.Lang.TypeDecl.kind
+          (O.txt (format_param_str p))
+      in
+      match params with
+      | [] -> O.noop
+      | [ x ] ->
+          let base = format_param_str x |> Syntax.Type.handle_format_params in
+          Type_expression.with_kind_annotation x.kind (O.txt base)
+      | lst ->
+          let left, right =
+            match delim with `parens -> ("(", ")") | `brackets -> ("[", "]")
+          in
+          O.txt left
+          ++ O.list lst ~sep:(O.txt ", ") ~f:format_param
+          ++ O.txt right
 
     let format_constraints constraints =
       O.list constraints ~f:(fun (t1, t2) ->
@@ -896,7 +933,12 @@ module Make (Syntax : SYNTAX) = struct
             let params = format_params l in
             Syntax.Type.handle_constructor_params (O.txt tyname) params
       in
-      let intro = keyword' ++ O.txt " " ++ tconstr in
+      let kind_annot =
+        match t.equation.kind with
+        | Default -> O.noop
+        | k -> O.txt " : " ++ Type_expression.kind_annotation k
+      in
+      let intro = keyword' ++ O.txt " " ++ tconstr ++ kind_annot in
       let constraints = format_constraints t.equation.constraints in
       let manifest, need_private, long_prefix =
         match t.equation.manifest with

src/loader/cmi.ml

@@ -490,6 +490,51 @@ let mark_class_declaration cld =
   List.iter mark_type_parameter cld.cty_params;
   mark_class_type cld.cty_params cld.cty_type
 
+#if defined OXCAML
+let read_parsetree_core_type (ct : Parsetree.core_type) =
+  let open TypeExpr in
+  match ct.ptyp_desc with
+  | Ptyp_var (s, _) -> Var s
+  | Ptyp_any _ -> Any
+  | _ -> failwith "invalid core type"
+
+let rec read_jkind_annotation (jk : Parsetree.jkind_annotation) =
+  let open KindAnnotation in
+  match jk.pjkind_desc with
+  | Pjk_default -> Default
+  | Pjk_abbreviation s -> Abbreviation s
+  | Pjk_mod (jk', modes) ->
+    let modes = List.map (fun (m : Parsetree.mode Location.loc) ->
+      let (Parsetree.Mode s) = m.txt in s) modes in
+    Mod (read_jkind_annotation jk', modes)
+  | Pjk_with (jk', cty, modalities) ->
+    let ty = read_parsetree_core_type cty in
+    let modalities = List.map (fun (m : Parsetree.modality Location.loc) ->
+      let (Parsetree.Modality s) = m.txt in s) modalities in
+    With (read_jkind_annotation jk', ty, modalities)
+  | Pjk_kind_of cty ->
+    Kind_of (read_parsetree_core_type cty)
+  | Pjk_product jks ->
+    Product (List.map read_jkind_annotation jks)
+
+let read_jkind_annotation = function
+  | None -> KindAnnotation.Default
+  | Some jk ->
+    match read_jkind_annotation jk with
+    | Abbreviation "value" -> Default
+    | k -> k
+
+let jkind_of_type_desc te =
+  match  te with
+  | Tvar { jkind; _ } | Tunivar { jkind; _ } ->
+      read_jkind_annotation jkind.annotation
+  | _ -> KindAnnotation.Default
+#else
+
+let jkind_of_type_desc _te = KindAnnotation.Default
+
+#endif
+
 let rec read_type_expr env typ =
   let open TypeExpr in
   let px = proxy typ in
@@ -549,10 +594,14 @@ let rec read_type_expr env typ =
       | Tpoly (typ, []) -> read_type_expr env typ
       | Tpoly (typ, tyl) ->
           let tyl = List.map Compat.repr tyl in
-          let vars = List.map name_of_type_repr tyl in
+          let vars_with_kinds = List.map (fun ty ->
+            let name = name_of_type_repr ty in
+            let kind = jkind_of_type_desc ty.desc in
+            (name, kind)
+          ) tyl in
           let typ = read_type_expr env typ in
             remove_names tyl;
-            Poly(vars, typ)
+            Poly(vars_with_kinds, typ)
       | Tunivar _ -> Var (name_of_type typ)
 #if OCAML_VERSION>=(5,4,0)
       | Tpackage {pack_path=p; pack_cstrs } ->
@@ -827,6 +876,7 @@ let read_type_parameter abstr var param =
     if name = "_" then Any
     else Var name
   in
+  let kind = jkind_of_type_desc (Compat.get_desc param) in
   let variance =
     if not (abstr || aliasable param) then None
     else begin
@@ -836,7 +886,7 @@ let read_type_parameter abstr var param =
         else None
       end in
   let injectivity = read_injectivity var in
-  {desc; variance; injectivity}
+  {desc; variance; injectivity; kind}
 
 let read_type_constraints env params =
   List.fold_right
@@ -898,7 +948,14 @@ let read_type_declaration env parent id decl =
     List.map2 (read_type_parameter abstr) decl.type_variance params
   in
   let private_ = (decl.type_private = Private) in
-  let equation = Equation.{params; manifest; constraints; private_} in
+  let kind =
+#if defined OXCAML
+    read_jkind_annotation decl.type_jkind.annotation
+#else
+    KindAnnotation.Default
+#endif
+  in
+  let equation = Equation.{params; manifest; constraints; private_; kind} in
   {id; source_loc; doc; canonical; equation; representation}
 
 let read_extension_constructor env parent id ext =

src/loader/cmi.mli

@@ -96,3 +96,9 @@ val read_extension_constructor : env ->
 val read_exception : env ->
   Paths.Identifier.Signature.t -> Ident.t ->
   Types.extension_constructor -> Odoc_model.Lang.Exception.t
+
+#if defined OXCAML
+val read_jkind_annotation :
+  Parsetree.jkind_annotation option ->
+  Odoc_model.Lang.KindAnnotation.t
+#endif

src/loader/cmti.ml

@@ -168,10 +168,11 @@ let rec read_core_type env container ctyp =
     | Ttyp_poly([], typ) -> read_core_type env container typ
 #if defined OXCAML
     | Ttyp_poly(vars, typ) ->
-      (* TODO: presumably want the layouts, eventually *)
-      Poly(List.map fst vars, read_core_type env container typ)
+      Poly(List.map (fun (name, jk) ->
+        (name, Cmi.read_jkind_annotation jk)
+      ) vars, read_core_type env container typ)
 #else
-    | Ttyp_poly(vars, typ) -> Poly(vars, read_core_type env container typ)
+    | Ttyp_poly(vars, typ) -> Poly(List.map (fun v -> (v, KindAnnotation.Default)) vars, read_core_type env container typ)
 #endif
 #if OCAML_VERSION >= (5,4,0)
     | Ttyp_package {tpt_path = pack_path; tpt_cstrs=pack_fields; _} ->
@@ -219,15 +220,14 @@ let read_value_description env parent vd =
 
 let read_type_parameter (ctyp, var_and_injectivity)  =
   let open TypeDecl in
-  let desc =
+  let desc, kind =
     match ctyp.ctyp_desc with
 #if defined OXCAML
-    (* TODO: presumably we want the layouts below, eventually *)
-    | Ttyp_var (None, _layout) -> Any
-    | Ttyp_var (Some s, _layout) -> Var s
+    | Ttyp_var (None, layout) -> Any, Cmi.read_jkind_annotation layout
+    | Ttyp_var (Some s, layout) -> Var s, Cmi.read_jkind_annotation layout
 #else
-    | Ttyp_any -> Any
-    | Ttyp_var s -> Var s
+    | Ttyp_any -> Any, KindAnnotation.Default
+    | Ttyp_var s -> Var s, KindAnnotation.Default
 #endif
     | _ -> assert false
   in
@@ -254,7 +254,7 @@ let read_type_parameter (ctyp, var_and_injectivity)  =
     var, injectivity
 #endif
   in
-    {desc; variance; injectivity}
+    {desc; variance; injectivity; kind}
 
 #if defined OXCAML
 let is_mutable = Types.is_mutable
@@ -346,7 +346,14 @@ let read_type_equation env container decl =
           read_core_type env container typ2))
       decl.typ_cstrs
   in
-    {params; private_; manifest; constraints}
+  let kind =
+#if defined OXCAML
+    Cmi.read_jkind_annotation decl.typ_jkind_annotation
+#else
+    KindAnnotation.Default
+#endif
+  in
+    {params; private_; manifest; constraints; kind}
 
 let read_type_declaration env parent decl =
   let open TypeDecl in

src/model/lang.ml

@@ -258,6 +258,7 @@ and TypeDecl : sig
     desc : param_desc;
     variance : variance option;
     injectivity : bool;
+    kind : KindAnnotation.t;
   }
 
   module Equation : sig
@@ -266,6 +267,7 @@ and TypeDecl : sig
       private_ : bool;
       manifest : TypeExpr.t option;
       constraints : (TypeExpr.t * TypeExpr.t) list;
+      kind : KindAnnotation.t;
     }
   end
 
@@ -417,6 +419,19 @@ and InstanceVariable : sig
 end =
   InstanceVariable
 
+(** {3 Kind annotations} *)
+
+and KindAnnotation : sig
+  type t =
+    | Default
+    | Abbreviation of string
+    | Mod of t * string list
+    | With of t * TypeExpr.t * string list
+    | Kind_of of TypeExpr.t
+    | Product of t list
+end =
+  KindAnnotation
+
 (** {3 Type expressions} *)
 
 and TypeExpr : sig
@@ -464,7 +479,7 @@ and TypeExpr : sig
     | Polymorphic_variant of TypeExpr.Polymorphic_variant.t
     | Object of TypeExpr.Object.t
     | Class of Path.ClassType.t * t list
-    | Poly of string list * t
+    | Poly of (string * KindAnnotation.t) list * t
     | Quote of t
     | Splice of t
     | Package of TypeExpr.Package.t

src/model_desc/lang_desc.ml

@@ -367,6 +367,21 @@ and typedecl_variance =
   Variant
     (function Pos -> C0 "Pos" | Neg -> C0 "Neg" | Bivariant -> C0 "Bivariant")
 
+and kind_annotation =
+  let open Lang.KindAnnotation in
+  Variant
+    (function
+    | Default -> C0 "Default"
+    | Abbreviation x -> C ("Abbreviation", x, string)
+    | Mod (x1, x2) -> C ("Mod", (x1, x2), Pair (kind_annotation, List string))
+    | With (x1, x2, x3) ->
+        C
+          ( "With",
+            (x1, x2, x3),
+            Triple (kind_annotation, typeexpr_t, List string) )
+    | Kind_of x -> C ("Kind_of", x, typeexpr_t)
+    | Product x -> C ("Product", x, List kind_annotation))
+
 and typedecl_param_desc =
   let open Lang.TypeDecl in
   Variant (function Any -> C0 "Any" | Var x -> C ("Var", x, string))
@@ -378,6 +393,7 @@ and typedecl_param =
       F ("desc", (fun t -> t.desc), typedecl_param_desc);
       F ("variance", (fun t -> t.variance), Option typedecl_variance);
       F ("injectivity", (fun t -> t.injectivity), bool);
+      F ("kind", (fun t -> t.kind), kind_annotation);
     ]
 
 and typedecl_equation =
@@ -391,6 +407,7 @@ and typedecl_equation =
         ( "constraints",
           (fun t -> t.constraints),
           List (Pair (typeexpr_t, typeexpr_t)) );
+      F ("kind", (fun t -> t.kind), kind_annotation);
     ]
 
 and typedecl_t =
@@ -667,7 +684,11 @@ and typeexpr_t =
     | Object x -> C ("Object", x, typeexpr_object)
     | Class (x1, x2) ->
         C ("Class", ((x1 :> Paths.Path.t), x2), Pair (path, List typeexpr_t))
-    | Poly (x1, x2) -> C ("Poly", (x1, x2), Pair (List string, typeexpr_t))
+    | Poly (x1, x2) ->
+        C
+          ( "Poly",
+            (x1, x2),
+            Pair (List (Pair (string, kind_annotation)), typeexpr_t) )
     | Quote x -> C ("Quote", x, typeexpr_t)
     | Splice x -> C ("Splice", x, typeexpr_t)
     | Package x -> C ("Package", x, typeexpr_package))

src/search/html.ml

@@ -55,7 +55,7 @@ let field_rhs ({ mutable_ = _; type_; parent_type = _ } : Entry.field_entry) =
   " : " ^ Text.of_type type_
 
 let typedecl_params ?(delim = `parens) params =
-  let format_param { TypeDecl.desc; variance; injectivity } =
+  let format_param { TypeDecl.desc; variance; injectivity; kind = _ } =
     let desc =
       match desc with TypeDecl.Any -> [ "_" ] | Var s -> [ "'"; s ]
     in