Tidy up code for 2024 day 10.

bin/day2410.ml

@@ -1,25 +1,48 @@
 type grid = { grid : string; width : int; height : int }
+(** Type to hold the grid in... We should be moving this to the library soon *)
 
+(** [grid_of_file fname] loads a grid from the given file *)
 let grid_of_file fname =
   let strs = Aoc.strings_of_file fname in
   let grid = List.fold_left String.cat "" strs in
   let width = String.length (List.hd strs) in
   let height = String.length grid / width in
+  assert (width * height = String.length grid);
   { grid; width; height }
 
+(** [grid_length grid] gives the length of the grid ([width * height]). *)
 let grid_length grid = String.length grid.grid
+
+(** [grid_get_by_idx grid idx] returns the element of [grid] at [idx]. *)
 let grid_get_by_idx grid idx = grid.grid.[idx]
+
+(** [grid_get_by_idx grid pos] returns the element of [grid] at [pos]. *)
 let grid_get_by_pos grid (x, y) = grid.grid.[x + (y * grid.width)]
+
+(** [grid_pos_if_idx grid idx] returns the [(x, y)] position of [idx] in [grid].
+*)
 let grid_pos_of_idx grid idx = (idx mod grid.width, idx / grid.width)
 
+(** [grid_pos_is_valid grid pos] returns [true] if and only if [pos] is a valid
+    position in [grid]. *)
 let grid_pos_is_valid grid (x, y) =
   x >= 0 && y >= 0 && x < grid.width && y < grid.height
 
+(** [next_char ch] returns the next character after [ch]. *)
 let next_char ch = Char.chr (1 + Char.code ch)
+
+(** [ten] really isn't the digit 10, but is the character after '9'. *)
 let ten = next_char '9'
+
+(** The character '1'. *)
 let one = '1'
 
+(** [find_trail grid pos0] returns a list of all end points of trails in [grid]
+    starting at [pos0]. [pos0] must point to a valid position that contains a
+    '0'. The same endpoint may be returned multiple times if there are multiple
+    routes to it. *)
 let find_trail grid pos0 =
+  assert (grid_get_by_pos grid pos0 = '0');
   let add_pos lst pos digit =
     if grid_pos_is_valid grid pos && grid_get_by_pos grid pos = digit then
       pos :: lst
@@ -41,6 +64,9 @@ let find_trail grid pos0 =
   in
   impl [ pos0 ] one
 
+(** [find_trails grid] returns the list of list of end-points of trails starting
+    at each position in [grid]. The [n]th element of the returned list
+    corresponds to the trails starting at index [n]. *)
 let find_trails grid =
   let rec impl acc idx =
     if idx >= grid_length grid then acc
@@ -50,12 +76,15 @@ let find_trails grid =
   in
   impl [] 0 |> List.rev
 
-let part1 grid =
-  find_trails grid
-  |> List.map (List.sort_uniq compare)
-  |> List.map List.length |> List.fold_left ( + ) 0
+(** [part sort_fn grid] returns a count of all trails in [grid], before counting
+    the trails for each grid index are sorted by [sort_fn]. *)
+let part sort_fn grid =
+  find_trails grid |> List.map sort_fn |> List.map List.length
+  |> List.fold_left ( + ) 0
 
-let part2 grid =
-  find_trails grid |> List.map List.length |> List.fold_left ( + ) 0
-
-let _ = Aoc.main grid_of_file [ (string_of_int, part1); (string_of_int, part2) ]
+let _ =
+  Aoc.main grid_of_file
+    [
+      (string_of_int, part (List.sort_uniq Stdlib.compare));
+      (string_of_int, part Fun.id);
+    ]