Storage: optimize fold

src/lib_storage/raw_store.ml

@@ -245,6 +245,21 @@ let cursor_next_lwt cursor acc f =
   | Error err -> lwt_fail_error err
   | Ok () -> Lwt.bind acc f
 
+let cursor_at_lwt cursor k acc f =
+  match Lmdb.cursor_at cursor (concat k) with
+  | Error KeyNotFound -> acc
+  | Error err -> lwt_fail_error err
+  | Ok () -> Lwt.bind acc f
+
+(* assumption: store path segments have only characters different than
+   the separator '/', which immediately precedes '0' *)
+let zero_char_str = String.make 1 (Char.chr (Char.code '/' + 1))
+let next_key_after_subdirs = function
+  | [] -> [ zero_char_str ]
+  | (_ :: _) as path ->
+      List.sub path (List.length path - 1) @
+      [List.last_exn path ^ zero_char_str]
+
 let fold t k ~init ~f =
   let base_len = List.length k in
   let rec inner ht cursor acc =
@@ -261,35 +276,33 @@ let fold t k ~init ~f =
         else begin
           let dir = list_sub kk_split 0 (succ base_len) in
           if Hashtbl.mem ht dir then
-            cursor_next_lwt cursor (Lwt.return acc) (inner ht cursor)
+            cursor_at_lwt cursor (next_key_after_subdirs dir)
+              (Lwt.return acc) (inner ht cursor)
           else begin
             Hashtbl.add ht dir () ;
             cursor_next_lwt cursor (f (`Dir dir) acc) (inner ht cursor)
           end
         end in
   with_rw_cursor_lwt t ~f:begin fun cursor ->
-    match Lmdb.cursor_at cursor (concat k) with
+    cursor_at_lwt cursor k
-    | Error KeyNotFound -> Lwt.return init
+      (Lwt.return init)
-    | Error err -> lwt_fail_error err
+      (fun acc ->
-    | Ok () ->
+         let ht = Hashtbl.create 31 in
-        let ht = Hashtbl.create 31 in
+         inner ht cursor acc)
-        inner ht cursor init
   end
 
 let fold_keys t k ~init ~f =
   with_rw_cursor_lwt t ~f:begin fun cursor ->
-    match Lmdb.cursor_at cursor (concat k) with
+    cursor_at_lwt cursor k
-    | Error KeyNotFound -> Lwt.return init
+      (Lwt.return init)
-    | Error err -> lwt_fail_error err
+      (let rec inner acc =
-    | Ok () ->
+         Lmdb.cursor_get cursor >>=? fun (kk, _v) ->
-        let rec inner acc =
+         let kk = MBytes.to_string kk in
-          Lmdb.cursor_get cursor >>=? fun (kk, _v) ->
+         let kk_split = split kk in
-          let kk = MBytes.to_string kk in
+         match is_child ~child:kk_split ~parent:k with
-          let kk_split = split kk in
+         | false -> Lwt.return acc
-          match is_child ~child:kk_split ~parent:k with
+         | true -> cursor_next_lwt cursor (f kk_split acc) inner
-          | false -> Lwt.return acc
+       in inner)
-          | true -> cursor_next_lwt cursor (f kk_split acc) inner
-        in inner init
   end
 
 let keys t =

src/lib_storage/test/assert.ml

@@ -33,6 +33,9 @@ let default_printer _ = ""
 let equal ?(eq=(=)) ?(prn=default_printer) ?(msg="") x y =
   if not (eq x y) then fail (prn x) (prn y) msg
 
+let equal_string ?msg s1 s2 =
+  equal ?msg ~prn:(fun s -> s) s1 s2
+
 let equal_string_option ?msg o1 o2 =
   let prn = function
     | None -> "None"
@@ -58,9 +61,19 @@ let make_equal_list eq prn ?(msg="") x y =
         () in
   iter 0 x y
 
+let equal_string_list ?msg l1 l2 =
+  make_equal_list ?msg (=) (fun x -> x) l1 l2
+
 let equal_string_list_list ?msg l1 l2 =
   let pr_persist l =
     let res =
       String.concat ";" (List.map (fun s -> Printf.sprintf "%S" s) l) in
     Printf.sprintf "[%s]" res in
   make_equal_list ?msg (=) pr_persist l1 l2
+
+let equal_key_dir_list ?msg l1 l2 =
+  make_equal_list ?msg (=)
+    (function
+      | `Key k -> "Key " ^ String.concat "/" k
+      | `Dir k -> "Dir " ^ String.concat "/" k)
+    l1 l2

src/lib_storage/test/test.ml

@@ -26,4 +26,5 @@
 let () =
   Alcotest.run "tezos-storage" [
     "context", Test_context.tests ;
+    "raw_store", Test_raw_store.tests ;
   ]

src/lib_storage/test/test_raw_store.ml

@@ -0,0 +1,86 @@
+(*****************************************************************************)
+(*                                                                           *)
+(* Open Source License                                                       *)
+(* Copyright (c) 2018 Dynamic Ledger Solutions, Inc. <contact@tezos.com>     *)
+(*                                                                           *)
+(* Permission is hereby granted, free of charge, to any person obtaining a   *)
+(* copy of this software and associated documentation files (the "Software"),*)
+(* to deal in the Software without restriction, including without limitation *)
+(* the rights to use, copy, modify, merge, publish, distribute, sublicense,  *)
+(* and/or sell copies of the Software, and to permit persons to whom the     *)
+(* Software is furnished to do so, subject to the following conditions:      *)
+(*                                                                           *)
+(* The above copyright notice and this permission notice shall be included   *)
+(* in all copies or substantial portions of the Software.                    *)
+(*                                                                           *)
+(* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR*)
+(* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,  *)
+(* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL   *)
+(* THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER*)
+(* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING   *)
+(* FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER       *)
+(* DEALINGS IN THE SOFTWARE.                                                 *)
+(*                                                                           *)
+(*****************************************************************************)
+
+open Raw_store
+
+let (>>=) = Lwt.bind
+let (>|=) = Lwt.(>|=)
+let (//) = Filename.concat
+
+let wrap_store_init f _ () =
+  Lwt_utils_unix.with_tempdir "tezos_test_" begin fun base_dir ->
+    let root = base_dir // "store" in
+    init ~mapsize:4_096_000L root >>= function
+    | Error _ -> Assert.fail_msg "wrap_store_init"
+    | Ok store -> f store
+  end
+
+let entries s k = fold s k ~init:[] ~f:(fun e acc -> Lwt.return (e :: acc)) >|= List.rev
+
+let test_fold st =
+  store st ["a"; "b"] (MBytes.of_string "Novembre") >>= fun _ ->
+  store st ["a"; "c"] (MBytes.of_string "Juin") >>= fun _ ->
+  store st ["a"; "d"; "e"] (MBytes.of_string "Septembre") >>= fun _ ->
+  store st ["f";] (MBytes.of_string "Avril") >>= fun _ ->
+  (* The code of '.' is just below the one of '/' ! *)
+  store st ["g";".12";"a"] (MBytes.of_string "Mai") >>= fun _ ->
+  store st ["g";".12";"b"] (MBytes.of_string "Février") >>= fun _ ->
+  store st ["g";"123";"456"] (MBytes.of_string "Mars") >>= fun _ ->
+  store st ["g";"1230"] (MBytes.of_string "Janvier") >>= fun _ ->
+
+  entries st [] >>= fun l ->
+  Assert.equal_key_dir_list ~msg:__LOC__ [`Dir ["a"]; `Key ["f"]; `Dir ["g"]] l ;
+
+  entries st ["0"] >>= fun l ->
+  Assert.equal_key_dir_list ~msg:__LOC__ [] l ;
+
+  entries st ["0"; "1"] >>= fun l ->
+  Assert.equal_key_dir_list ~msg:__LOC__ [] l ;
+
+  entries st ["a"] >>= fun l ->
+  Assert.equal_key_dir_list ~msg:__LOC__ [`Key ["a"; "b"]; `Key ["a"; "c"]; `Dir ["a"; "d"]] l ;
+
+  entries st ["a"; "d"] >>= fun l ->
+  Assert.equal_key_dir_list ~msg:__LOC__ [`Key ["a"; "d"; "e"]] l ;
+
+  entries st ["f"] >>= fun l ->
+  Assert.equal_key_dir_list ~msg:__LOC__ [] l ;
+
+  entries st ["f"; "z"] >>= fun l ->
+  Assert.equal_key_dir_list ~msg:__LOC__ [] l ;
+
+  entries st ["g"] >>= fun l ->
+  Assert.equal_key_dir_list ~msg:__LOC__ [`Dir ["g";".12"]; `Dir ["g";"123"]; `Key ["g";"1230"]] l ;
+
+  entries st ["g";"123"] >>= fun l ->
+  Assert.equal_key_dir_list ~msg:__LOC__ [`Key ["g";"123";"456"]] l ;
+
+  entries st ["z"] >>= fun l ->
+  Assert.equal_key_dir_list ~msg:__LOC__ [] l ;
+
+  Lwt.return_unit
+
+let tests =
+  [Alcotest_lwt.test_case "fold" `Quick (wrap_store_init test_fold)]