Return a specialized function so we don't break type safety.

[virt-df.git] / lib / diskimage.ml

diff --git a/lib/diskimage.ml b/lib/diskimage.ml
index 71f56b4..697cb76 100644 (file)
--- a/lib/diskimage.ml
+++ b/lib/diskimage.ml
@@ -303,12 +303,15 @@ let print_binary_tree leaf_printer node_printer tree =
   print tree;
   eprintf "}\n";
 
+type owner =
+    [ `Filesystem of filesystem
+    | `Partitions of partitions
+    | `PhysicalVolume of pv ]
+
 (* A segment describes the owner of a range of disk addresses. *)
 type segment =
     int63 * int63 *                    (* disk offset, size of segment *)
-      [ `Filesystem of filesystem
-      | `Partitions of partitions
-      | `PhysicalVolume of pv ] *      (* owner *)
+      owner *                          (* owner *)
       int63                            (* owner offset *)
 
 type interval = int63 * int63          (* start point, end point (bytes) *)
@@ -657,3 +660,23 @@ let create_ownership machine =
 
   (* Return the ownership structure. *)
   ownership
+
+let get_owners_lookup machine ownership disk =
+  (* Get the correct tree. *)
+  let tree = List.assoc disk ownership in
+
+  fun offset ->
+    let rec query = function
+      | Leaf (_, segments) -> segments
+      | Node ((Leaf ((_, leftend), _) | Node (_, ((_, leftend), _), _) as left),
+             (_, segments),
+             right) ->
+         let subsegments =
+           if offset < leftend then query left else query right in
+         segments @ subsegments
+    in
+    let owners = query tree in
+
+    List.map (
+      fun (_, _, owner, owner_offset) -> (owner, offset -^ owner_offset)
+    ) owners