Add support for unit tests via $[test] attribute

src/lib/jib_compile.ml

@@ -131,6 +131,7 @@ 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;
@@ -181,6 +182,7 @@ 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;
@@ -2803,6 +2805,16 @@ 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
+
   let toplevel_lets_of_ast ast =
     let toplevel_lets_of_def = function
       | DEF_aux (DEF_let (LB_aux (LB_val (pat, _), _)), _) -> pat_ids pat
@@ -2852,7 +2864,14 @@ 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
-    let roots = Specialize.get_initial_calls () |> List.map (fun id -> Callgraph.Function id) |> NodeSet.of_list 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. *)
+    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
+    in
     let roots = IdSet.fold (fun id roots -> NodeSet.add (Callgraph.Type id) roots) C.preserve_types roots in
     let roots = NodeSet.add (Callgraph.Type (mk_id "exception")) roots in
     let roots =
@@ -2861,6 +2880,7 @@ module Make (C : CONFIG) = struct
     let roots =
       NodeSet.union (toplevel_lets_of_ast ast |> List.map (fun id -> Callgraph.Letbind id) |> NodeSet.of_list) roots
     in
+
     let g = G.prune roots NodeSet.empty g in
     let ast = Callgraph.filter_ast NodeSet.empty g ast in
 

src/lib/jib_compile.mli

@@ -87,6 +87,7 @@ 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

@@ -4662,6 +4662,15 @@ let check_fundef_lazy env def_annot (FD_aux (FD_function (recopt, tannot_opt, fu
     | 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"
+    )
+  end;
   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

src/sail_c_backend/c_backend.ml

@@ -2393,6 +2393,48 @@ module Codegen (Config : CODEGEN_CONFIG) = struct
         |> List.map string |> separate hardline
       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 "};"
+      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 "};"
+      in
+
+      (* A simple function to run the unit tests. It isn't called from anywhere
+         by default and you don't need to use it - you can use SAIL_TESTS directly
+         in your own custom test runner. *)
+      let model_test =
+        [
+          Printf.sprintf "%svoid model_test()" (static ());
+          "{";
+          "  for (size_t i = 0; SAIL_TESTS[i] != 0 && SAIL_TEST_NAMES[i] != 0; ++i) {";
+          "    model_init();";
+          "    printf(\"Testing %s\\n\", SAIL_TEST_NAMES[i]);";
+          "    SAIL_TESTS[i](UNIT);";
+          "    printf(\"Pass\\n\");";
+          "    model_fini();";
+          "  }";
+          "}";
+        ]
+        |> List.map string |> separate hardline
+      in
+
       let model_main =
         let extra_pre =
           List.filter_map
@@ -2439,7 +2481,8 @@ module Codegen (Config : CODEGEN_CONFIG) = struct
                  model_init ^^ hlhl ^^ model_fini ^^ hlhl ^^ model_pre_exit ^^ hlhl ^^ model_default_main ^^ hlhl
                else empty
              )
-          ^^ model_main ^^ hardline ^^ end_extern_cpp ^^ hardline
+          ^^ model_main ^^ hlhl ^^ unit_test_functions ^^ hlhl ^^ unit_test_names ^^ hlhl ^^ model_test ^^ hardline
+          ^^ end_extern_cpp ^^ hardline
           )
       )
     with Type_error.Type_error (l, err) ->

test/c/test_attribute.expect

@@ -0,0 +1 @@
+ok

test/c/test_attribute.sail

@@ -0,0 +1,24 @@
+default Order dec
+
+$[test]
+function test_0() -> unit = {
+  print_endline("test_0");
+}
+
+$[test]
+val test_1 : unit -> unit
+
+function test_1() = {
+  print_endline("test_1");
+}
+
+$[test]
+val test_2 : unit -> unit
+
+function test_2() -> unit = {
+  print_endline("test_2");
+}
+
+function main() -> unit = {
+  print_endline("ok");
+}

test/typecheck/fail/test_not_unit_0.expect

@@ -0,0 +1 @@
+test_bad must have type unit -> unit

test/typecheck/fail/test_not_unit_0.sail

@@ -0,0 +1,11 @@
+default Order dec
+
+$[test]
+function test_good() -> unit = {
+  print_endline("test_good");
+}
+
+$[test]
+function test_bad() -> int = 5
+
+function main() -> unit = ()

test/typecheck/fail/test_not_unit_1.expect

@@ -0,0 +1 @@
+test_bad must have type unit -> unit

test/typecheck/fail/test_not_unit_1.sail

@@ -0,0 +1,6 @@
+default Order dec
+
+$[test]
+function test_bad(x : int) -> unit = ()
+
+function main() -> unit = ()

test/typecheck/fail/test_not_unit_2.expect

@@ -0,0 +1 @@
+test_bad must have type unit -> unit

test/typecheck/fail/test_not_unit_2.sail

@@ -0,0 +1,6 @@
+default Order dec
+
+$[test]
+function test_bad(x : unit, y : unit) -> unit = ()
+
+function main() -> unit = ()

test/typecheck/fail/test_not_unit_3.expect

@@ -0,0 +1 @@
+test_bad must have type unit -> unit

test/typecheck/fail/test_not_unit_3.sail

@@ -0,0 +1,8 @@
+default Order dec
+
+$[test]
+val test_bad : unit -> int
+
+function test_bad() = 5
+
+function main() -> unit = ()

test/typecheck/fail/test_not_unit_4.expect

@@ -0,0 +1 @@
+test attribute must be on function declaration

test/typecheck/fail/test_not_unit_4.sail

@@ -0,0 +1,8 @@
+default Order dec
+
+val test_bad : unit -> unit
+
+$[test]
+function test_bad() = ()
+
+function main() -> unit = ()