unikernel/duniverse/base/test/allocation/test_hashtbl_allocation.ml

@@ -0,0 +1,90 @@
+open! Base
+open Expect_test_helpers_core
+
+let () = Int_conversions.sexp_of_int_style := `Underscores
+
+let%expect_test "find_and_call_1_and_2" =
+  let test x =
+    let t = Hashtbl.create (module Int) ~size:16 ~growth_allowed:false in
+    for i = 0 to x - 1 do
+      Hashtbl.add_exn t ~key:i ~data:(i * 7)
+    done;
+    let if_found a b = assert (a = b) in
+    let if_not_found a b =
+      assert (a = x);
+      assert (b = x * 7)
+    in
+    require_no_allocation [%here] (fun () ->
+      for i = 0 to x do
+        Hashtbl.find_and_call1 t i ~a:(i * 7) ~if_found ~if_not_found
+      done);
+    let if_found ~key ~data:a b =
+      assert (a = b);
+      assert (key = a / 7)
+    in
+    let if_not_found a b =
+      assert (a = x);
+      assert (b = x * 7)
+    in
+    require_no_allocation [%here] (fun () ->
+      for i = 0 to x do
+        Hashtbl.findi_and_call1 t i ~a:(i * 7) ~if_found ~if_not_found
+      done);
+    let if_found a b c =
+      assert (a = b);
+      assert (b = c / 2)
+    in
+    let if_not_found a b c =
+      assert (a = x);
+      assert (b = x * 7);
+      assert (c = x * 14)
+    in
+    require_no_allocation [%here] (fun () ->
+      for i = 0 to x do
+        Hashtbl.find_and_call2 t i ~a:(i * 7) ~b:(i * 14) ~if_found ~if_not_found
+      done);
+    let if_found ~key ~data:a b c =
+      assert (a = b);
+      assert (b = c / 2);
+      assert (key = a / 7)
+    in
+    let if_not_found a b c =
+      assert (a = x);
+      assert (b = x * 7);
+      assert (c = x * 14)
+    in
+    require_no_allocation [%here] (fun () ->
+      for i = 0 to x do
+        Hashtbl.findi_and_call2 t i ~a:(i * 7) ~b:(i * 14) ~if_found ~if_not_found
+      done);
+    print_s (Int.sexp_of_t x)
+  in
+  (* try various load factors, to exercise all branches of matching on the structure of
+     the avl tree *)
+  test 1;
+  test 3;
+  test 10;
+  test 17;
+  test 25;
+  test 29;
+  test 33;
+  test 3133;
+  [%expect {|
+    1
+    3
+    10
+    17
+    25
+    29
+    33
+    3_133
+    |}]
+;;
+
+let%expect_test ("find_or_add shouldn't allocate" [@tags "no-js"]) =
+  let default = Fn.const () in
+  let t = Hashtbl.create (module Int) ~size:16 ~growth_allowed:false in
+  Hashtbl.add_exn t ~key:100 ~data:();
+  require_no_allocation [%here] (fun () -> Hashtbl.find_or_add t 100 ~default);
+  [%expect {| |}]
+;;