From 392201445bd3a210bf601f1f73ace73f37e3c0b8 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Mon, 12 Jun 2023 22:07:36 -0400
Subject: [PATCH] test: add t_tree_futs

exercises futures pretty hard.
---
 test/dune           |  4 +--
 test/t_tree_futs.ml | 67 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 test/t_tree_futs.ml

diff --git a/test/dune b/test/dune
index 490da383..2b2e807b 100644
--- a/test/dune
+++ b/test/dune
@@ -1,3 +1,3 @@
 (tests
- (names t_fib t_bench1 t_fib_rec t_futs1)
- (libraries moonpool))
+ (names t_fib t_bench1 t_fib_rec t_futs1 t_tree_futs)
+ (libraries moonpool trace tracy-client.trace))
diff --git a/test/t_tree_futs.ml b/test/t_tree_futs.ml
new file mode 100644
index 00000000..4fba1837
--- /dev/null
+++ b/test/t_tree_futs.ml
@@ -0,0 +1,67 @@
+open Moonpool
+
+let ( let@ ) = ( @@ )
+
+type 'a tree =
+  | Leaf of 'a
+  | Node of 'a tree Fut.t * 'a tree Fut.t
+
+let rec mk_tree ~pool n : _ tree Fut.t =
+  let@ _sp = Trace.with_span ~__FILE__ ~__LINE__ "mk-tree" in
+  if n <= 1 then
+    Fut.return (Leaf 1)
+  else
+    let open (val Fut.infix pool) in
+    let l =
+      Fut.spawn ~on:pool (fun () -> mk_tree ~pool (n - 1)) |> Fut.join ~on:pool
+    and r =
+      Fut.spawn ~on:pool (fun () -> mk_tree ~pool (n - 1)) |> Fut.join ~on:pool
+    in
+
+    Fut.return @@ Node (l, r)
+
+let rec rev ~pool (t : 'a tree Fut.t) : 'a tree Fut.t =
+  let@ _sp = Trace.with_span ~__FILE__ ~__LINE__ "rev" in
+  let open (val Fut.infix pool) in
+  t >>= function
+  | Leaf n -> Fut.return (Leaf n)
+  | Node (l, r) ->
+    let l = rev ~pool l and r = rev ~pool r in
+    Fut.spawn ~on:pool (fun () -> Node (r, l))
+
+let rec sum ~pool (t : int tree Fut.t) : int Fut.t =
+  let@ _sp = Trace.with_span ~__FILE__ ~__LINE__ "sum" in
+  let open (val Fut.infix pool) in
+  t >>= function
+  | Leaf n -> Fut.return n
+  | Node (l, r) ->
+    let* l = sum ~pool l and* r = sum ~pool r in
+    Fut.spawn ~on:pool (fun () -> l + r)
+
+let run ~pool n : (int * int) Fut.t =
+  let@ _sp = Trace.with_span ~__FILE__ ~__LINE__ "run" in
+  let open (val Fut.infix pool) in
+  let t = Fut.return n >>= mk_tree ~pool in
+  let t' = rev ~pool t in
+  let sum_t = sum ~pool t in
+  let sum_t' = sum ~pool t' in
+  Fut.both sum_t sum_t'
+
+let default_n = 16
+
+let () =
+  (*
+  Tracy_client_trace.setup ();
+   *)
+  let pool = Pool.create ~per_domain:1 ~min:2 () in
+
+  let start = Unix.gettimeofday () in
+  let n = try int_of_string (Sys.getenv "N") with _ -> default_n in
+
+  Printf.printf "n=%d\n%!" n;
+  let n1, n2 = run ~pool n |> Fut.wait_block_exn in
+  Printf.printf "n: %d, n': %d (in %.2fs)\n%!" n1 n2
+    (Unix.gettimeofday () -. start);
+  assert (n1 = 1 lsl (n - 1));
+  assert (n2 = 1 lsl (n - 1));
+  ()