Optimize in-memory repo parsing

Author: futurGH

mist/lib/mst.ml

@@ -118,6 +118,86 @@ let ( >>? ) lazy_opt_lwt f =
   let%lwt result = Lazy.force lazy_opt_lwt in
   f result
 
+(* extracts leaves from a block map *)
+let leaves_from_blocks (blocks : Block_map.t) (root : Cid.t) :
+    (string * Cid.t) list =
+  let leaves = ref [] in
+  let stack = Stack.create () in
+  Stack.push (root, "") stack ;
+  while not (Stack.is_empty stack) do
+    let cid, prefix = Stack.pop stack in
+    match Block_map.get cid blocks with
+    | None ->
+        () (* missing block probably a record *)
+    | Some bytes -> (
+      try
+        let node = decode_block_raw bytes in
+        (* proess left subtree *)
+        ( match node.l with
+        | Some left_cid ->
+            Stack.push (left_cid, prefix) stack
+        | None ->
+            () ) ;
+        (* process entries in reverse order so they come out in correct order *)
+        let last_key = ref prefix in
+        List.iter
+          (fun (entry : entry_raw) ->
+            let key_prefix =
+              if entry.p = 0 then ""
+              else if entry.p <= String.length !last_key then
+                String.sub !last_key 0 entry.p
+              else !last_key
+            in
+            let full_key = key_prefix ^ Bytes.to_string entry.k in
+            last_key := full_key ;
+            leaves := (full_key, entry.v) :: !leaves ;
+            (* push right subtree to stack *)
+            match entry.t with
+            | Some right_cid ->
+                Stack.push (right_cid, full_key) stack
+            | None ->
+                () )
+          node.e
+      with Invalid_argument _ -> () )
+  done ;
+  List.sort (fun (k1, _) (k2, _) -> String.compare k1 k2) !leaves
+
+(* extracts just mst node cids (non-leaf blocks) from a block map *)
+let mst_node_cids_from_blocks (blocks : Block_map.t) (root : Cid.t) : Cid.t list
+    =
+  let nodes = ref [] in
+  let visited = ref Cid.Set.empty in
+  let stack = Stack.create () in
+  Stack.push root stack ;
+  while not (Stack.is_empty stack) do
+    let cid = Stack.pop stack in
+    if not (Cid.Set.mem cid !visited) then (
+      visited := Cid.Set.add cid !visited ;
+      match Block_map.get cid blocks with
+      | None ->
+          ()
+      | Some bytes -> (
+        try
+          let node = decode_block_raw bytes in
+          nodes := cid :: !nodes ;
+          (* add all children to stack *)
+          ( match node.l with
+          | Some left_cid ->
+              Stack.push left_cid stack
+          | None ->
+              () ) ;
+          List.iter
+            (fun (entry : entry_raw) ->
+              match entry.t with
+              | Some right_cid ->
+                  Stack.push right_cid stack
+              | None ->
+                  () )
+            node.e
+        with Invalid_argument _ -> () ) )
+  done ;
+  !nodes
+
 module type Intf = sig
   module Store : Writable_blockstore
 

pegasus/lib/repository.ml

@@ -592,20 +592,9 @@ let import_car t (stream : Car.stream) : (t, exn) Lwt_result.t =
           failwith ("invalid commit: " ^ e)
     in
     if commit.did <> t.did then failwith "did does not match commit did" ;
-    (* create in-memory mst to walk *)
-    let mem_bs = Mist.Storage.Memory_blockstore.create ~blocks:all_blocks () in
-    let mem_mst : Mem_mst.t = {blockstore= mem_bs; root= commit.data} in
-    let%lwt leaves = Mem_mst.leaves_of_root mem_mst in
-    let leaf_cids =
-      List.fold_left
-        (fun acc (_, cid) -> Cid.Set.add cid acc)
-        Cid.Set.empty leaves
-    in
-    (* get mst nodes by filtering out leaves and commit from all blocks *)
+    let leaves = Mist.Mst.leaves_from_blocks all_blocks commit.data in
     let mst_node_cids =
-      Block_map.keys all_blocks
-      |> List.filter (fun cid ->
-          (not (Cid.equal cid root)) && not (Cid.Set.mem cid leaf_cids) )
+      Mist.Mst.mst_node_cids_from_blocks all_blocks commit.data
     in
     (* collect mst node blocks for insert *)
     let mst_blocks =
@@ -620,23 +609,23 @@ let import_car t (stream : Car.stream) : (t, exn) Lwt_result.t =
     in
     (* collect record data for insert *)
     let since = Tid.now () in
-    let blob_refs : (string * Cid.t) list ref = ref [] in
-    let record_data =
-      List.map
-        (fun (path, cid) ->
+    let record_data, blob_refs =
+      List.fold_left
+        (fun (acc_data, acc_refs) (path, cid) ->
           match Block_map.get cid all_blocks with
           | Some data ->
               let record = Lex.of_cbor data in
               let record_refs =
                 Util.find_blob_refs record
                 |> List.map (fun (br : Mist.Blob_ref.t) -> (path, br.ref))
               in
-              blob_refs := record_refs @ !blob_refs ;
-              (path, cid, data, since)
+              ( (path, cid, data, since) :: acc_data
+              , List.rev_append record_refs acc_refs )
           | None ->
               failwith ("missing record block: " ^ Cid.to_string cid) )
-        leaves
+        ([], []) leaves
     in
+    let record_data = List.rev record_data in
     let%lwt _ =
       Util.use_pool t.db.db (fun conn ->
           Util.transact conn (fun () ->
@@ -645,7 +634,7 @@ let import_car t (stream : Car.stream) : (t, exn) Lwt_result.t =
               let$! () = User_store.Bulk.put_blocks mst_blocks conn in
               let$! () =
                 [%rapper execute {sql| DELETE FROM records |sql}] () conn
-                in
+              in
               let$! () = User_store.Bulk.put_records record_data conn in
               let$! () = User_store.Bulk.put_blob_refs !blob_refs conn in
               Lwt.return_ok () ) )