Adjust attribute handling for Sail unit tests

src/lib/jib_compile.ml

@@ -131,7 +131,6 @@ type ctx = {
   registers : ctyp Bindings.t;
   letbinds : int list;
   letbind_ids : IdSet.t;
-  unit_test_ids : IdSet.t;
   no_raw : bool;
   no_static : bool;
   coverage_override : bool;
@@ -182,7 +181,6 @@ let initial_ctx ?for_target env effect_info =
     registers = Bindings.empty;
     letbinds = [];
     letbind_ids = IdSet.empty;
-    unit_test_ids = IdSet.empty;
     no_raw = false;
     no_static = false;
     coverage_override = true;
@@ -2806,14 +2804,15 @@ module Make (C : CONFIG) = struct
     (if reverse then List.rev ctype_defs else ctype_defs) @ cdefs
 
   let unit_tests_of_ast ast =
-    let unit_tests_of_def = function
-      | DEF_aux (DEF_fundef (FD_aux (FD_function (_, _, [FCL_aux (FCL_funcl (id, _), _)]), _)), def_annot)
-        when Option.is_some (get_def_attribute "test" def_annot) ->
-          IdSet.singleton id
-      | _ -> IdSet.empty
-    in
-    let unit_tests_of_defs defs = List.fold_left IdSet.union IdSet.empty (List.map unit_tests_of_def defs) in
-    unit_tests_of_defs ast.defs |> IdSet.elements
+    List.fold_left
+      (fun ids -> function
+        | DEF_aux (DEF_val (VS_aux (VS_val_spec (_, id, _), _)), def_annot)
+          when Option.is_some (get_def_attribute "test" def_annot) ->
+            IdSet.add id ids
+        | _ -> ids
+      )
+      IdSet.empty ast.defs
+    |> IdSet.elements
 
   let toplevel_lets_of_ast ast =
     let toplevel_lets_of_def = function
@@ -2864,10 +2863,10 @@ module Make (C : CONFIG) = struct
     let module G = Graph.Make (Callgraph.Node) in
     let g = Callgraph.graph_of_ast ast in
     let module NodeSet = Set.Make (Callgraph.Node) in
-    (* Get the list of unit tests (functions with $[test]). We can't do this in
-       compile_def because they would have already been dead-code elimintate by then. *)
+    (* Get the list of unit tests (valspecs with $[test]), so we can
+       add them to the list of roots to avoid pruning them as
+       dead-code. *)
     let unit_tests = unit_tests_of_ast ast in
-    let ctx = { ctx with unit_test_ids = unit_tests |> IdSet.of_list } in
 
     let roots =
       Specialize.get_initial_calls () @ unit_tests |> List.map (fun id -> Callgraph.Function id) |> NodeSet.of_list

src/lib/jib_compile.mli

@@ -87,7 +87,6 @@ type ctx = {
   registers : ctyp Bindings.t;
   letbinds : int list;
   letbind_ids : IdSet.t;
-  unit_test_ids : IdSet.t;
   no_raw : bool;
   no_static : bool;
   coverage_override : bool;

src/lib/type_check.ml

@@ -4594,6 +4594,15 @@ let check_funcls_complete ?global_env l env funcls typ =
 
 let empty_tannot_opt = Typ_annot_opt_aux (Typ_annot_opt_none, Parse_ast.Unknown)
 
+let check_test_attribute def_annot typ =
+  if Option.is_some (get_def_attribute "test" def_annot) then (
+    match typ with
+    | Typ_aux (Typ_fn ([arg], ret), _) when is_unit_typ arg && is_unit_typ ret -> ()
+    | _ ->
+        Reporting.err_general def_annot.loc "Functions with the $[test] attribute must have type 'unit -> unit'"
+        |> raise
+  )
+
 let check_fundef_lazy env def_annot (FD_aux (FD_function (recopt, tannot_opt, funcls), (l, _))) =
   let id =
     match
@@ -4657,20 +4666,21 @@ let check_fundef_lazy env def_annot (FD_aux (FD_function (recopt, tannot_opt, fu
         let err_l = Option.fold ~none:l ~some:(fun val_l -> Hint ("val here", val_l, l)) have_val_spec in
         typ_error err_l "function does not have a function type"
   in
-  begin
-    match have_val_spec with
-    | Some vs_l -> check_tannot_opt ~def_type:"function" vs_l env vtyp_ret tannot_opt
-    | None -> ()
-  end;
-  (* Check $[test] functions have type unit -> unit. *)
-  (* TODO: Does the annotation go on the type declaration or the function definition? *)
-  begin
-    if Option.is_some (get_def_attribute "test" def_annot) then (
-      match (vtyp_args, vtyp_ret) with
-      | [arg], ret when is_unit_typ arg && is_unit_typ ret -> ()
-      | _ -> typ_error l "$[test] functions must have type: unit -> unit"
+
+  (* If we have a val_spec, check that any annotations on the function are consistent *)
+  ( match have_val_spec with
+  | Some vs_l -> (
+      check_tannot_opt ~def_type:"function" vs_l env vtyp_ret tannot_opt;
+      match get_def_attribute "test" def_annot with
+      | Some (l, _) ->
+          "Function with separate val prototype has a $[test] attribute, it should be attached there instead."
+          |> Reporting.err_general (Hint ("val declaration here", vs_l, l))
+          |> raise
+      | None -> ()
     )
-  end;
+  | None -> check_test_attribute def_annot typ
+  );
+
   typ_debug (lazy ("Checking fundef " ^ string_of_id id ^ " has type " ^ string_of_bind (quant, typ)));
   let funcl_env =
     if Option.is_some have_val_spec then Env.add_typquant l quant env
@@ -4703,7 +4713,7 @@ let check_fundef_lazy env def_annot (FD_aux (FD_function (recopt, tannot_opt, fu
          then (funcls, fun attrs -> attrs)
          else check_funcls_complete l funcl_env funcls typ
        in
-       let def_annot = fix_body_visibility (update_attr def_annot) in
+       let def_annot = fix_body_visibility (update_attr def_annot) |> remove_def_attribute "test" in
        DEF_aux (DEF_fundef (FD_aux (FD_function (recopt, empty_tannot_opt, funcls), (l, empty_tannot))), def_annot)
       )
   in
@@ -4777,6 +4787,7 @@ let check_val_spec env def_annot (VS_aux (vs, (l, _))) =
         (* !opt_expand_valspec controls whether the actual valspec in
            the AST is expanded, the val_spec type stored in the
            environment is always expanded and uses typq' and typ' *)
+        check_test_attribute def_annot typ;
         let typq, typ = if !opt_expand_valspec then (typq', typ') else (typq, typ) in
         let vs = VS_val_spec (TypSchm_aux (TypSchm_ts (typq, typ), ts_l), id, exts) in
         (vs, id, typq', typ', env)

src/sail_c_backend/c_backend.ml

@@ -2252,6 +2252,16 @@ module Codegen (Config : CODEGEN_CONFIG) = struct
     | _ :: ast -> c_ast_registers ~early ast
     | [] -> []
 
+  let get_unit_tests cdefs =
+    List.fold_left
+      (fun ids -> function
+        | CDEF_aux (CDEF_val (id, _, _, _), def_annot) when Option.is_some (get_def_attribute "test" def_annot) ->
+            IdSet.add id ids
+        | _ -> ids
+      )
+      IdSet.empty cdefs
+    |> IdSet.elements
+
   let compile_ast env effect_info basename ast =
     try
       let cdefs, ctx = jib_of_ast env effect_info ast in
@@ -2378,7 +2388,7 @@ module Codegen (Config : CODEGEN_CONFIG) = struct
         |> List.map string |> separate hardline
       in
 
-      let model_default_main =
+      let model_main =
         [
           Printf.sprintf "%sint model_main(int argc, char *argv[])" (static ());
           "{";
@@ -2393,27 +2403,25 @@ module Codegen (Config : CODEGEN_CONFIG) = struct
         |> List.map string |> separate hardline
       in
 
+      let unit_tests = get_unit_tests cdefs in
+
       (* TODO: Formatting here isn't quite right. What are the arguments to jump? *)
       let unit_test_functions =
-        string "unit (*const SAIL_TESTS[])(unit) = {"
-        ^^ hardline
-        ^^ jump 2 2
-             (IdSet.fold (fun id acc -> codegen_function_id id ^^ string "," ^^ hardline ^^ acc) ctx.unit_test_ids empty
-             ^^ string "0" ^^ hardline
-             )
-        ^^ string "};"
+        ["unit (*const SAIL_TESTS[])(unit) = {"]
+        @ Util.map_last
+            (fun is_last id -> sprintf "  %s%s" (sgen_function_id id) (if is_last then "" else ","))
+            unit_tests
+        @ ["};"]
+        |> separate_map hardline string
       in
 
       let unit_test_names =
-        string "const char* const SAIL_TEST_NAMES[] = {"
-        ^^ hardline
-        ^^ jump 2 2
-             (IdSet.fold
-                (fun id acc -> string ("\"" ^ String.escaped (string_of_id id) ^ "\"") ^^ string "," ^^ hardline ^^ acc)
-                ctx.unit_test_ids empty
-             ^^ string "0" ^^ hardline
-             )
-        ^^ string "};"
+        ["const char* const SAIL_TEST_NAMES[] = {"]
+        @ Util.map_last
+            (fun is_last id -> sprintf "  \"%s\"%s" (String.escaped (string_of_id id)) (if is_last then "" else ","))
+            unit_tests
+        @ ["};"]
+        |> separate_map hardline string
       in
 
       (* A simple function to run the unit tests. It isn't called from anywhere
@@ -2435,7 +2443,7 @@ module Codegen (Config : CODEGEN_CONFIG) = struct
         |> List.map string |> separate hardline
       in
 
-      let model_main =
+      let actual_main =
         let extra_pre =
           List.filter_map
             (function CDEF_aux (CDEF_pragma ("c_in_main", arg), _) -> Some ("  " ^ arg) | _ -> None)
@@ -2478,10 +2486,10 @@ module Codegen (Config : CODEGEN_CONFIG) = struct
           ^^ (if Config.generate_header then hardline ^^ Printf.ksprintf string "#include \"%s.h\"" basename else empty)
           ^^ hlhl ^^ docs ^^ hlhl
           ^^ ( if not Config.no_rts then
-                 model_init ^^ hlhl ^^ model_fini ^^ hlhl ^^ model_pre_exit ^^ hlhl ^^ model_default_main ^^ hlhl
+                 model_init ^^ hlhl ^^ model_fini ^^ hlhl ^^ model_pre_exit ^^ hlhl ^^ model_main ^^ hlhl
                else empty
              )
-          ^^ model_main ^^ hlhl ^^ unit_test_functions ^^ hlhl ^^ unit_test_names ^^ hlhl ^^ model_test ^^ hardline
+          ^^ actual_main ^^ hlhl ^^ unit_test_functions ^^ hlhl ^^ unit_test_names ^^ hlhl ^^ model_test ^^ hardline
           ^^ end_extern_cpp ^^ hardline
           )
       )

test/typecheck/fail/test_not_unit_0.expect

@@ -1 +1,5 @@
-test_bad must have type unit -> unit
+Error:
+fail/test_not_unit_0.sail:11.0-30:
+11 |function test_bad() -> int = 5
+   |^----------------------------^
+   | Functions with the $[test] attribute must have type 'unit -> unit'

test/typecheck/fail/test_not_unit_0.sail

@@ -1,5 +1,7 @@
 default Order dec
 
+$include <prelude.sail>
+
 $[test]
 function test_good() -> unit = {
   print_endline("test_good");

test/typecheck/fail/test_not_unit_1.expect

@@ -1 +1,5 @@
-test_bad must have type unit -> unit
+Error:
+fail/test_not_unit_1.sail:4.0-39:
+4 |function test_bad(x : int) -> unit = ()
+  |^-------------------------------------^
+  | Functions with the $[test] attribute must have type 'unit -> unit'

test/typecheck/fail/test_not_unit_2.expect

@@ -1 +1,5 @@
-test_bad must have type unit -> unit
+Error:
+fail/test_not_unit_2.sail:4.0-50:
+4 |function test_bad(x : unit, y : unit) -> unit = ()
+  |^------------------------------------------------^
+  | Functions with the $[test] attribute must have type 'unit -> unit'

test/typecheck/fail/test_not_unit_3.expect

@@ -1 +1,5 @@
-test_bad must have type unit -> unit
+Error:
+fail/test_not_unit_3.sail:4.0-26:
+4 |val test_bad : unit -> int
+  |^------------------------^
+  | Functions with the $[test] attribute must have type 'unit -> unit'

test/typecheck/fail/test_not_unit_4.expect

@@ -1 +1,8 @@
-test attribute must be on function declaration
+Error:
+fail/test_not_unit_4.sail:3.4-12:
+3 |val test_bad : unit -> unit
+  |    ^------^ val declaration here
+fail/test_not_unit_4.sail:5.0-7:
+5 |$[test]
+  |^-----^
+  | Function with separate val prototype has a $[test] attribute, it should be attached there instead.