Tidy up 2024 day 14 code.

bin/day2414.ml

@@ -1,5 +1,7 @@
 module IntMap = Map.Make (Int)
 
+(** [parse_robot s] returns a robot description [((x, y), (dx, dy))] parsed from
+    the string [s]. *)
 let parse_robot s =
   let re =
     Str.regexp {|p=\(-?[0-9]+\),\(-?[0-9]+\) v=\(-?[0-9]+\),\(-?[0-9]+\)|}
@@ -10,23 +12,41 @@ let parse_robot s =
     ( int_of_string (Str.matched_group 3 s),
       int_of_string (Str.matched_group 4 s) ) )
 
-let parse_robots = List.map parse_robot
-let robots_of_file fname = Aoc.strings_of_file fname |> parse_robots
-let width = 101
-let height = 103
-let secs = 100
+(** [robots_of_file fname] returns a list of robots parsed from the file
+    [fname]. *)
+let robots_of_file fname = Aoc.strings_of_file fname |> List.map parse_robot
 
+(** Grid width *)
+let width = 101
+
+(** Grid height *)
+let height = 103
+
+(** Number of seconds to run part 1 for*)
+let secs1 = 100
+
+(** Maximum number of seconds to run part 2 for *)
+let secs2 = 1000000
+
+(** [normalize_velocity robot] returns a robot where the velocity has been
+    normalized to be non-negative in both directions. *)
 let normalize_velocity (p, (dx, dy)) =
   (p, ((dx + width) mod width, (dy + height) mod height))
 
+(** [calc_pos_after secs r] returns the [(x, y)] position of a robot after
+    [secs] seconds. *)
 let calc_pos_after secs ((x, y), (dx, dy)) =
   let x' = (x + (secs * dx)) mod width in
   let y' = (y + (secs * dy)) mod height in
   (x', y')
 
+(** [in_a_quadrant pos] returns true if [pos] is in a quadrant. *)
 let in_a_quadrant (x, y) = x <> width / 2 && y <> height / 2
+
+(** [update_count n] Is used by [IntMap.update] to increment a count. *)
 let update_count = function None -> Some 1 | Some x -> Some (x + 1)
 
+(** [get_quadrant p] returns the quadrant ID that the position [p] is in. *)
 let get_quadrant (x, y) =
   if x < width / 2 && y < height / 2 then 1
   else if x > width / 2 && y < height / 2 then 2
@@ -34,10 +54,14 @@ let get_quadrant (x, y) =
   else if x > width / 2 && y > height / 2 then 3
   else failwith "get_quadrant"
 
-let loc_counts map p =
+(** [quadrant_counts map p] updates the quadrant count map [map] with [p]. Keys
+    to [map] are quadrant IDs, and the values are the number of robots in that
+    quadrant. *)
+let quadrant_counts map p =
   let idx = get_quadrant p in
   IntMap.update idx update_count map
 
+(** [print_locs lst] prints the grid layout of the robots given in list. *)
 let print_locs lst =
   let a = Array.make_matrix height width '.' in
   let rec impl = function
@@ -54,22 +78,30 @@ let print_locs lst =
       print_newline ())
     a
 
+(** [part1 robots] solves part1 for the list [robots]. *)
 let part1 robots =
   let counts =
     robots
     |> List.map normalize_velocity
-    |> List.map (calc_pos_after secs)
+    |> List.map (calc_pos_after secs1)
     |> List.filter in_a_quadrant
-    |> List.fold_left loc_counts IntMap.empty
+    |> List.fold_left quadrant_counts IntMap.empty
   in
   IntMap.fold (fun _ v acc -> acc * v) counts 1
 
+(** [find_tree max_n lst] tries to find the Christmas tree picture by iterating
+    through the various steps robots move in. Returns the number of iterations
+    before finding the tree. *)
 let find_tree max_n lst =
+  (* We assume that the picture will occur when every robot is in a unique
+     location. *)
   let num_robots = List.length lst in
   let rec impl n =
     if n > max_n then failwith "None found"
     else
       let poses = List.map (calc_pos_after n) lst in
+      (* If every tree is in a unique location then sort_uniq will not remove
+         any elements from the list. *)
       if List.length (List.sort_uniq Aoc.IntPair.compare poses) = num_robots
       then (
         print_locs poses;
@@ -78,9 +110,10 @@ let find_tree max_n lst =
   in
   impl 0
 
+(** [part2 robots] solves part 2 for the list of robots. *)
 let part2 robots =
   let robots = List.map normalize_velocity robots in
-  find_tree 100000 robots
+  find_tree secs2 robots
 
 let _ =
   Aoc.main robots_of_file [ (string_of_int, part1); (string_of_int, part2) ]