perf(pkg): index immediate revision directory entries

bench/micro/dune

@@ -55,3 +55,17 @@
  (allow_overlapping_dependencies)
  (modules path_bench_main)
  (libraries path_bench core_bench.inline_benchmarks))
+
+(library
+ (name rev_store_bench)
+ (modules rev_store_bench)
+ (library_flags -linkall)
+ (preprocess
+  (pps ppx_bench))
+ (libraries dune_pkg stdune core_bench.inline_benchmarks))
+
+(executable
+ (name rev_store_bench_main)
+ (allow_overlapping_dependencies)
+ (modules rev_store_bench_main)
+ (libraries rev_store_bench core_bench.inline_benchmarks))

bench/micro/rev_store_bench.ml

@@ -0,0 +1,23 @@
+open Stdune
+module At_rev = Dune_pkg.Rev_store.At_rev
+
+let make_paths file_count =
+  List.init file_count ~f:(fun i ->
+    let directory = i mod 100 in
+    Path.Local.of_string (sprintf "dir-%03d/file-%06d" directory i))
+;;
+
+let%bench_fun
+    ("directory_entries create" [@indexed file_count = [ 1_000; 10_000; 100_000 ]])
+  =
+  let paths = make_paths file_count in
+  fun () -> ignore (At_rev.For_tests.make_directory_entries paths)
+;;
+
+let%bench_fun
+    ("directory_entries immediate" [@indexed file_count = [ 1_000; 10_000; 100_000 ]])
+  =
+  let entries = make_paths file_count |> At_rev.For_tests.make_directory_entries in
+  let path = Path.Local.of_string "dir-000" in
+  fun () -> ignore (At_rev.For_tests.directory_entries entries ~recursive:false path)
+;;

bench/micro/rev_store_bench_main.ml

@@ -0,0 +1 @@
+Inline_benchmarks_public.Runner.main ~libname:"rev_store_bench"

bench/micro/runner.sh

@@ -7,6 +7,7 @@ declare -A benchmarks=(
   [thread_pool]="thread_pool_bench:thread_pool_bench_main"
   [digest]="digest_bench:digest_bench_main"
   [path]="path_bench:path_bench_main"
+  [rev_store]="rev_store_bench:rev_store_bench_main"
 )
 
 if [[ -z "${benchmarks[$1]}" ]]; then

src/dune_pkg/rev_store.ml

@@ -838,11 +838,60 @@ end
 module At_rev = struct
   type repo = t
 
+  module Directory_entries = struct
+    type directory =
+      { immediate : File.Set.t
+      ; recursive : File.Set.t
+      }
+
+    type t = directory Path.Local.Table.t
+
+    let empty = { immediate = File.Set.empty; recursive = File.Set.empty }
+
+    let create files =
+      let directories = Path.Local.Table.create (File.Set.cardinal files) in
+      File.Set.iter files ~f:(fun file ->
+        match File.path file |> Path.Local.parent with
+        | None -> ()
+        | Some parent ->
+          let { immediate; recursive } =
+            Path.Local.Table.find_or_add directories parent ~f:(Fun.const empty)
+          in
+          Path.Local.Table.set
+            directories
+            parent
+            { immediate = File.Set.add immediate file
+            ; recursive = File.Set.add recursive file
+            };
+          let rec add_to_ancestors = function
+            | None -> ()
+            | Some ancestor ->
+              let { immediate; recursive } =
+                Path.Local.Table.find_or_add directories ancestor ~f:(Fun.const empty)
+              in
+              Path.Local.Table.set
+                directories
+                ancestor
+                { immediate; recursive = File.Set.add recursive file };
+              add_to_ancestors (Path.Local.parent ancestor)
+          in
+          add_to_ancestors (Path.Local.parent parent));
+      directories
+    ;;
+
+    let find directories ~recursive path =
+      match Path.Local.Table.find directories path with
+      | None -> File.Set.empty
+      | Some { immediate; recursive = recursive_entries } ->
+        if recursive then recursive_entries else immediate
+    ;;
+  end
+
   type t =
     { repo : repo
     ; revision : Object.t
     ; files : File.Set.t
-    ; recursive_directory_entries : File.Set.t Path.Local.Table.t
+    ; directory_entries : Directory_entries.t
     ; submodules : Object.t Path.Local.Map.t
     }
 
@@ -1021,74 +1070,33 @@ module At_rev = struct
       >>| List.cons files
       >>| File.Set.union_all
     in
-    let recursive_directory_entries =
-      let recursive_directory_entries =
-        Path.Local.Table.create (File.Set.cardinal files)
-      in
-      (* Build a table mapping each directory path to the set of files under it
-         in the directory hierarchy. *)
-      File.Set.iter files ~f:(fun file ->
-        (* Add [file] to the set of files under each directory which is an
-           ancestor of [file]. *)
-        let rec loop = function
-          | None -> ()
-          | Some parent ->
-            let recursive_directory_entries_of_parent =
-              Path.Local.Table.find_or_add
-                recursive_directory_entries
-                parent
-                ~f:(Fun.const File.Set.empty)
-            in
-            let recursive_directory_entries_of_parent =
-              File.Set.add recursive_directory_entries_of_parent file
-            in
-            Path.Local.Table.set
-              recursive_directory_entries
-              parent
-              recursive_directory_entries_of_parent;
-            loop (Path.Local.parent parent)
-        in
-        loop (File.path file |> Path.Local.parent));
-      recursive_directory_entries
-    in
-    { repo; revision; files; recursive_directory_entries; submodules = commit_paths }
+    let directory_entries = Directory_entries.create files in
+    { repo; revision; files; directory_entries; submodules = commit_paths }
   ;;
 
-  let content
-        { repo; revision; files = _; recursive_directory_entries = _; submodules = _ }
-        path
-    =
+  let content { repo; revision; files = _; directory_entries = _; submodules = _ } path =
     show repo [ `Path (revision, path) ]
   ;;
 
-  let directory_entries_recursive t path =
-    Path.Local.Table.find t.recursive_directory_entries path
-    |> Option.value ~default:File.Set.empty
+  let directory_entries { directory_entries; _ } ~recursive path =
+    Directory_entries.find directory_entries ~recursive path
   ;;
 
-  let directory_entries_immediate t path =
-    (* TODO: there are much better ways of implementing this:
-       1. using libgit or ocamlgit
-       2. possibly using [$ git archive] *)
-    File.Set.filter t.files ~f:(fun (file : File.t) ->
-      match Path.Local.parent (File.path file) with
-      | None -> false
-      | Some p -> Path.Local.equal p path)
-  ;;
+  module For_tests = struct
+    type directory_entries = Directory_entries.t
 
-  let directory_entries t ~recursive path =
-    (if recursive then directory_entries_recursive else directory_entries_immediate)
-      t
-      path
-  ;;
+    let make_directory_entries paths =
+      let hash = Object.of_sha1_unsafe (String.make 40 '0') in
+      List.fold_left paths ~init:File.Set.empty ~f:(fun files path ->
+        File.Set.add files (File.Direct { path; size = 0; hash }))
+      |> Directory_entries.create
+    ;;
+
+    let directory_entries = Directory_entries.find
+  end
 
   let check_out
-        { repo = { dir; _ }
-        ; revision
-        ; files = _
-        ; recursive_directory_entries = _
-        ; submodules
-        }
+        { repo = { dir; _ }; revision; files = _; directory_entries = _; submodules }
         ~target
     =
     let git = git () in

src/dune_pkg/rev_store.mli

@@ -92,6 +92,19 @@ module At_rev : sig
   module Config : sig
     val parse : string -> (string * string option * string * string) option
   end
+
+  (** Build directory entries without creating a Git repository. *)
+  module For_tests : sig
+    type directory_entries
+
+    val make_directory_entries : Path.Local.t list -> directory_entries
+
+    val directory_entries
+      :  directory_entries
+      -> recursive:bool
+      -> Path.Local.t
+      -> File.Set.t
+  end
 end
 
 (** Resolve the revision in the given remote. The [revision] can be any

test/expect-tests/dune_pkg/rev_store_tests.ml

@@ -17,6 +17,36 @@ let () =
   Dune_tests_common.init ()
 ;;
 
+let%expect_test "immediate and recursive directory entries" =
+  let entries =
+    [ "root"; "a/one"; "a/two"; "a/b/three"; "a/b/c/four"; "other/five" ]
+    |> List.map ~f:Path.Local.of_string
+    |> Rev_store.At_rev.For_tests.make_directory_entries
+  in
+  let print ~recursive path =
+    Rev_store.At_rev.For_tests.directory_entries
+      entries
+      ~recursive
+      (Path.Local.of_string path)
+    |> Rev_store.File.Set.iter ~f:(fun file ->
+      Rev_store.File.path file |> Path.Local.to_string |> print_endline)
+  in
+  print ~recursive:false ".";
+  print ~recursive:false "a";
+  print ~recursive:true "a";
+  print ~recursive:false "missing";
+  [%expect
+    {|
+    root
+    a/one
+    a/two
+    a/b/c/four
+    a/b/three
+    a/one
+    a/two
+    |}]
+;;
+
 let%expect_test "fetching non-existent object twice returns consistent results" =
   (* This test checks that fetching a non-existent object twice returns
      consistent errors. *)