From fff3a551bfaf3e8dfdc57dcd7b5f84ca15260a57 Mon Sep 17 00:00:00 2001
From: Antonio Nuno Monteiro <anmonteiro@gmail.com>
Date: Fri, 21 Aug 2026 15:43:36 -0700
Subject: [PATCH] perf(memo): flatten dependency sections directly

Signed-off-by: Antonio Nuno Monteiro <anmonteiro@gmail.com>
---
 src/memo/deps.ml                              | 34 +++++++++++--------
 test/expect-tests/memo/deps_representation.ml | 28 +++++++++------
 2 files changed, 37 insertions(+), 25 deletions(-)

diff --git a/src/memo/deps.ml b/src/memo/deps.ml
index a180070d1e5..add1a414b49 100644
--- a/src/memo/deps.ml
+++ b/src/memo/deps.ml
@@ -15,14 +15,15 @@ module Static = struct
     | Seq of 'node t Array.Immutable.t
     | Par of 'node t Array.Immutable.t
 
-  (* Flatten a chronological list of sections into a sequence, flattening nested [Seq]s:
-     (x ; (y ; z)) = (x ; y ; z). *)
-  let flatten_seqs (sections : 'node t list) : 'node t =
+  (* Flatten a reverse-chronological list of sections into a sequence, flattening nested
+     [Seq]s: (x ; (y ; z)) = (x ; y ; z). *)
+  let flatten_rev_seqs (sections : 'node t list) : 'node t =
     let flat =
-      List.concat_map sections ~f:(function
-        | Empty -> []
-        | Seq arr -> Array.Immutable.to_list arr
-        | (Singleton _ | Par _) as t -> [ t ])
+      List.fold_left sections ~init:[] ~f:(fun acc section ->
+        match section with
+        | Empty -> acc
+        | Seq arr -> Array.Immutable.fold_right arr ~init:acc ~f:List.cons
+        | (Singleton _ | Par _) as section -> section :: acc)
     in
     match flat with
     | [] -> Empty
@@ -38,13 +39,13 @@ module Static = struct
       if i < 0
       then acc
       else (
-        let elements =
+        let acc =
           match f i with
-          | Empty -> []
-          | Par arr -> Array.Immutable.to_list arr
-          | (Singleton _ | Seq _) as t -> [ t ]
+          | Empty -> acc
+          | Par arr -> Array.Immutable.fold_right arr ~init:acc ~f:List.cons
+          | (Singleton _ | Seq _) as section -> section :: acc
         in
-        loop (i - 1) (elements @ acc))
+        loop (i - 1) acc)
     in
     match loop (num_threads - 1) [] with
     | [] -> Empty
@@ -79,9 +80,12 @@ module Dynamic = struct
     | section -> section :: t
   ;;
 
-  (* The list is most-recent-first, so reverse it to chronological order before flattening
-     the sequence. *)
-  let to_static (t : 'node t) : 'node Static.t = Static.flatten_seqs (List.rev t)
+  let to_static (t : 'node t) : 'node Static.t =
+    match t with
+    | [] -> Static.Empty
+    | [ section ] -> section
+    | _ :: _ :: _ -> Static.flatten_rev_seqs t
+  ;;
 end
 
 (* Note that dependencies should be checked in the order in which they were depended on to
diff --git a/test/expect-tests/memo/deps_representation.ml b/test/expect-tests/memo/deps_representation.ml
index 1d103fc4c26..552ee236fb3 100644
--- a/test/expect-tests/memo/deps_representation.ml
+++ b/test/expect-tests/memo/deps_representation.ml
@@ -52,20 +52,28 @@ let%expect_test "dependency structure of Memo combinators" =
       ; Singleton (Some "c", ())
       ]
     |}];
-  (* A single-thread parallel section is not wrapped in [Par], and nested [Seq]s are
-     flattened, so the result is one flat [Seq] in chronological order. *)
+  (* Nested sequential and parallel sections are flattened without changing their order. *)
+  let parallel ms = Memo.map (Memo.all_concurrently ms) ~f:ignore in
   print_deps
-    "flatten_seqs"
+    "flatten_nested_seq"
     (let* () = read a in
-     let* (_ : unit list) =
-       Memo.parallel_map [ () ] ~f:(fun () ->
-         let* () = read b in
-         read c)
-     in
-     Memo.return ());
+     parallel
+       [ (let* () = read b in
+          read c)
+       ; Memo.return ()
+       ]);
   [%expect
     {|
-    flatten_seqs: Seq
+    flatten_nested_seq: Seq
+      [ Singleton (Some "a", ())
+      ; Singleton (Some "b", ())
+      ; Singleton (Some "c", ())
+      ]
+    |}];
+  print_deps "flatten_nested_par" (parallel [ parallel [ read a; read b ]; read c ]);
+  [%expect
+    {|
+    flatten_nested_par: Par
       [ Singleton (Some "a", ())
       ; Singleton (Some "b", ())
       ; Singleton (Some "c", ())