Allow parsing doc comments and attributes on function clauses

After de-scattering, top-level doc comments on function clauses will end
up attached to the function clause AST nodes themselves, so we now print
and parse such comments (and attributes) appropriately

scattered_function_doc.sail

@@ -0,0 +1,11 @@
+default Order dec
+
+val foo : int -> unit
+
+scattered function foo
+
+/*! First clause doc comment */
+function clause foo(0) = ()
+
+/*! Second clause doc comment */
+function clause foo(_) = ()

src/lib/chunk_ast.ml

@@ -103,7 +103,7 @@ and chunk =
       rec_opt : chunks option;
       typq_opt : chunks option;
       return_typ_opt : chunks option;
-      funcls : pexp_chunks list;
+      funcls : (chunks * pexp_chunks) list;
     }
   | Val of { id : id; extern_opt : extern option; typq_opt : chunks option; typ : chunks }
   | Enum of { id : id; enum_functions : chunks list option; members : chunks list }
@@ -207,7 +207,9 @@ let rec prerr_chunk indent = function
         | None -> ()
       end;
       List.iteri
-        (fun i funcl ->
+        (fun i (funcl_header, funcl) ->
+          Printf.eprintf "%s  header %d:\n" indent i;
+          Queue.iter (prerr_chunk (indent ^ "    ")) funcl_header;
           Printf.eprintf "%s  pat %d:\n" indent i;
           Queue.iter (prerr_chunk (indent ^ "    ")) funcl.pat;
           begin
@@ -1057,7 +1059,18 @@ and chunk_pexp ?delim comments (Pat_aux (aux, l)) =
       ignore (pop_trailing_comment comments exp_chunks (ending_line_num l));
       { funcl_space = true; pat = pat_chunks; guard = Some guard_chunks; body = exp_chunks }
 
-let chunk_funcl comments (FCL_aux (FCL_funcl (_, pexp), _)) = chunk_pexp comments pexp
+let chunk_funcl comments funcl =
+  let chunks = Queue.create () in
+  let rec chunk_funcl' comments (FCL_aux (aux, _)) =
+    match aux with
+    | FCL_attribute (attr, arg, funcl) ->
+        Queue.add (Atom (Printf.sprintf "$[%s %s]" attr arg)) chunks;
+        Queue.add (Spacer (false, 1)) chunks;
+        chunk_funcl' comments funcl
+    | FCL_doc (_, funcl) -> chunk_funcl' comments funcl
+    | FCL_funcl (_, pexp) -> chunk_pexp comments pexp
+  in
+  (chunks, chunk_funcl' comments funcl)
 
 let chunk_quant_item comments chunks last = function
   | QI_aux (QI_id kopt, l) ->

src/lib/chunk_ast.mli

@@ -95,7 +95,7 @@ and chunk =
       rec_opt : chunks option;
       typq_opt : chunks option;
       return_typ_opt : chunks option;
-      funcls : pexp_chunks list;
+      funcls : (chunks * pexp_chunks) list;
     }
   | Val of { id : Parse_ast.id; extern_opt : Parse_ast.extern option; typq_opt : chunks option; typ : chunks }
   | Enum of { id : Parse_ast.id; enum_functions : chunks list option; members : chunks list }

src/lib/format_sail.ml

@@ -670,12 +670,14 @@ module Make (Config : CONFIG) = struct
     let guarded_pat, body = doc_pexp_chunks_pair opts pexp in
     separate space [guarded_pat; string "=>"; body]
 
-  and doc_funcl return_typ_opt opts pexp =
+  and doc_funcl return_typ_opt opts (header, pexp) =
     let return_typ =
       match return_typ_opt with
       | Some chunks -> space ^^ prefix_parens indent (string "->") (doc_chunks opts chunks) ^^ space
       | None -> space
     in
+    doc_chunks opts header
+    ^^
     match pexp.guard with
     | None ->
         (if pexp.funcl_space then space else empty)

src/lib/initial_check.ml

@@ -1081,8 +1081,14 @@ let to_ast_typschm_opt ctx (P.TypSchm_opt_aux (aux, l)) : tannot_opt ctx_out =
       let typq, ctx = to_ast_typquant ctx tq in
       (Typ_annot_opt_aux (Typ_annot_opt_some (typq, to_ast_typ ctx typ), l), ctx)
 
-let to_ast_funcl ctx (P.FCL_aux (fcl, l) : P.funcl) : uannot funcl =
+let rec to_ast_funcl doc attrs ctx (P.FCL_aux (fcl, l) : P.funcl) : uannot funcl =
   match fcl with
+  | P.FCL_attribute (attr, arg, fcl) -> to_ast_funcl doc ((attr, arg, l) :: attrs) ctx fcl
+  | P.FCL_doc (doc_comment, fcl) -> begin
+      match doc with
+      | Some _ -> raise (Reporting.err_general l "Function clause has multiple documentation comments")
+      | None -> to_ast_funcl (Some doc_comment) attrs ctx fcl
+    end
   | P.FCL_funcl (id, pexp) ->
       FCL_aux (FCL_funcl (to_ast_id ctx id, to_ast_case ctx pexp), (mk_def_annot l, empty_uannot))
 
@@ -1099,12 +1105,14 @@ let to_ast_impl_funcls ctx (P.FCL_aux (fcl, l) : P.funcl) : uannot funcl list =
             targets
       | None -> [FCL_aux (FCL_funcl (to_ast_id ctx id, to_ast_case ctx pexp), (mk_def_annot l, empty_uannot))]
     )
+  | _ -> raise (Reporting.err_general l "Attributes or documentation comment not permitted here")
 
 let to_ast_fundef ctx (P.FD_aux (fd, l) : P.fundef) : uannot fundef =
   match fd with
   | P.FD_function (rec_opt, tannot_opt, _, funcls) ->
       let tannot_opt, ctx = to_ast_tannot_opt ctx tannot_opt in
-      FD_aux (FD_function (to_ast_rec ctx rec_opt, tannot_opt, List.map (to_ast_funcl ctx) funcls), (l, empty_uannot))
+      FD_aux
+        (FD_function (to_ast_rec ctx rec_opt, tannot_opt, List.map (to_ast_funcl None [] ctx) funcls), (l, empty_uannot))
 
 let rec to_ast_mpat ctx (P.MP_aux (mpat, l)) =
   MP_aux
@@ -1165,7 +1173,7 @@ let to_ast_scattered ctx (P.SD_aux (aux, l)) =
     | P.SD_function (rec_opt, tannot_opt, _, id) ->
         let tannot_opt, _ = to_ast_tannot_opt ctx tannot_opt in
         (None, SD_function (to_ast_rec ctx rec_opt, tannot_opt, to_ast_id ctx id), ctx)
-    | P.SD_funcl funcl -> (None, SD_funcl (to_ast_funcl ctx funcl), ctx)
+    | P.SD_funcl funcl -> (None, SD_funcl (to_ast_funcl None [] ctx funcl), ctx)
     | P.SD_variant (id, parse_typq) ->
         let id = to_ast_id ctx id in
         let typq, typq_ctx = to_ast_typquant ctx parse_typq in

src/lib/parse_ast.ml

@@ -308,7 +308,12 @@ type rec_opt_aux =
     Rec_none (* no termination measure *)
   | Rec_measure of pat * exp (* recursive with termination measure *)
 
-type funcl_aux = (* Function clause *)
+type funcl = FCL_aux of funcl_aux * l
+
+and funcl_aux =
+  (* Function clause *)
+  | FCL_attribute of string * string * funcl
+  | FCL_doc of string * funcl
   | FCL_funcl of id * pexp
 
 type type_union = Tu_aux of type_union_aux * l
@@ -326,8 +331,6 @@ type effect_opt = Effect_opt_aux of effect_opt_aux * l
 
 type rec_opt = Rec_aux of rec_opt_aux * l
 
-type funcl = FCL_aux of funcl_aux * l
-
 type subst_aux =
   (* instantiation substitution *)
   | IS_typ of kid * atyp (* instantiate a type variable with a type *)

src/lib/parser.mly

@@ -231,6 +231,9 @@ let fix_extern typschm = function
   | None -> None
   | Some extern -> Some { extern with pure = typschm_is_pure typschm }
 
+let funcl_annot fs fcl =
+  List.fold_right (fun f fcl -> f fcl) fs fcl
+
 let effect_deprecated l =
   Reporting.warn ~once_from:__POS__ "Deprecated" l "Explicit effect annotations are deprecated. They are no longer used and can be removed."
 
@@ -242,7 +245,7 @@ let warn_extern_effect l =
 
 let set_syntax_deprecated l =
   Reporting.warn ~once_from:__POS__ "Deprecated" l "Old set syntax, {|1, 2, 3|} can now be written as {1, 2, 3}."
-  
+
 %}
 
 /*Terminals with no content*/
@@ -892,6 +895,12 @@ vector_update_list:
   | vector_update Comma vector_update_list
     { $1 :: $3 }
 
+funcl_annotation:
+  | attr = Attribute
+    { (fun funcl -> FCL_aux (FCL_attribute (fst attr, snd attr, funcl), loc $startpos(attr) $endpos(attr))) }
+  | doc = Doc
+    { (fun funcl -> FCL_aux (FCL_doc (doc, funcl), loc $startpos(doc) $endpos(doc))) }
+
 funcl_patexp:
   | pat Eq exp
     { mk_pexp (Pat_exp ($1, $3)) $startpos $endpos }
@@ -913,6 +922,8 @@ funcl_patexp_typ:
     { (mk_pexp (Pat_when ($5, $7, $12)) $startpos $endpos, mk_tannot $2 $10 $startpos $endpos($10)) }
 
 funcl:
+  | funcl_annotation+ id funcl_patexp
+    { funcl_annot $1 (mk_funcl (FCL_funcl ($2, $3)) $startpos($2) $endpos) }
   | id funcl_patexp
     { mk_funcl (FCL_funcl ($1, $2)) $startpos $endpos }
 
@@ -923,8 +934,12 @@ funcls2:
     { $1 :: $3 }
 
 funcls:
+  | funcl_annotation+ id funcl_patexp_typ
+    { let pexp, tannot = $3 in ([funcl_annot $1 (mk_funcl (FCL_funcl ($2, pexp)) $startpos($2) $endpos)], tannot) }
   | id funcl_patexp_typ
     { let pexp, tannot = $2 in ([mk_funcl (FCL_funcl ($1, pexp)) $startpos $endpos], tannot) }
+  | funcl_annotation+ id funcl_patexp And funcls2
+    { (funcl_annot $1 (mk_funcl (FCL_funcl ($2, $3)) $startpos($2) $endpos($3)) :: $5, mk_tannotn) }
   | id funcl_patexp And funcls2
     { (mk_funcl (FCL_funcl ($1, $2)) $startpos $endpos($2) :: $4, mk_tannotn) }
 

src/lib/pretty_print_sail.ml

@@ -583,7 +583,9 @@ and doc_vector_update = function
   | VU_range (high, low, value) ->
       separate space [doc_atomic_exp high; string ".."; doc_atomic_exp low; equals; doc_exp value]
 
-let doc_funcl (FCL_aux (FCL_funcl (id, Pat_aux (pexp, _)), _)) =
+let doc_funcl (FCL_aux (FCL_funcl (id, Pat_aux (pexp, _)), (def_annot, _))) =
+  doc_def_annot def_annot
+  ^^
   match pexp with
   | Pat_exp (pat, exp) -> group (separate space [doc_id id; doc_pat pat; equals; doc_exp_as_block exp])
   | Pat_when (pat, wh, exp) ->

test/typecheck/scattered_function_doc.sail

@@ -0,0 +1,11 @@
+default Order dec
+
+val foo : int -> unit
+
+scattered function foo
+
+/*! First clause doc comment */
+function clause foo(0) = ()
+
+/*! Second clause doc comment */
+function clause foo(_) = ()