Tidy up 2024 day 11 code.

bin/day2411.ml

@@ -1,13 +1,15 @@
 module IntMap = Map.Make (Int)
 
-let load_file fname =
-  match In_channel.with_open_text fname In_channel.input_line with
-  | Some x -> x
-  | None -> failwith "load_file"
-
+(** [apply_n n fn arg] is equivalent to [(fn (fn ... (fn (fn arg))))] where [fn]
+    is called [n] times.*)
 let rec apply_n n fn arg = if n <= 0 then arg else apply_n (n - 1) fn (fn arg)
+
+(** [update_count n o] returns [Some n] if [o] is [None], or [Some (n + x)] if
+    [o] is [Some x]. *)
 let update_count n = function None -> Some n | Some x -> Some (x + n)
 
+(** [map_of_ints lst] returns an [IntMap] with key-value pairs counting how many
+    times each integer appears in [lst]. *)
 let map_of_ints =
   let rec impl acc = function
     | [] -> acc
@@ -15,24 +17,21 @@ let map_of_ints =
   in
   impl IntMap.empty
 
-(** [calc_blink_rec acc lst] returns an updated map based off [acc] with the
-    result of apply a blink step to the stones in [lst]. Entries in [lst] are
-    pairs of [(stone id, count)]. [acc] and the resulting map have keys which
-    are stone id, and values which are count. *)
-let rec calc_blink_rec acc = function
-  | [] -> acc
-  | (0, n) :: t -> calc_blink_rec (IntMap.update 1 (update_count n) acc) t
-  | (x, n) :: t when Aoc.digits10 x mod 2 = 0 ->
+(** [map_stone id n map] calculates how a collection of stones changes in a
+    blink, updating [map] with the result. [id] is the ID of the stone to
+    update, [n] is how many stones there are with that ID. *)
+let map_stone id n acc =
+  match id with
+  | 0 -> IntMap.update 1 (update_count n) acc
+  | x when Aoc.digits10 x mod 2 = 0 ->
       let pow = Aoc.pow10 (Aoc.digits10 x / 2) in
       let acc = IntMap.update (x / pow) (update_count n) acc in
       let acc = IntMap.update (x mod pow) (update_count n) acc in
-      calc_blink_rec acc t
-  | (x, n) :: t ->
-      calc_blink_rec (IntMap.update (x * 2024) (update_count n) acc) t
+      acc
+  | x -> IntMap.update (x * 2024) (update_count n) acc
 
-(** [calc_blink map] calculates how a collection of stones changes in a blink.
-    [map] is a map with key of stone ID, and value number of times that stone
-    appears. The result is a map with similar key, value pairs.
+(** [calc_blink map] calculates how a collection of stones changes in a blink,
+    returning the result.
 
     This improves performance because we find that the transformation stones go
     through ends up producing repeated numbers. e.g.: 0 -> 1 -> 2024 -> 20 24 ->
@@ -41,13 +40,17 @@ let rec calc_blink_rec acc = function
     We also note that despite the problem description saying stones stay in
     order, the result we are asked for (number of stones) does not require them
     to be in order. *)
-let calc_blink map = IntMap.to_list map |> calc_blink_rec IntMap.empty
+let calc_blink map = IntMap.fold map_stone map IntMap.empty
 
-(** [part n str] returns the number of stones after [n] blinks, given an initial
-    string [str] of space seperated stone IDs. *)
-let part n str =
-  let map = Aoc.ints_of_string str |> map_of_ints |> apply_n n calc_blink in
+(** [part n lst] returns the number of stones after [n] blinks, given an initial
+    list, [lst], of stone IDs. *)
+let part n lst =
+  let map = map_of_ints lst |> apply_n n calc_blink in
   IntMap.fold (fun _ v acc -> v + acc) map 0
 
+(** [ints_of_file fname] returns the integers listed on the first line of
+    [fname]. *)
+let ints_of_file fname = Aoc.string_of_file fname |> Aoc.ints_of_string
+
 let _ =
-  Aoc.main load_file [ (string_of_int, part 25); (string_of_int, part 75) ]
+  Aoc.main ints_of_file [ (string_of_int, part 25); (string_of_int, part 75) ]

lib/aoc.ml

@@ -6,6 +6,11 @@ let distance1 a b = abs (a - b)
 let strings_of_file fname =
   In_channel.with_open_text fname In_channel.input_lines
 
+let string_of_file fname =
+  match In_channel.with_open_text fname In_channel.input_line with
+  | Some x -> x
+  | None -> failwith "Aoc.string_of_file"
+
 let main prep parts =
   try
     match Sys.argv with

lib/aoc.mli

@@ -9,6 +9,9 @@ val strings_of_file : string -> string list
 (** [strings_from_file fname] returns a list of strings from the file [fname].
     Each string represents a line from the file. *)
 
+val string_of_file : string -> string
+(** [string_of_file fname] returns the first line in [fname]. *)
+
 val log10i : int -> int
 (** [log10i n] returns the floor of [(log10 (float_of_int n))]. [n] must be
     positive. *)