refactor: switch to new API for recording deps (#9552)

Signed-off-by: Rudi Grinberg <me@rgrinberg.com>

src/dune_engine/action_builder.ml

@@ -1,20 +1,32 @@
 open Import
 
 type 'a eval_mode =
-  | Lazy : unit eval_mode
-  | Eager : Dep.Fact.t eval_mode
+  | Lazy : Dep.Set.t eval_mode
+  | Eager : Dep.Facts.t eval_mode
 
-type 'a thunk = { f : 'm. 'm eval_mode -> ('a * 'm Dep.Map.t) Memo.t } [@@unboxed]
+type 'a thunk = { f : 'm. 'm eval_mode -> ('a * 'm) Memo.t } [@@unboxed]
 
 module Deps_or_facts = struct
-  let union : type a. a eval_mode -> a Dep.Map.t -> a Dep.Map.t -> a Dep.Map.t =
+  let empty : type m. m eval_mode -> m = function
+    | Lazy -> Dep.Set.empty
+    | Eager -> Dep.Facts.empty
+  ;;
+
+  let return : type a m. a -> m eval_mode -> a * m = fun a mode -> a, empty mode
+
+  let union : type m. m eval_mode -> m -> m -> m =
     fun mode a b ->
     match mode with
     | Lazy -> Dep.Set.union a b
     | Eager -> Dep.Facts.union a b
   ;;
 
-  let union_all mode l = List.fold_left l ~init:Dep.Map.empty ~f:(union mode)
+  let union_all : type m. m eval_mode -> m list -> m =
+    fun mode list ->
+    match mode with
+    | Lazy -> Dep.Set.union_all list
+    | Eager -> Dep.Facts.union_all list
+  ;;
 end
 
 module T = struct
@@ -23,7 +35,7 @@ module T = struct
   module M = struct
     type 'a t = 'a thunk
 
-    let return x = { f = (fun _mode -> Memo.return (x, Dep.Map.empty)) }
+    let return x = { f = (fun mode -> Memo.return (Deps_or_facts.return x mode)) }
 
     let map t ~f =
       { f =
@@ -63,14 +75,59 @@ module T = struct
     }
   ;;
 
-  let of_memo m =
+  let of_memo memo =
     { f =
-        (fun _mode ->
-          let+ x = m in
-          x, Dep.Map.empty)
+        (fun mode ->
+          let+ x = memo in
+          x, Deps_or_facts.empty mode)
     }
   ;;
 
+  let record res (deps : Dep.Set.t) ~f =
+    let f : type m. m eval_mode -> (_ * m) Memo.t =
+      fun mode ->
+      let open Memo.O in
+      match mode with
+      | Lazy -> Memo.return (res, deps)
+      | Eager ->
+        let+ facts = Dep.Facts.record_facts deps ~f in
+        res, facts
+    in
+    { f }
+  ;;
+
+  let record_success memo =
+    let f : type m. m eval_mode -> (unit * m) Memo.t =
+      fun mode ->
+      let open Memo.O in
+      match mode with
+      | Lazy -> Memo.return ((), Dep.Set.empty)
+      | Eager ->
+        let+ () = memo in
+        (), Dep.Facts.empty
+    in
+    { f }
+  ;;
+
+  module Expert = struct
+    let record_dep_on_source_file_exn res ?loc (src_path : Path.Source.t) =
+      let f : type m. m eval_mode -> (_ * m) Memo.t =
+        fun mode ->
+        let (path : Path.t) = Path.source src_path in
+        let dep = Dep.file path in
+        match mode with
+        | Lazy -> Memo.return (res, Dep.Set.singleton dep)
+        | Eager ->
+          let open Memo.O in
+          let+ digest =
+            Fs_memo.file_digest_exn ?loc (Path.Outside_build_dir.In_source_dir src_path)
+          in
+          res, Dep.Facts.singleton dep (Dep.Fact.file path digest)
+      in
+      { f }
+    ;;
+  end
+
   module O = struct
     let ( >>> ) a b =
       { f =
@@ -127,11 +184,11 @@ let of_thunk t = t
 let run t mode = t.f mode
 
 let force_lazy_or_eager
-  : type a b.
-    a eval_mode
-    -> (b * Dep.Set.t) Memo.Lazy.t Lazy.t
-    -> (b * Dep.Facts.t) Memo.Lazy.t
-    -> (b * a Dep.Map.t) Memo.t
+  : type a m.
+    m eval_mode
+    -> (a * Dep.Set.t) Memo.Lazy.t Lazy.t
+    -> (a * Dep.Facts.t) Memo.Lazy.t
+    -> (a * m) Memo.t
   =
   fun mode lazy_ eager ->
   match mode with
@@ -221,7 +278,7 @@ let create_memo name ~input ?cutoff ?human_readable_description f =
   { lazy_; eager }
 ;;
 
-let exec_memo : type i o m. (i, o) memo -> i -> m eval_mode -> (o * m Dep.Map.t) Memo.t =
+let exec_memo : type i o m. (i, o) memo -> i -> m eval_mode -> (o * m) Memo.t =
   fun memo i mode ->
   match mode with
   | Lazy -> Memo.exec (Lazy.force memo.lazy_) i
@@ -234,8 +291,8 @@ let goal t =
   { f =
       (fun mode ->
         let open Memo.O in
-        let+ a, (_irrelevant_for_goals : _ Dep.Map.t) = t.f mode in
-        a, Dep.Map.empty)
+        let+ a, _facts_are_irrelevant_for_goals = t.f mode in
+        a, Deps_or_facts.empty mode)
   }
 ;;
 

src/dune_engine/action_builder.mli

@@ -65,7 +65,10 @@ val exec_memo : ('i, 'o) memo -> 'i -> 'o t
     but the contents of [p] is irrelevant. *)
 val goal : 'a t -> 'a t
 
-(** If you're thinking of using [Process.run] here, check that: (i) you don't in
+(** An action builder with no dependencies. Consider passing [Memo.of_thunk] to delay
+    forcing the computation until the action's dependencies need to be determined.
+
+    If you're thinking of using [Process.run] here, check that: (i) you don't in
     fact need [Command.run], and that (ii) [Process.run] only reads the declared
     build rule dependencies. *)
 val of_memo : 'a Memo.t -> 'a t
@@ -85,23 +88,47 @@ val of_memo_join : 'a t Memo.t -> 'a t
     dependencies, using [Eager] mode will increase parallelism. If you only want
     to know the set of dependencies, using [Lazy] will avoid unnecessary work. *)
 type 'a eval_mode =
-  | Lazy : unit eval_mode
-  | Eager : Dep.Fact.t eval_mode
+  | Lazy : Dep.Set.t eval_mode
+  | Eager : Dep.Facts.t eval_mode
 
 (** Execute an action builder. *)
-val run : 'a t -> 'b eval_mode -> ('a * 'b Dep.Map.t) Memo.t
+val run : 'a t -> 'b eval_mode -> ('a * 'b) Memo.t
 
 (** {1 Low-level} *)
 
-type 'a thunk = { f : 'm. 'm eval_mode -> ('a * 'm Dep.Map.t) Memo.t } [@@unboxed]
+type 'a thunk = { f : 'm. 'm eval_mode -> ('a * 'm) Memo.t } [@@unboxed]
 
 val of_thunk : 'a thunk -> 'a t
 
 module Deps_or_facts : sig
-  val union : 'a eval_mode -> 'a Dep.Map.t -> 'a Dep.Map.t -> 'a Dep.Map.t
-  val union_all : 'a eval_mode -> 'a Dep.Map.t list -> 'a Dep.Map.t
+  val union : 'm eval_mode -> 'm -> 'm -> 'm
+  val union_all : 'm eval_mode -> 'm list -> 'm
+end
+
+(** Record the given set as dependencies of the action produced by the action builder. *)
+val record : 'a -> Dep.Set.t -> f:(Dep.t -> Dep.Fact.t Memo.t) -> 'a t
+
+(** Record a given Memo computation as a "dependency" of the action builder, i.e., require
+    that it must succeed. Consider passing [Memo.of_thunk] to delay forcing the computation
+    until the action's dependencies need to be determined. *)
+val record_success : unit Memo.t -> unit t
+
+module Expert : sig
+  (** Like [record] but records a dependency on a *source* file. Evaluating the resulting
+      [t] in the [Eager] mode will raise a user error if the file can't be digested.
+
+      This function is in the [Expert] module because depending on files in the source
+      directory is usually a mistake. As of 2023-11-14, we use this function only for
+      setting up the rules that copy files from the source to the build directory. *)
+  val record_dep_on_source_file_exn
+    :  'a
+    -> ?loc:(unit -> Loc.t option Memo.t)
+    -> Path.Source.t
+    -> 'a t
 end
 
+(** {1 Evaluation} *)
+
 (** Evaluate a [t] and collect the set of its dependencies. This avoids doing the build
     work required for finding the facts about those dependencies, so you should use this
     function if you don't need the facts. *)

src/dune_engine/build_system.ml

@@ -131,8 +131,8 @@ module type Rec = sig
 
   val build_file : Path.t -> Digest.t Memo.t
   val build_dir : Path.t -> (Digest.t * Digest.t Targets.Produced.t) Memo.t
+  val build_dep : Dep.t -> Dep.Fact.t Memo.t
   val build_deps : Dep.Set.t -> Dep.Facts.t Memo.t
-  val eval_deps : 'a Action_builder.eval_mode -> Dep.Set.t -> 'a Dep.Map.t Memo.t
   val execute_rule : Rule.t -> rule_execution_result Memo.t
 
   val execute_action
@@ -200,13 +200,6 @@ end = struct
 
   let build_deps deps = Dep.Map.parallel_map deps ~f:(fun dep () -> build_dep dep)
 
-  let eval_deps : type a. a Action_builder.eval_mode -> Dep.Set.t -> a Dep.Map.t Memo.t =
-    fun mode deps ->
-    match mode with
-    | Lazy -> Memo.return deps
-    | Eager -> build_deps deps
-  ;;
-
   let select_sandbox_mode (config : Sandbox_config.t) ~loc ~sandboxing_preference =
     (* Rules with (mode patch-back-source-tree) are special and are not affected
        by sandboxing preferences. *)
@@ -663,12 +656,7 @@ end = struct
         ~info:(if Loc.is_none loc then Internal else From_dune_file loc)
         ~targets:(Targets.File.create target)
         ~mode:Standard
-        (Action_builder.of_thunk
-           { f =
-               (fun mode ->
-                 let+ deps = eval_deps mode deps in
-                 act.action, deps)
-           })
+        (Action_builder.record act.action deps ~f:build_dep)
     in
     let+ { deps = _; targets = _ } =
       execute_rule_impl
@@ -853,19 +841,15 @@ end = struct
               ]))
   ;;
 
-  let dep_on_anonymous_action (x : Rule.Anonymous_action.t Action_builder.t)
-    : _ Action_builder.t
+  let execute_anonymous_action action =
+    let* action, facts = Action_builder.evaluate_and_collect_facts action in
+    execute_action action ~observing_facts:facts
+  ;;
+
+  let dep_on_anonymous_action (action : Rule.Anonymous_action.t Action_builder.t)
+    : unit Action_builder.t
     =
-    Action_builder.of_thunk
-      { f =
-          (fun (type m) (mode : m Action_builder.eval_mode) ->
-            match mode with
-            | Lazy -> Memo.return ((), Dep.Map.empty)
-            | Eager ->
-              let* action, facts = Action_builder.evaluate_and_collect_facts x in
-              let+ () = execute_action action ~observing_facts:facts in
-              (), Dep.Map.empty)
-      }
+    Action_builder.record_success (Memo.of_thunk_apply execute_anonymous_action action)
   ;;
 
   let dep_on_alias_definition (definition : Rules.Dir_rules.Alias_spec.item) =
@@ -1036,6 +1020,7 @@ end
 
 include Exported
 
+let record_deps (deps : Dep.Set.t) = Action_builder.record () deps ~f:build_dep
 let eval_pred = Pred.eval
 let build_pred = Pred.build
 

src/dune_engine/build_system.mli

@@ -18,6 +18,9 @@ val file_exists : Path.t -> bool Memo.t
 (** Build a set of dependencies and return learned facts about them. *)
 val build_deps : Dep.Set.t -> Dep.Facts.t Memo.t
 
+(** Record the given set as dependencies of the action produced by the action builder. *)
+val record_deps : Dep.Set.t -> unit Action_builder.t
+
 (** [eval_pred glob] returns the set of filenames in [File_selector.dir glob] that matches
     [File_selector.predicate glob], including both sources and generated files. *)
 val eval_pred : File_selector.t -> Filename_set.t Memo.t

src/dune_engine/load_rules.ml

@@ -267,15 +267,6 @@ let no_rule_found ~loc fn =
       [ "fn", Path.Build.to_dyn fn ]
 ;;
 
-let eval_source_file : type a. a Action_builder.eval_mode -> Path.Source.t -> a Memo.t =
-  fun mode path ->
-  match mode with
-  | Lazy -> Memo.return ()
-  | Eager ->
-    let+ d = Fs_memo.file_digest_exn (In_source_dir path) in
-    Dep.Fact.file (Path.source path) d
-;;
-
 module rec Load_rules : sig
   val load_dir : dir:Path.t -> Loaded.t Memo.t
   val is_under_directory_target : Path.t -> bool Memo.t
@@ -286,28 +277,28 @@ module rec Load_rules : sig
 end = struct
   open Load_rules
 
+  let copy_source_action ~src_path ~build_path : Action.Full.t Action_builder.t =
+    let action =
+      Action.Full.make
+        (Action.copy (Path.source src_path) build_path)
+        (* Sandboxing this action doesn't make much sense: if we can copy [src_path] to
+           the sandbox, we might as well copy it to the build directory directly. *)
+        ~sandbox:Sandbox_config.no_sandboxing
+    in
+    Action_builder.Expert.record_dep_on_source_file_exn
+      action
+      ~loc:Current_rule_loc.get
+      src_path
+  ;;
+
   let create_copy_rules ~dir ~ctx_dir ~non_target_source_filenames =
     Filename.Set.to_list_map non_target_source_filenames ~f:(fun filename ->
-      let path = Path.Source.relative dir filename in
-      let ctx_path = Path.Build.append_source ctx_dir path in
-      let build =
-        Action_builder.of_thunk
-          { f =
-              (fun mode ->
-                let+ fact = eval_source_file mode path in
-                let path = Path.source path in
-                ( Action.Full.make
-                    (Action.copy path ctx_path)
-                    (* There's an [assert false] in [prepare_managed_paths]
-                       that blows up if we try to sandbox this. *)
-                    ~sandbox:Sandbox_config.no_sandboxing
-                , Dep.Map.singleton (Dep.file path) fact ))
-          }
-      in
+      let src_path = Path.Source.relative dir filename in
+      let build_path = Path.Build.append_source ctx_dir src_path in
       Rule.make
-        ~info:(Source_file_copy path)
-        ~targets:(Targets.File.create ctx_path)
-        build)
+        ~info:(Source_file_copy src_path)
+        ~targets:(Targets.File.create build_path)
+        (copy_source_action ~src_path ~build_path))
   ;;
 
   let compile_rules ~dir ~source_dirs rules =

src/dune_rules/action_builder.ml

@@ -3,7 +3,7 @@ include Dune_engine.Action_builder
 open O
 module With_targets = With_targets
 
-let register_action_deps : type a. a eval_mode -> Dep.Set.t -> a Dep.Map.t Memo.t =
+let register_action_deps : type a. a eval_mode -> Dep.Set.t -> a Memo.t =
   fun mode deps ->
   match mode with
   | Eager -> Build_system.build_deps deps
@@ -42,13 +42,8 @@ let dyn_paths paths = dyn_deps (paths >>| fun (x, paths) -> x, Dep.Set.of_files
 let dyn_paths_unit paths = dyn_deps (paths >>| fun paths -> (), Dep.Set.of_files paths)
 
 let contents p =
-  of_thunk
-    { f =
-        (fun _mode ->
-          let open Memo.O in
-          let+ x = Build_system.read_file p in
-          x, Dep.Map.empty)
-    }
+  let* () = Dep.file p |> Dep.Set.singleton |> Build_system.record_deps in
+  of_memo (Build_system.read_file p)
 ;;
 
 let lines_of p = contents p >>| String.split_lines
@@ -129,8 +124,7 @@ let paths_existing paths =
        if_file_exists file ~then_:(path file) ~else_:(return ())))
 ;;
 
-let paths_matching
-  : type a. File_selector.t -> a eval_mode -> (Filename_set.t * a Dep.Map.t) Memo.t
+let paths_matching : type a. File_selector.t -> a eval_mode -> (Filename_set.t * a) Memo.t
   =
   fun g mode ->
   let open Memo.O in