perf(memo): avoid fibers for current sequential dependencies

bench/micro/memo_bench/benchmarks.ml

@@ -118,6 +118,24 @@ let%bench_fun "20-reads (create and compute)" = twenty_reads.create_and_compute
 let%bench_fun "20-reads (incr and recompute)" = twenty_reads.incr_and_recompute ()
 let%bench_fun "20-reads (restore from cache)" = twenty_reads.restore_from_cache ()
 
+let many_sequential_deps =
+  Case.create (fun () ->
+    let v = Var.create 0 in
+    let deps =
+      Memo.memoize
+        (List.fold_left
+           ~init:(Memo.return 0)
+           (List.init 1_000 ~f:(fun _i -> ()))
+           ~f:(fun acc () ->
+             Memo.bind acc ~f:(fun acc -> Memo.map (Var.read v) ~f:(fun v -> acc + v))))
+    in
+    v, Memo.bind (Var.read v) ~f:(fun _ -> deps))
+;;
+
+let%bench_fun "1000 repeated current deps (restore from cache)" =
+  many_sequential_deps.restore_from_cache ()
+;;
+
 let clique =
   Case.create (fun () ->
     let v = Var.create 0 in

src/memo/deps.ml

@@ -79,6 +79,19 @@ module Dynamic = struct
   let to_static (t : 'node t) : 'node Static.t = Static.flatten_seqs (List.rev t)
 end
 
+module Check = struct
+  type 'cycle t =
+    | Unchanged
+    | Changed
+    | Deferred of 'cycle Changed_or_not.t Fiber.t
+
+  let[@inline] to_fiber = function
+    | Unchanged -> Fiber.return Changed_or_not.Unchanged
+    | Changed -> Fiber.return Changed_or_not.Changed
+    | Deferred fiber -> fiber
+  ;;
+end
+
 (* Note that dependencies should be checked in the order in which they were depended on to
    avoid recomputations of dependencies that are no longer relevant, and to eliminate
    spurious dependency cycles. This is why [changed_or_not] checks sequential sections in
@@ -89,7 +102,7 @@ let changed_or_not (t : 'node t) ~f =
     | Empty -> Fiber.return Changed_or_not.Unchanged
     | Singleton node ->
       Counter.add Metrics.Restore.edges 1;
-      f ~ok_to_recompute_eagerly node
+      Check.to_fiber (f ~ok_to_recompute_eagerly node)
     | Seq arr -> seq arr 0
     | Par arr ->
       Fiber.map_reduce_array
@@ -100,16 +113,23 @@ let changed_or_not (t : 'node t) ~f =
           match section with
           | Static.Singleton node ->
             Counter.add Metrics.Restore.edges 1;
-            f ~ok_to_recompute_eagerly:true node
+            Check.to_fiber (f ~ok_to_recompute_eagerly:true node)
           | other -> loop ~ok_to_recompute_eagerly:false other)
   and seq arr index =
     if index < Array.Immutable.length arr
-    then
-      loop ~ok_to_recompute_eagerly:false (Array.Immutable.get arr index)
-      >>= function
-      | Changed_or_not.Unchanged -> seq arr (index + 1)
-      | (Changed | Cancelled _) as res -> Fiber.return res
+    then (
+      match Array.Immutable.get arr index with
+      | Static.Singleton node ->
+        Counter.add Metrics.Restore.edges 1;
+        (match f ~ok_to_recompute_eagerly:false node with
+         | Check.Unchanged -> seq arr (index + 1)
+         | Check.Changed -> Fiber.return Changed_or_not.Changed
+         | Check.Deferred fiber -> fiber >>= continue_seq arr (index + 1))
+      | other -> loop ~ok_to_recompute_eagerly:false other >>= continue_seq arr (index + 1))
     else Fiber.return Changed_or_not.Unchanged
+  and continue_seq arr index = function
+    | Changed_or_not.Unchanged -> seq arr index
+    | (Changed | Cancelled _) as result -> Fiber.return result
   in
   loop ~ok_to_recompute_eagerly:false t
 ;;

src/memo/deps.mli

@@ -11,6 +11,13 @@ type 'node t
 
 val empty : 'node t
 
+module Check : sig
+  type 'cycle t =
+    | Unchanged
+    | Changed
+    | Deferred of 'cycle Changed_or_not.t Fiber.t
+end
+
 (** Like [t] but supports cheap appending of new dependencies. *)
 module Dynamic : sig
   type 'node static := 'node t
@@ -33,12 +40,14 @@ end
     in order and the check stops early at the first [Changed]/[Cancelled]; parallel
     sections are checked concurrently and their results combined.
 
-    [ok_to_recompute_eagerly] is [true] when the dependency is a direct child of a parallel
-    section, telling [f] that it may eagerly recompute a dependency without a cutoff in
-    parallel with its siblings (instead of deferring the recomputation). *)
+    [f] can return [Unchanged] or [Changed] directly when checking a dependency requires no
+    fiber, or [Deferred fiber] otherwise. [ok_to_recompute_eagerly] is [true] when the
+    dependency is a direct child of a parallel section, telling [f] that it may eagerly
+    recompute a dependency without a cutoff in parallel with its siblings (instead of
+    deferring the recomputation). *)
 val changed_or_not
   :  'node t
-  -> f:(ok_to_recompute_eagerly:bool -> 'node -> 'cycle Changed_or_not.t Fiber.t)
+  -> f:(ok_to_recompute_eagerly:bool -> 'node -> 'cycle Check.t)
   -> 'cycle Changed_or_not.t Fiber.t
 
 module For_debugging : sig

src/memo/exec.ml

@@ -43,11 +43,15 @@ let cancelled ~dependency_cycle : Collect_errors_monoid.t =
   }
 ;;
 
+let dep_has_changed ~(node : _ Dep_node.t) ~(dep : _ Dep_node.t) =
+  match Run.compare (Dep_node.last_changed_at dep) (Dep_node.last_validated_at node) with
+  | Gt -> true
+  | Eq | Lt -> false
+;;
+
 (* [Changed] if [dep] is newer than [node] and [Unchanged] otherwise. *)
 let dep_changed_or_not ~(node : _ Dep_node.t) ~(dep : _ Dep_node.t) : _ Changed_or_not.t =
-  match Run.compare (Dep_node.last_changed_at dep) (Dep_node.last_validated_at node) with
-  | Gt -> Changed
-  | Eq | Lt -> Unchanged
+  if dep_has_changed ~node ~dep then Changed else Unchanged
 ;;
 
 let rec restore_from_cache
@@ -75,45 +79,48 @@ let rec restore_from_cache
       node.deps
       ~f:(fun[@inline] ~ok_to_recompute_eagerly (Dep_node.T dep) ->
         (* If the [Run.is_current] check succeeds then the node must have been [Cached] in
-           the current run, so there is no need to restore it (which would allocate a
-           fiber). We can compare the timestamps directly. *)
+           the current run, so there is no need to restore it or allocate a fiber. We can
+           compare the timestamps directly. *)
         if Run.is_current (Dep_node.last_validated_at dep)
-        then Fiber.return (dep_changed_or_not ~node ~dep)
+        then
+          if dep_has_changed ~node ~dep then Deps.Check.Changed else Deps.Check.Unchanged
         else
-          consider_and_restore_from_cache_without_adding_dep dep
-          >>= function
-          | Unchanged ->
-            (* Here [dep_changed_or_not] can return [Changed] if the [node] was skipped in
-               the previous run, i.e., it was unreachable, while the [dep] wasn't skipped
-               and changed. *)
-            Fiber.return (dep_changed_or_not ~node ~dep)
-          | Cancelled { dependency_cycle } ->
-            Fiber.return (Changed_or_not.Cancelled { dependency_cycle })
-          | Changed ->
-            (match Spec.has_cutoff dep.spec with
-             | false when not ok_to_recompute_eagerly ->
-               (* If [dep] has no cutoff and [ok_to_recompute_eagerly] is not set, it is
-                  sufficient to check whether [dep] is up to date. We are in the [Changed]
-                  branch, which means [dep] is not up to date, and we therefore must
-                  recompute the [node]. *)
-               Fiber.return Changed_or_not.Changed
-             | _ ->
-               (* If [dep] has a cutoff predicate, it is not sufficient to check whether it
-                  is up to date: even if it isn't, after we recompute it, the resulting
-                  value may remain unchanged, allowing us to skip recomputing the [node].
+          Deps.Check.Deferred
+            (consider_and_restore_from_cache_without_adding_dep dep
+             >>= function
+             | Unchanged ->
+               (* Here [dep_changed_or_not] can return [Changed] if the [node] was skipped
+                  in the previous run, i.e., it was unreachable, while the [dep] wasn't
+                  skipped and changed. *)
+               Fiber.return (dep_changed_or_not ~node ~dep)
+             | Cancelled { dependency_cycle } ->
+               Fiber.return (Changed_or_not.Cancelled { dependency_cycle })
+             | Changed ->
+               (match Spec.has_cutoff dep.spec with
+                | false when not ok_to_recompute_eagerly ->
+                  (* If [dep] has no cutoff and [ok_to_recompute_eagerly] is not set, it is
+                     sufficient to check whether [dep] is up to date. We are in the
+                     [Changed] branch, which means [dep] is not up to date, and we therefore
+                     must recompute the [node]. *)
+                  Fiber.return Changed_or_not.Changed
+                | _ ->
+                  (* If [dep] has a cutoff predicate, it is not sufficient to check whether
+                     it is up to date: even if it isn't, after we recompute it, the resulting
+                     value may remain unchanged, allowing us to skip recomputing the [node].
 
-                  If [dep] has no cutoff but [ok_to_recompute_eagerly] is set (which could
-                  happen if [dep] is a direct child of a [Par] node), we eagerly recompute
-                  [dep] so that its computation runs in parallel with its siblings, instead
-                  of being deferred to the compute phase where it might run sequentially. We
-                  still report [Changed] in this case since there is no cutoff to check. *)
-               consider_and_compute_without_adding_dep dep
-               >>| (function
-                | Ok () ->
-                  (match Spec.has_cutoff dep.spec with
-                   | false -> Changed_or_not.Changed
-                   | true -> dep_changed_or_not ~node ~dep)
-                | Error dependency_cycle -> Cancelled { dependency_cycle })))
+                     If [dep] has no cutoff but [ok_to_recompute_eagerly] is set (which could
+                     happen if [dep] is a direct child of a [Par] node), we eagerly recompute
+                     [dep] so that its computation runs in parallel with its siblings,
+                     instead of being deferred to the compute phase where it might run
+                     sequentially. We still report [Changed] in this case since there is no
+                     cutoff to check. *)
+                  consider_and_compute_without_adding_dep dep
+                  >>| (function
+                   | Ok () ->
+                     (match Spec.has_cutoff dep.spec with
+                      | false -> Changed_or_not.Changed
+                      | true -> dep_changed_or_not ~node ~dep)
+                   | Error dependency_cycle -> Cancelled { dependency_cycle }))))
 
 and compute : 'i 'o. ('i, 'o) Dep_node.t -> unit Fiber.t =
   fun node ->

test/expect-tests/memo/deps_representation.ml

@@ -79,3 +79,62 @@ let%expect_test "dependency structure of Memo combinators" =
   print_deps "nested_empty" (seq (par (seq (par e))));
   [%expect {| nested_empty: Empty |}]
 ;;
+
+let%expect_test "restore repeated sequential dependencies" =
+  let value = ref 1 in
+  let leaf =
+    Memo.lazy_node ~name:"leaf" ~cutoff:Int.equal (fun () ->
+      printfn "compute leaf";
+      Memo.return !value)
+  in
+  let top =
+    Memo.lazy_node ~name:"top" ~cutoff:Int.equal (fun () ->
+      printfn "compute top";
+      let rec loop count sum =
+        match count with
+        | 0 -> Memo.return sum
+        | count ->
+          let* value = read leaf in
+          loop (count - 1) (sum + value)
+      in
+      loop 3 0)
+  in
+  let run_top () = printfn "top = %d" (run (read top)) in
+  run_top ();
+  [%expect
+    {|
+    compute top
+    compute leaf
+    top = 3
+    |}];
+  Memo.reset Memo.Invalidation.empty;
+  run_top ();
+  [%expect {| top = 3 |}];
+  Memo.reset (Memo.Node.invalidate ~reason:Memo.Invalidation.Reason.Test leaf);
+  run_top ();
+  [%expect
+    {|
+    compute leaf
+    top = 3
+    |}];
+  value := 2;
+  Memo.reset (Memo.Node.invalidate ~reason:Memo.Invalidation.Reason.Test leaf);
+  run_top ();
+  [%expect
+    {|
+    compute leaf
+    compute top
+    top = 6
+    |}];
+  value := 3;
+  Memo.reset (Memo.Node.invalidate ~reason:Memo.Invalidation.Reason.Test leaf);
+  printfn "leaf = %d" (run (read leaf));
+  run_top ();
+  [%expect
+    {|
+    compute leaf
+    leaf = 3
+    compute top
+    top = 9
+    |}]
+;;