Optimize Tree.length for irmin 3.7

src/irmin/tree.ml

@@ -419,6 +419,7 @@ module Make (P : Backend.S) = struct
       mutable map : map option;
       mutable ptr : ptr_option;
       mutable findv_cache : map option;
+      mutable length : int Lazy.t option;
       env : Env.t;
     }
 
@@ -481,7 +482,7 @@ module Make (P : Backend.S) = struct
              Portable_dirty (v, m))
       |> sealv
 
-    let of_v ~env v =
+    let of_v ?length ~env v =
       let ptr, map, value =
         match v with
         | Map m -> (Ptr_none, Some m, None)
@@ -490,12 +491,15 @@ module Make (P : Backend.S) = struct
         | Value _ | Portable_dirty _ | Pruned _ -> (Ptr_none, None, None)
       in
       let findv_cache = None in
-      let info = { ptr; map; value; findv_cache; env } in
+      let info = { ptr; map; value; findv_cache; env; length } in
       { v; info }
 
     let of_map m = of_v (Map m)
     let of_key repo k = of_v (Key (repo, k))
-    let of_value ?updates repo v = of_v (Value (repo, v, updates))
+
+    let of_value ?length ?updates repo v =
+      of_v ?length (Value (repo, v, updates))
+
     let of_portable_dirty p updates = of_v (Portable_dirty (p, updates))
     let pruned h = of_v (Pruned h)
 
@@ -1120,7 +1124,7 @@ module Make (P : Backend.S) = struct
     let empty_hash = hash ~cache:false (empty ())
     let singleton k v = of_map (StepMap.singleton k v)
 
-    let length ~cache t =
+    let slow_length ~cache t =
       match
         (Scan.cascade t
            [
@@ -1150,6 +1154,14 @@ module Make (P : Backend.S) = struct
             | Pnode x -> P.Node_portable.length x |> Lwt.return)
       | Pruned h -> pruned_hash_exn "length" h
 
+    let length ~cache t =
+      match t.info.length with
+      | Some (lazy len) -> Lwt.return len
+      | None ->
+          let+ len = slow_length ~cache t in
+          t.info.length <- Some (Lazy.from_val len);
+          len
+
     let is_empty ~cache t =
       match
         (Scan.cascade t
@@ -1470,6 +1482,27 @@ module Make (P : Backend.S) = struct
       in
       aux_uniq ~path acc 0 t Lwt.return
 
+    let incremental_length t step up n updates =
+      match t.info.length with
+      | None -> None
+      | Some len ->
+          Some
+            (lazy
+              (let len = Lazy.force len in
+               let exists =
+                 match StepMap.find_opt step updates with
+                 | Some (Add _) -> true
+                 | Some Remove -> false
+                 | None -> (
+                     match P.Node.Val.find n step with
+                     | None -> false
+                     | Some _ -> true)
+               in
+               match up with
+               | Add _ when not exists -> len + 1
+               | Remove when exists -> len - 1
+               | _ -> len))
+
     let update t step up =
       let env = t.info.env in
       let of_map m =
@@ -1483,7 +1516,9 @@ module Make (P : Backend.S) = struct
       let of_value repo n updates =
         let updates' = StepMap.add step up updates in
         if updates == updates' then t
-        else of_value ~env repo n ~updates:updates'
+        else
+          let length = incremental_length t step up n updates in
+          of_value ?length ~env repo n ~updates:updates'
       in
       let of_portable n updates =
         let updates' = StepMap.add step up updates in
@@ -1623,7 +1658,8 @@ module Make (P : Backend.S) = struct
 
   let of_backend_node repo n =
     let env = Env.empty () in
-    Node.of_value ~env repo n
+    let length = lazy (P.Node.Val.length n) in
+    Node.of_value ~length ~env repo n
 
   let dump ppf = function
     | `Node n -> Fmt.pf ppf "node: %a" Node.dump n

test/irmin-pack/test_tree.ml

@@ -28,6 +28,7 @@ type ('key, 'value) op =
   | Del of 'key
   | Find of 'key
   | Find_tree of 'key
+  | Length of 'key * int
 
 module Make (Conf : Irmin_pack.Conf.S) = struct
   module Store = struct
@@ -85,15 +86,25 @@ module Make (Conf : Irmin_pack.Conf.S) = struct
     | Del k -> Tree.remove tree k
     | Find k -> find tree k
     | Find_tree k -> find_tree tree k
-
-  let run ops tree =
+    | Length (k, len_expected) ->
+        let+ len_tree = Tree.length tree k in
+        Alcotest.(check int)
+          (Fmt.str "expected tree length at %a" Fmt.(Dump.list string) k)
+          len_expected len_tree;
+        tree
+
+  let run_disjoint ops tree =
     let run_one op =
       let* _ = run_one tree op in
       Lwt.return_unit
     in
     let+ () = Lwt_list.iter_s run_one ops in
     (tree, ())
 
+  let run ops tree =
+    let+ tree = Lwt_list.fold_left_s run_one tree ops in
+    (tree, ())
+
   let proof_of_ops repo hash ops : _ Lwt.t =
     let+ t, () = Store.Tree.produce_proof repo hash (run ops) in
     t
@@ -265,6 +276,16 @@ let check_equivalence tree proof op =
         (Fmt.str "same tree at %a" Fmt.(Dump.list string) k)
         v_tree v_proof;
       (tree, proof)
+  | Length (k, len_expected) ->
+      let* len_tree = Tree.length tree k in
+      Alcotest.(check int)
+        (Fmt.str "expected tree length at %a" Fmt.(Dump.list string) k)
+        len_expected len_tree;
+      let+ len_proof = Tree.length proof k in
+      Alcotest.(check int)
+        (Fmt.str "same tree length at %a" Fmt.(Dump.list string) k)
+        len_tree len_proof;
+      (tree, proof)
 
 let test_proofs ctxt ops =
   let tree = ctxt.tree in
@@ -291,17 +312,18 @@ let test_proofs ctxt ops =
     Lwt_list.fold_left_s
       (fun (tree, proof) op -> check_equivalence tree proof op)
       (tree, tree_proof)
-      [
-        Add ([ "00" ], "0");
-        Add ([ "00" ], "1");
-        Del [ "00" ];
-        Find [ "00" ];
-        Add ([ "00" ], "0");
-        Add ([ "00" ], "1");
-        Find [ "00" ];
-        Find_tree [ "01" ];
-        Find_tree [ "z"; "o"; "o" ];
-      ]
+      (ops
+      @ [
+          Add ([ "00" ], "0");
+          Add ([ "00" ], "1");
+          Del [ "00" ];
+          Find [ "00" ];
+          Add ([ "00" ], "0");
+          Add ([ "00" ], "1");
+          Find [ "00" ];
+          Find_tree [ "01" ];
+          Find_tree [ "z"; "o"; "o" ];
+        ])
   in
   Lwt.return_unit
 
@@ -323,6 +345,40 @@ let test_small_inode () =
   let ops = [ Add ([ "00" ], ""); Del [ "01" ] ] in
   test_proofs ctxt ops
 
+let test_length_proof () =
+  let bindings = bindings fewer_steps in
+  let size = List.length fewer_steps in
+  let* ctxt = init_tree bindings in
+  let ops =
+    [
+      Length ([], size) (* initial size *);
+      Add ([ "01" ], "0");
+      Length ([], size) (* "01" was already accounted for *);
+      Add ([ "01" ], "1");
+      Length ([], size) (* adding it again doesn't change the length *);
+      Add ([ "new" ], "0");
+      Length ([], size + 1) (* "new" is a new file, so the length changes *);
+      Add ([ "new" ], "1");
+      Length ([], size + 1) (* adding it again doesn't change the length *);
+      Del [ "inexistant" ];
+      Length ([], size + 1)
+      (* removing an inexistant object doesn't change the length *);
+      Del [ "00" ];
+      Length ([], size) (* but removing the existing "00" does *);
+      Del [ "00" ];
+      Length ([], size) (* removing "00" twice doesn't change the length *);
+      Del [ "new" ];
+      Length ([], size - 1) (* removing the fresh "new" does *);
+      Del [ "new" ];
+      Length ([], size - 1) (* but only once *);
+      Add ([ "new" ], "2");
+      Length ([], size) (* adding "new" again does *);
+      Add ([ "00" ], "2");
+      Length ([], size + 1) (* adding "00" again does too *);
+    ]
+  in
+  test_proofs ctxt ops
+
 let test_deeper_proof () =
   let* ctxt =
     let tree = Tree.empty () in
@@ -704,6 +760,8 @@ let tests =
       (fun _switch -> test_large_inode);
     Alcotest_lwt.test_case "test Merkle proof for small inodes" `Quick
       (fun _switch -> test_small_inode);
+    Alcotest_lwt.test_case "test Merkle proof for Tree.length" `Quick
+      (fun _switch -> test_length_proof);
     Alcotest_lwt.test_case "test deeper Merkle proof" `Quick (fun _switch ->
         test_deeper_proof);
     Alcotest_lwt.test_case "test large Merkle proof" `Slow (fun _switch ->