Convert directly to array

liam923 · liam923 · commit fc1857f3247f · 2024-08-05T12:35:18.000-04:00
diff --git a/ocaml/file_formats/cms_format.ml b/ocaml/file_formats/cms_format.ml
@@ -120,7 +120,7 @@ let save_cms filename modname binary_annots sourcefile initial_env shape =
             cms_source_digest = source_digest;
             cms_initial_env = if Cmt_format.need_to_clear_env
               then Env.keep_only_summary initial_env else initial_env;
-            cms_uid_to_loc = cms_uid_to_loc |> Shape.Uid.Tbl.to_list |> Array.of_list;
+            cms_uid_to_loc = cms_uid_to_loc |> Shape.Uid.Tbl.to_array;
             cms_uid_to_attributes;
             cms_impl_shape = shape;
             cms_ident_occurrences
diff --git a/ocaml/utils/identifiable.ml b/ocaml/utils/identifiable.ml
@@ -82,6 +82,7 @@ module type Tbl = sig
   val of_map : 'a Map.Make(T).t -> 'a t
   val memoize : 'a t -> (key -> 'a) -> key -> 'a
   val map : 'a t -> ('a -> 'b) -> 'b t
+  val to_array : 'a t -> (T.t * 'a) array
 end
 
 module Pair (A : Thing) (B : Thing) : Thing with type t = A.t * B.t = struct
@@ -226,6 +227,28 @@ module Make_tbl (T : Thing) = struct
 
   let map t f =
     of_map (T_map.map f (to_map t))
+
+  let to_array t =
+    let build_array key value acc =
+      match acc with
+      | None ->
+        (* acc is None if we're on the first element. If this is the case, allocate the
+          array, duplicating the first value for each entry *)
+        let arr = Array.make (length t) (key, value) in
+        Some (arr, 1)
+      | Some (arr, i) ->
+        (* If acc is Some, then the array has been allocated and we want to write the
+           i-th value *)
+        arr.(i) <- (key, value);
+        Some (arr, i + 1)
+    in
+    let result = fold build_array t None in
+    match result with
+    | Some (arr, _) -> arr
+    | None ->
+      (* The result is None iff the map is empty, in which case we return an empty
+         array *)
+      [||]
 end
 
 module type S = sig
diff --git a/ocaml/utils/identifiable.mli b/ocaml/utils/identifiable.mli
@@ -97,6 +97,7 @@ module type Tbl = sig
   val of_map : 'a Map.Make(T).t -> 'a t
   val memoize : 'a t -> (key -> 'a) -> key -> 'a
   val map : 'a t -> ('a -> 'b) -> 'b t
+  val to_array : 'a t -> (T.t * 'a) array
 end
 
 module type S = sig
