vhd-tool: Add read_headers_interval command

This command returns a more efficient representation of allocated clusters
(when compared to read_headers), utilizing a sparse interval format instead of
returning every single allocated cluster.

This is the more efficient option, decreasing the filesize and memory usage in
vhd-tool, but it's currently under a feature flag, so it's added as a new
command instead of replacing read_headers immediately.

Cram test for read_headers is still passing, so this refactoring has preserved
the legacy format.

Signed-off-by: Andrii Sultanov <andriy.sultanov@vates.tech>

Author: last-genius

ocaml/libs/vhd/vhd_format/f.ml

@@ -2903,24 +2903,10 @@ functor
 
       let raw ?from (vhd : fd Vhd.t) = raw_common ?from vhd
 
-      let vhd_blocks_to_json (t : fd Vhd.t) =
+      let vhd_blocks_to_json_aux (t : fd Vhd.t) blocks =
         let block_size_sectors_shift =
           t.Vhd.header.Header.block_size_sectors_shift
         in
-        let max_table_entries = Vhd.used_max_table_entries t in
-
-        let include_block = include_block None t in
-
-        let blocks =
-          Seq.init max_table_entries Fun.id
-          |> Seq.filter_map (fun i ->
-              if include_block i then
-                Some (`Int i)
-              else
-                None
-          )
-          |> List.of_seq
-        in
         let json =
           `Assoc
             [
@@ -2934,6 +2920,52 @@ functor
         let json_string = Yojson.to_string json in
         print_string json_string ; return ()
 
+      let vhd_blocks_to_json (t : fd Vhd.t) =
+        let max_table_entries = Vhd.used_max_table_entries t in
+        let blocks =
+          Seq.init max_table_entries Fun.id
+          |> Seq.filter_map (fun i ->
+              if include_block None t i then
+                Some (`Int i)
+              else
+                None
+          )
+          |> List.of_seq
+        in
+        vhd_blocks_to_json_aux t blocks
+
+      let vhd_blocks_to_json_interval (t : fd Vhd.t) =
+        let max_table_entries = Vhd.used_max_table_entries t in
+        let blocks, last_block =
+          Seq.init max_table_entries Fun.id
+          |> Seq.fold_left
+               (fun (acc, left_block) i ->
+                 if include_block None t i then
+                   match left_block with
+                   | Some _ ->
+                       (acc, left_block)
+                   | None ->
+                       (acc, Some i)
+                 else
+                   match left_block with
+                   | Some x ->
+                       (`List [`Int x; `Int (i - 1)] :: acc, None)
+                   | None ->
+                       (acc, None)
+               )
+               ([], None)
+        in
+        (* Close off the interval we were tracking we ran off the end of the seq *)
+        let blocks =
+          match last_block with
+          | Some x ->
+              `List [`Int x; `Int (max_table_entries - 1)] :: blocks
+          | None ->
+              blocks
+        in
+        let blocks = List.rev blocks in
+        vhd_blocks_to_json_aux t blocks
+
       let vhd_common ?from ?raw ?(emit_batmap = false) (t : fd Vhd.t) =
         let block_size_sectors_shift =
           t.Vhd.header.Header.block_size_sectors_shift
@@ -3173,6 +3205,8 @@ functor
         Vhd_input.vhd_common ?from ~raw vhd
 
       let blocks_json = Vhd_input.vhd_blocks_to_json
+
+      let blocks_json_interval = Vhd_input.vhd_blocks_to_json_interval
     end
 
     (* Create a VHD stream from data on t, using `include_block` guide us which blocks have data *)

ocaml/libs/vhd/vhd_format/f.mli

@@ -470,6 +470,8 @@ module From_file : functor (F : S.FILE) -> sig
         [from] into [t] *)
 
     val blocks_json : fd Vhd.t -> unit t
+
+    val blocks_json_interval : fd Vhd.t -> unit t
   end
 
   module Raw_input : sig

ocaml/vhd-tool/cli/main.ml

@@ -385,19 +385,34 @@ let stream_cmd =
   , Cmd.info "stream" ~sdocs:_common_options ~doc ~man
   )
 
+let vhd_source =
+  let doc = Printf.sprintf "Path to the VHD file" in
+  Arg.(required & pos 0 (some file) None & info [] ~doc)
+
 let read_headers_cmd =
   let doc =
     {|Parse VHD headers and output allocated blocks information in JSON format \
      like: {"virtual_size": X, "cluster_bits": X, "data_clusters": [1,2,3]}|}
   in
-  let source =
-    let doc = Printf.sprintf "Path to the VHD file" in
-    Arg.(required & pos 0 (some file) None & info [] ~doc)
-  in
-  ( Term.(ret (const Impl.read_headers $ common_options_t $ source))
+  ( Term.(
+      ret
+        (const (Impl.read_headers ~legacy:true) $ common_options_t $ vhd_source)
+    )
   , Cmd.info "read_headers" ~sdocs:_common_options ~doc
   )
 
+let read_headers_interval_cmd =
+  let doc =
+    {|Parse VHD headers and output allocated blocks intervals information in JSON format \
+     like: {"virtual_size": X, "cluster_bits": X, "data_clusters": [[1,13],[17,17],[19,272]]|}
+  in
+  ( Term.(
+      ret
+        (const (Impl.read_headers ~legacy:false) $ common_options_t $ vhd_source)
+    )
+  , Cmd.info "read_headers_interval" ~sdocs:_common_options ~doc
+  )
+
 let cmds =
   [
     info_cmd
@@ -408,6 +423,7 @@ let cmds =
   ; serve_cmd
   ; stream_cmd
   ; read_headers_cmd
+  ; read_headers_interval_cmd
   ]
   |> List.map (fun (t, i) -> Cmd.v i t)
 

ocaml/vhd-tool/src/impl.ml

@@ -1168,11 +1168,15 @@ let stream_t common args ?(progress = no_progress_bar) () =
     args.StreamCommon.tar_filename_prefix args.StreamCommon.good_ciphersuites
     args.StreamCommon.verify_cert
 
-let read_headers common source =
+let read_headers common source ~legacy =
   let path = [Filename.dirname source] in
   let thread =
     retry common 3 (fun () -> Vhd_IO.openchain ~path source false) >>= fun t ->
-    Vhd_IO.close t >>= fun () -> Hybrid_input.blocks_json t
+    Vhd_IO.close t >>= fun () ->
+    if legacy then
+      Hybrid_input.blocks_json t
+    else
+      Hybrid_input.blocks_json_interval t
   in
   Lwt_main.run thread ; `Ok ()
 

ocaml/vhd-tool/src/impl.mli

@@ -36,7 +36,7 @@ val stream :
   Common.t -> StreamCommon.t -> [> `Error of bool * string | `Ok of unit]
 
 val read_headers :
-  Common.t -> string -> [> `Error of bool * string | `Ok of unit]
+  Common.t -> string -> legacy:bool -> [> `Error of bool * string | `Ok of unit]
 
 val serve :
      Common.t