Tidy up code for 2024 day 9

Not happy that mutable state is being used.  But this seems to be the
simplest way, and is efficient.

bin/day2409.ml

@@ -3,26 +3,22 @@ let load_file fname =
   | Some x -> x
   | None -> failwith "load_file"
 
-(*
-  let disk_print disk =
-    let prnt c = if c = -1 then print_char '.' else print_int c in
-    Array.iter prnt disk;
-    print_newline();
-*)
+(** [disk_size disk_str] returns the size of the disk represented by the string
+    [disk_str]. See AoC 2024 day 9 for description of string format. *)
 let disk_size disk_str =
   let rec impl acc disk_str =
-    let len = String.length disk_str in
-    match len with
+    match String.length disk_str with
     | 0 -> acc
     | len ->
         let h = int_of_string (String.sub disk_str 0 1) in
-        flush_all ();
         let t = String.sub disk_str 1 (len - 1) in
-        flush_all ();
         impl (acc + h) t
   in
   impl 0 disk_str
 
+(** [disk_init disk_str] returns a disk which represents the description given
+    by [disk_str]. The returned disk is a mutable array with elements being [-1]
+    for free space and [id >= 0] for file with the given ID. *)
 let disk_init disk_str =
   let size = disk_size disk_str in
   let disk = Array.make size (-1) in
@@ -48,6 +44,8 @@ let disk_init disk_str =
   impl 0 0 disk_str;
   disk
 
+(** [disk_defrag disk] defrags [disk] given 2024/09/part 1 rules. Updates disk
+    in place. *)
 let disk_defrag disk =
   let rec impl front back =
     if front >= back then ()
@@ -61,11 +59,91 @@ let disk_defrag disk =
   impl 0 (Array.length disk - 1);
   disk
 
-let calc_checksum disk =
+(** [find_largest_id disk] returns the largest ID on a fragmented disk [disk].
+*)
+let find_largest_id disk =
+  let rec impl pos =
+    if pos = -1 then failwith "find_largest_id.impl"
+    else if disk.(pos) = -1 then impl (pos - 1)
+    else disk.(pos)
+  in
+  impl (Array.length disk - 1)
+
+(** [find_id disk id search_pos] returns a pair [(start, length)] of the file
+    [id] on [disk]. We only look backwards from [search_pos]. *)
+let rec find_id disk id search_pos =
+  let rec count_length len pos =
+    if pos < 0 then (pos, len)
+    else if disk.(pos) <> id then (pos + 1, len)
+    else count_length (len + 1) (pos - 1)
+  in
+  if search_pos < 0 then failwith "find_id"
+  else if disk.(search_pos) <> id then find_id disk id (search_pos - 1)
+  else count_length 0 search_pos
+
+(** [find_space disk len] finds [len] elements of free space on [disk] starting
+    from 0. Returns [None] if no free-space found, or [Some pos] giving the
+    position of the start of the found space. *)
+let find_space disk len =
+  let rec free_length acc pos =
+    if pos >= Array.length disk then acc
+    else if disk.(pos) <> -1 then acc
+    else free_length (acc + 1) (pos + 1)
+  in
+  let rec impl pos =
+    if pos >= Array.length disk then None
+    else if disk.(pos) <> -1 then impl (pos + 1)
+    else if free_length 0 pos < len then impl (pos + 1)
+    else Some pos
+  in
+  impl 0
+
+(** [move_file disk id src dest] moves the file [id] on [disk] from [src] to
+    [dest].*)
+let rec move_file disk id src dest =
+  if src >= Array.length disk then ()
+  else if dest >= Array.length disk then ()
+  else if disk.(src) <> id then ()
+  else if disk.(dest) <> -1 then failwith "move_block"
+  else (
+    disk.(dest) <- disk.(src);
+    disk.(src) <- -1;
+    move_file disk id (src + 1) (dest + 1);
+    ())
+
+(** [file_defrag disk id search_pos] Locates the file [id] on [disk] and defrags
+    it if possible. We start searching down from [search_pos]. *)
+let file_defrag disk id search_pos =
+  let pos, len = find_id disk id search_pos in
+  match find_space disk len with
+  | None -> pos - 1
+  | Some x ->
+      if pos >= x then move_file disk id pos x;
+      pos - 1
+
+(** Defrag a whole disk according to 2024/09/part 2 rules. *)
+let disk_defrag_whole disk =
+  let max_id = find_largest_id disk in
+  let rec impl id pos =
+    if id = 0 then disk else impl (id - 1) (file_defrag disk id pos)
+  in
+  impl max_id (Array.length disk - 1)
+
+(** [disk_checksum disk] Calculates the checksum for [disk]. *)
+let disk_checksum disk =
   let rec impl acc idx =
-    if disk.(idx) = -1 then acc else impl (acc + (disk.(idx) * idx)) (idx + 1)
+    if idx = Array.length disk then acc
+    else if disk.(idx) = -1 then impl acc (idx + 1)
+    else impl (acc + (disk.(idx) * idx)) (idx + 1)
   in
   impl 0 0
 
-let part1 str = disk_init str |> disk_defrag |> calc_checksum
-let _ = Aoc.main load_file [ (string_of_int, part1) ]
+(** [part algo str] defrags the disk represented by [str] using algorithm
+    [algo]. *)
+let part algo str = disk_init str |> algo |> disk_checksum
+
+let _ =
+  Aoc.main load_file
+    [
+      (string_of_int, part disk_defrag); (string_of_int, part disk_defrag_whole);
+    ]