Zero alloc annotation for OxCaml

CHANGES.md

@@ -6,6 +6,7 @@
   invocation, eliminating the need for shell scripting (@davesnx, #1387)
 - Support for OxCaml (@lukemaurer, @art-w, #1399)
 - OCaml 5.5.0 support (@panglesd, @xvw, #1406)
+- Support for OxCaml zero alloc definitions (@Leonidas-from-XIV, #)
 
 ### Fixed
 - Fix compile-time crashing bugs #930 and #1385 (@jonludlam, #1400)

src/document/generator.ml

@@ -988,13 +988,38 @@ module Make (Syntax : SYNTAX) = struct
         | External _ -> ([ "external" ], Syntax.Type.External.semicolon)
       in
       let name = Paths.Identifier.name t.id in
+      let zero_alloc =
+        match
+          List.find
+            (function Odoc_model.Lang.Value.Zero_alloc _ -> true)
+            t.ext_attr
+        with
+        | exception Not_found -> O.noop
+        | Zero_alloc Ignore -> O.noop
+        | Zero_alloc (Respect {opt; strict; arity; custom_error_message}) ->
+            let ext_arg = match opt, strict with
+              | Some (), None -> " opt"
+              | None, Some () -> " strict"
+              | _, _ -> ""
+            in
+            let ext_arg = match arity with
+              | None -> ext_arg
+              | Some n -> ext_arg ^ (Printf.sprintf "arity %d" n)
+            in
+            let ext_arg = match custom_error_message with
+              | None -> ext_arg
+              | Some s -> ext_arg ^ (Printf.sprintf "custom_error_message %S" s)
+            in
+            let ext_attr = Printf.sprintf "[@@zero_alloc%s]" ext_arg in
+            O.cut ++ O.txt " " ++ O.txt ext_attr
+      in
       let content =
         O.documentedSrc
           (O.box_hv
           @@ O.keyword Syntax.Value.variable_keyword
              ++ O.txt " " ++ O.txt name
              ++ O.txt Syntax.Type.annotation_separator
-             ++ O.cut ++ type_expr t.type_
+             ++ O.cut ++ type_expr t.type_ ++ zero_alloc
              ++ if semicolon then O.txt ";" else O.noop)
       in
       let attr = [ "value" ] @ extra_attr in

src/loader/cmi.ml

@@ -757,7 +757,8 @@ let read_value_description ({ident_env ; warnings_tag} as env) parent id vd =
         External primitives
     | _ -> assert false
   in
-  Value { Value.id; source_loc; doc; type_; value }
+  let ext_attr = List.filter_map Doc_attr.known_attribute vd.val_attributes in
+  Value { Value.id; source_loc; doc; type_; value; ext_attr }
 
 #if defined OXCAML
 let is_mutable = Types.is_mutable

src/loader/cmt.ml

@@ -53,7 +53,9 @@ let rec read_pattern env parent doc pat =
           Cmi.mark_type_expr pat.pat_type;
           let type_ = Cmi.read_type_expr env pat.pat_type in
           let value = Abstract in
-          [Value {id; source_loc; doc; type_; value}]
+          (* TODO read ext_attr out of id *)
+          let ext_attr = [] in
+          [Value {id; source_loc; doc; type_; value; ext_attr}]
 #if OCAML_VERSION < (5,2, 0)
     | Tpat_alias(pat, id, _) ->
 #elif defined OXCAML
@@ -68,7 +70,8 @@ let rec read_pattern env parent doc pat =
           Cmi.mark_type_expr pat.pat_type;
           let type_ = Cmi.read_type_expr env pat.pat_type in
           let value = Abstract in
-          Value {id; source_loc; doc; type_; value} :: read_pattern env parent doc pat
+          let ext_attr = [] in
+          Value {id; source_loc; doc; type_; value; ext_attr} :: read_pattern env parent doc pat
     | Tpat_constant _ -> []
     | Tpat_tuple pats ->
 #if OCAML_VERSION >= (5, 4, 0) || defined OXCAML

src/loader/cmti.ml

@@ -233,7 +233,8 @@ let read_value_description env parent vd =
     | [] -> Value.Abstract
     | primitives -> External primitives
   in
-  Value { Value.id; source_loc; doc; type_; value }
+  let ext_attr = List.filter_map Doc_attr.known_attribute vd.val_attributes in
+  Value { Value.id; source_loc; doc; type_; value; ext_attr }
 
 let read_type_parameter (ctyp, var_and_injectivity)  =
   let open TypeDecl in

src/loader/doc_attr.ml

@@ -74,6 +74,108 @@ let attribute_unpack = function
   | { Location.txt = name; loc }, attr_payload -> (name, attr_payload, loc)
 #endif
 
+module Zero_alloc = struct
+  let constant_of_expression (pexp_desc : Parsetree.expression_desc) =
+    match pexp_desc with
+    | Pexp_constant (Pconst_integer (i, _))  -> Some (`Integer (int_of_string i))
+    | Pexp_constant (Pconst_string (s, _, _)) -> Some (`String s)
+    | _ -> None
+
+  let parse_argument_of_keyword (pstr_desc : Parsetree.structure_item_desc) =
+    match pstr_desc with
+    | Pstr_eval ({pexp_desc; _}, _) -> constant_of_expression pexp_desc
+    | _ -> None
+
+  let identifier_of_expression (pexp_desc : Parsetree.expression_desc) =
+    match pexp_desc with
+    | Pexp_ident {txt=longident; _} -> (
+      match longident with
+      | Lident s -> Some s
+      | _ -> None)
+    | _ -> None
+
+  let keyword_of_structure_arg (pstr_desc: Parsetree.structure_item_desc) =
+    match pstr_desc with
+    | Pstr_eval ({pexp_desc; _}, _) -> (
+      match identifier_of_expression pexp_desc with
+      | Some "arity" -> Some `Arity
+      | Some "custom_error_message" -> Some `Custom_error_message
+      | Some "ignore" -> Some `Ignore
+      | Some "strict" -> Some `Strict
+      | Some "opt" -> Some `Opt
+      | _ -> None)
+    | _ -> None
+
+  let rec arguments_of_structure_items so_far (structure_items : Parsetree.structure) =
+    match structure_items with
+    | [] -> so_far
+    | {pstr_desc; _} :: structure_items ->
+        match (so_far : Lang.Value.Zero_alloc.t) with
+        | Ignore -> Ignore
+        | Respect respect ->
+          match keyword_of_structure_arg pstr_desc with
+          | None -> so_far
+          | Some `Ignore -> Lang.Value.Zero_alloc.Ignore
+          | Some `Strict ->
+            let so_far = Lang.Value.Zero_alloc.Respect {respect with strict = Some ()} in
+            arguments_of_structure_items so_far structure_items
+          | Some `Opt ->
+            let so_far = Lang.Value.Zero_alloc.Respect {respect with opt = Some ()} in
+            arguments_of_structure_items so_far structure_items
+          | Some `Arity -> (
+            match structure_items with
+            | [] ->
+              (* this is an error *)
+              so_far
+            | {pstr_desc; _} :: structure_items -> (
+              match parse_argument_of_keyword pstr_desc with
+              | Some (`Integer n) ->
+                let arity = Some n in
+                let so_far = Lang.Value.Zero_alloc.Respect {respect with arity} in
+                arguments_of_structure_items so_far structure_items
+              | _ ->
+                (* this is an error *)
+                so_far))
+          | Some `Custom_error_message -> (
+              match structure_items with
+              | [] ->
+                (* this is an error *)
+                so_far
+              | {pstr_desc; _} :: structure_items -> (
+                match parse_argument_of_keyword pstr_desc with
+                | Some (`String s) ->
+                  let custom_error_message = Some s in
+                  let so_far = Lang.Value.Zero_alloc.Respect {respect with custom_error_message} in
+                  arguments_of_structure_items so_far structure_items
+                | _ ->
+                  (* this is an error *)
+                  so_far))
+
+  let arguments_of_payload (payload : Parsetree.payload) =
+    match payload with
+    | PStr structure_items ->
+      (* start with completely empty arguments *)
+      let so_far = Lang.Value.Zero_alloc.Respect {
+        opt = None;
+        strict = None;
+        arity = None;
+        custom_error_message = None}
+      in
+      Some (arguments_of_structure_items so_far structure_items)
+    | _ -> None
+end
+
+let known_attribute attr =
+  let name, payload, _ = attribute_unpack attr in
+  match name with
+#if defined OXCAML
+  | "zero_alloc" -> (
+    match Zero_alloc.arguments_of_payload payload with
+    | Some zalloc_arg -> Some (Lang.Value.Zero_alloc zalloc_arg)
+    | None -> None)
+#endif
+  | _ -> None
+
 type payload = string * Location.t
 
 type parsed_attribute =

src/loader/doc_attr.mli

@@ -93,3 +93,4 @@ type parsed_attribute =
   ]
 
 val parse_attribute : Parsetree.attribute -> parsed_attribute option
+val known_attribute : Parsetree.attribute -> Lang.Value.attr option

src/model/lang.ml

@@ -321,12 +321,26 @@ end =
 and Value : sig
   type value = Abstract | External of string list
 
+  module Zero_alloc : sig
+    type respect = {
+          opt : unit option;
+          strict : unit option;
+          arity : int option;
+          custom_error_message : string option
+        }
+    type t =
+      | Ignore
+      | Respect of respect
+  end
+  type attr = Zero_alloc of Zero_alloc.t
+
   type t = {
     id : Identifier.Value.t;
     source_loc : Identifier.SourceLocation.t option;
     value : value;
     doc : Comment.docs;
     type_ : TypeExpr.t;
+    ext_attr : attr list;
   }
 end =
   Value

src/odoc/extract_code.mli

@@ -1,3 +1,5 @@
+(** [extract ~dst ~input ~names ~line_directives ~warnings_options] extracts source code from document blocks.
+    @param dst File path to write to *)
 val extract :
   dst:string option ->
   input:string ->

src/xref2/lang_of.ml

@@ -692,6 +692,8 @@ and value_ map parent id v =
     doc = docs (parent :> Identifier.LabelParent.t) v.doc;
     type_ = type_expr map (parent :> Identifier.LabelParent.t) v.type_;
     value = v.value;
+    (* TODO this needs to be fixed *)
+    ext_attr = [];
   }
 
 and typ_ext map parent t =

test/generators/cases/oxcaml.mli

@@ -1,2 +1,18 @@
 val f : int -> ('a . 'a -> 'a) -> unit
 (** Polymorphic arguments require parentheses *)
+
+val add : bool -> int -> int -> int [@@zero_alloc]
+(** Zero allocation bindings have an extension attribute attached.
+    See https://oxcaml.org/documentation/miscellaneous-extensions/zero_alloc_check/
+ *)
+
+val add_opt : bool -> int -> int -> int [@@zero_alloc opt]
+(** Like [add] but with an [opt] attribute.
+ *)
+
+val add_strict : bool -> int -> int -> int [@@zero_alloc strict]
+(** Like [add] but with a [strict] attribute.
+ *)
+
+val[@zero_alloc] f : int -> int
+(** Alternative syntax for zero alloc annotation *)

test/generators/html/Oxcaml.html

@@ -15,23 +15,68 @@
    <h1>Module <code><span>Oxcaml</span></code></h1>
   </header>
   <div class="odoc-content">
+   <div class="odoc-spec">
+    <div class="spec value anchored" id="val-add">
+     <a href="#val-add" class="anchor"></a>
+     <code>
+      <span><span class="keyword">val</span> add : 
+       <span>bool <span class="arrow">&#45;&gt;</span></span> 
+       <span>int <span class="arrow">&#45;&gt;</span></span> 
+       <span>int <span class="arrow">&#45;&gt;</span></span> int 
+       [@@zero_alloc]
+      </span>
+     </code>
+    </div>
+    <div class="spec-doc">
+     <p>Zero allocation bindings have an extension attribute attached.
+       See 
+      https://oxcaml.org/documentation/miscellaneous-extensions/zero_alloc_check/
+     </p>
+    </div>
+   </div>
+   <div class="odoc-spec">
+    <div class="spec value anchored" id="val-add_opt">
+     <a href="#val-add_opt" class="anchor"></a>
+     <code>
+      <span><span class="keyword">val</span> add_opt : 
+       <span>bool <span class="arrow">&#45;&gt;</span></span> 
+       <span>int <span class="arrow">&#45;&gt;</span></span> 
+       <span>int <span class="arrow">&#45;&gt;</span></span> int 
+       [@@zero_alloc opt]
+      </span>
+     </code>
+    </div>
+    <div class="spec-doc">
+     <p>Like <code>add</code> but with an <code>opt</code> attribute.</p>
+    </div>
+   </div>
+   <div class="odoc-spec">
+    <div class="spec value anchored" id="val-add_strict">
+     <a href="#val-add_strict" class="anchor"></a>
+     <code>
+      <span><span class="keyword">val</span> add_strict : 
+       <span>bool <span class="arrow">&#45;&gt;</span></span> 
+       <span>int <span class="arrow">&#45;&gt;</span></span> 
+       <span>int <span class="arrow">&#45;&gt;</span></span> int 
+       [@@zero_alloc strict]
+      </span>
+     </code>
+    </div>
+    <div class="spec-doc">
+     <p>Like <code>add</code> but with a <code>strict</code> attribute.</p>
+    </div>
+   </div>
    <div class="odoc-spec">
     <div class="spec value anchored" id="val-f">
      <a href="#val-f" class="anchor"></a>
      <code>
       <span><span class="keyword">val</span> f : 
-       <span>int <span class="arrow">&#45;&gt;</span></span> 
-       <span>
-        <span>('a. 
-         <span><span class="type-var">'a</span> 
-          <span class="arrow">&#45;&gt;</span>
-         </span> <span class="type-var">'a</span>)
-        </span> <span class="arrow">&#45;&gt;</span>
-       </span> unit
+       <span>int <span class="arrow">&#45;&gt;</span></span> int 
+       [@@zero_alloc]
       </span>
      </code>
     </div>
-    <div class="spec-doc"><p>Polymorphic arguments require parentheses</p>
+    <div class="spec-doc"><p>Alternative syntax for zero alloc annotation</p>
     </div>
    </div>
   </div>

test/generators/latex/Oxcaml.tex

@@ -1,5 +1,11 @@
 \section{Module \ocamlinlinecode{Oxcaml}}\label{Oxcaml}%
-\label{Oxcaml--val-f}\ocamlcodefragment{\ocamltag{keyword}{val} f : int \ocamltag{arrow}{$\rightarrow$} ('a.\allowbreak{} \ocamltag{type-var}{'a} \ocamltag{arrow}{$\rightarrow$} \ocamltag{type-var}{'a}) \ocamltag{arrow}{$\rightarrow$} unit}\begin{ocamlindent}Polymorphic arguments require parentheses\end{ocamlindent}%
+\label{Oxcaml--val-add}\ocamlcodefragment{\ocamltag{keyword}{val} add : bool \ocamltag{arrow}{$\rightarrow$} int \ocamltag{arrow}{$\rightarrow$} int \ocamltag{arrow}{$\rightarrow$} int [@@zero\_\allowbreak{}alloc]}\begin{ocamlindent}Zero allocation bindings have an extension attribute attached. See https://oxcaml.org/documentation/miscellaneous-extensions/zero\_alloc\_check/\end{ocamlindent}%
+\medbreak
+\label{Oxcaml--val-add_opt}\ocamlcodefragment{\ocamltag{keyword}{val} add\_\allowbreak{}opt : bool \ocamltag{arrow}{$\rightarrow$} int \ocamltag{arrow}{$\rightarrow$} int \ocamltag{arrow}{$\rightarrow$} int [@@zero\_\allowbreak{}alloc opt]}\begin{ocamlindent}Like \ocamlinlinecode{add} but with an \ocamlinlinecode{opt} attribute.\end{ocamlindent}%
+\medbreak
+\label{Oxcaml--val-add_strict}\ocamlcodefragment{\ocamltag{keyword}{val} add\_\allowbreak{}strict : bool \ocamltag{arrow}{$\rightarrow$} int \ocamltag{arrow}{$\rightarrow$} int \ocamltag{arrow}{$\rightarrow$} int [@@zero\_\allowbreak{}alloc strict]}\begin{ocamlindent}Like \ocamlinlinecode{add} but with a \ocamlinlinecode{strict} attribute.\end{ocamlindent}%
+\medbreak
+\label{Oxcaml--val-f}\ocamlcodefragment{\ocamltag{keyword}{val} f : int \ocamltag{arrow}{$\rightarrow$} int [@@zero\_\allowbreak{}alloc]}\begin{ocamlindent}Alternative syntax for zero alloc annotation\end{ocamlindent}%
 \medbreak
 
 

test/generators/man/Oxcaml.3o

@@ -11,10 +11,31 @@ Oxcaml
 .SH Documentation
 .sp 
 .nf 
-\f[CB]val\fR f : int \f[CB]\->\fR ('a\. \f[CB]'a\fR \f[CB]\->\fR \f[CB]'a\fR) \f[CB]\->\fR unit
+\f[CB]val\fR add : bool \f[CB]\->\fR int \f[CB]\->\fR int \f[CB]\->\fR int [@@zero_alloc]
 .fi 
 .br 
 .ti +2
-Polymorphic arguments require parentheses
+Zero allocation bindings have an extension attribute attached\. See https://oxcaml\.org/documentation/miscellaneous-extensions/zero_alloc_check/
+.nf 
+.sp 
+\f[CB]val\fR add_opt : bool \f[CB]\->\fR int \f[CB]\->\fR int \f[CB]\->\fR int [@@zero_alloc opt]
+.fi 
+.br 
+.ti +2
+Like add but with an opt attribute\.
+.nf 
+.sp 
+\f[CB]val\fR add_strict : bool \f[CB]\->\fR int \f[CB]\->\fR int \f[CB]\->\fR int [@@zero_alloc strict]
+.fi 
+.br 
+.ti +2
+Like add but with a strict attribute\.
+.nf 
+.sp 
+\f[CB]val\fR f : int \f[CB]\->\fR int [@@zero_alloc]
+.fi 
+.br 
+.ti +2
+Alternative syntax for zero alloc annotation
 .nf