Add support for unit tests via $[test] attribute

src/lib/jib_compile.ml

@@ -129,6 +129,7 @@ type ctx = {
   locals : (mut * ctyp) Bindings.t;
   letbinds : int list;
   letbind_ids : IdSet.t;
+  unit_test_ids : IdSet.t;
   no_raw : bool;
   coverage_override : bool;
   def_annot : unit def_annot option;
@@ -176,6 +177,7 @@ let initial_ctx ?for_target env effect_info =
     locals = Bindings.empty;
     letbinds = [];
     letbind_ids = IdSet.empty;
+    unit_test_ids = IdSet.empty;
     no_raw = false;
     coverage_override = true;
     def_annot = None;
@@ -2402,6 +2404,14 @@ 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
@@ -2414,7 +2424,13 @@ 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 = NodeSet.add (Callgraph.Type (mk_id "exception")) roots in
     let roots =
       Bindings.fold (fun typ_id _ roots -> NodeSet.add (Callgraph.Type typ_id) roots) (Env.get_enums ctx.tc_env) roots

src/lib/jib_compile.mli

@@ -85,6 +85,7 @@ type ctx = {
   locals : (mut * ctyp) Bindings.t;
   letbinds : int list;
   letbind_ids : IdSet.t;
+  unit_test_ids : IdSet.t;
   no_raw : bool;
   coverage_override : bool;
   def_annot : unit def_annot option;

src/lib/type_check.ml

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

@@ -2225,13 +2225,47 @@ let compile_ast env effect_info output_chan c_includes ast =
     let end_extern_cpp = separate hardline (List.map string [""; "#ifdef __cplusplus"; "}"; "#endif"]) in
     let hlhl = hardline ^^ 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
+
     Document.to_string
       (preamble ^^ hlhl ^^ docs ^^ hlhl
       ^^ ( if not !opt_no_rts then
              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
       )
     |> output_string output_chan
   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 = ()