reorganize day 17 code a bit

src/day17.ml

@@ -11,16 +11,6 @@ let dir_of_char = function
 
 let next_stream_idx stream idx = (idx + 1) mod Array.length stream
 
-type shape =
-  | Row
-  | Plus
-  | Angle
-  | Col
-  | Square
-
-let next_shape_idx idx = (idx + 1) mod 5
-let get_shape idx = [| Row; Plus; Angle; Col; Square |].(idx)
-
 (*[{
                        #
          #       #     #
@@ -30,6 +20,15 @@ let get_shape idx = [| Row; Plus; Angle; Col; Square |].(idx)
  row    plus   angle   col   square
 
 }]*)
+type shape =
+  | Row
+  | Plus
+  | Angle
+  | Col
+  | Square
+
+let get_shape idx = [| Row; Plus; Angle; Col; Square |].(idx mod 5)
+
 let shape_points shape pos =
   let points =
     match shape with
@@ -42,35 +41,38 @@ let shape_points shape pos =
   points |> List.map (fun point -> Vec2.(pos + of_tuple point))
 ;;
 
-type cache_key =
-  { stream_idx : int
-  ; rock_idx : int
+type loop_key =
+  { rock : shape
+  ; stream_idx : int
   ; normalized_topo : int array
   }
 
+type loop =
+  { rock_idx : int
+  ; height : int
+  }
+
 type cave =
   { grid : char Grid.t
-  ; loop_cache : (cache_key, int * int) Hashtbl.t
   ; mutable grid_topo : int array
   ; mutable grid_height : int
-  ; mutable height : int
-  ; mutable rock_idx : int
+  ; total_rocks : int
+  ; mutable rocks : int
   ; mutable stream_idx : int
-  ; mutable dropped_rocks : int
-  ; mutable remaining_rocks : int
+  ; loop_cache : (loop_key, loop) Hashtbl.t
+  ; mutable height : int
   }
 
-let init_cave max_height remaining_rocks =
+let init_cave max_height total_rocks =
   let grid = Grid.init ~width:7 ~height:max_height (Fun.const '.') in
   { grid
   ; height = -1
   ; grid_height = -1
-  ; rock_idx = 0
-  ; dropped_rocks = 0
+  ; rocks = 0
   ; stream_idx = 0
   ; grid_topo = Array.init cave_width (Fun.const 0)
   ; loop_cache = Hashtbl.create max_height
-  ; remaining_rocks
+  ; total_rocks
   }
 ;;
 
@@ -99,7 +101,7 @@ let can_place_shape cave shape pos =
 ;;
 
 let drop_shape cave stream =
-  let shape = get_shape cave.rock_idx in
+  let shape = get_shape cave.rocks in
   let rec drop pos =
     let stream_pos = Vec2.(pos + stream.(cave.stream_idx)) in
     cave.stream_idx <- next_stream_idx stream cave.stream_idx;
@@ -109,9 +111,7 @@ let drop_shape cave stream =
     then drop drop_pos
     else place_shape cave shape pos
   in
-  cave.rock_idx <- next_shape_idx cave.rock_idx;
-  cave.dropped_rocks <- cave.dropped_rocks + 1;
-  cave.remaining_rocks <- cave.remaining_rocks - 1;
+  cave.rocks <- cave.rocks + 1;
   drop Vec2.{ x = 2; y = Int.(cave.grid_height + 4) }
 ;;
 
@@ -121,24 +121,24 @@ let normalize_topo topo =
 ;;
 
 let rec drop_shapes cave stream =
-  let cache_key =
-    { rock_idx = cave.rock_idx
+  let loop_key =
+    { rock = get_shape cave.rocks
     ; stream_idx = cave.stream_idx
     ; normalized_topo = normalize_topo cave.grid_topo
     }
   in
-  match Hashtbl.get cave.loop_cache cache_key with
-  | Some (cached_height, cached_dropped_rocks)
-    when let rock_delta = cave.dropped_rocks - cached_dropped_rocks in
-         rock_delta < cave.remaining_rocks ->
-    let rock_delta = cave.dropped_rocks - cached_dropped_rocks in
-    let height_delta = cave.height - cached_height in
-    cave.height <- cave.height + height_delta;
-    cave.dropped_rocks <- cave.dropped_rocks + rock_delta;
-    cave.remaining_rocks <- cave.remaining_rocks - rock_delta;
+  let remaining_rocks cave = cave.total_rocks - cave.rocks in
+  let loop_rocks loop = cave.rocks - loop.rock_idx in
+  let loop_fits loop = loop_rocks loop < remaining_rocks cave in
+  match Hashtbl.get cave.loop_cache loop_key with
+  | Some loop when loop_fits loop ->
+    let loop_height = cave.height - loop.height in
+    cave.height <- cave.height + loop_height;
+    cave.rocks <- cave.rocks + loop_rocks loop;
     drop_shapes cave stream
-  | _ when cave.remaining_rocks > 0 ->
-    Hashtbl.add cave.loop_cache cache_key (cave.height, cave.dropped_rocks);
+  | _ when remaining_rocks cave > 0 ->
+    let loop = { height = cave.height; rock_idx = cave.rocks } in
+    Hashtbl.add cave.loop_cache loop_key loop;
     drop_shape cave stream;
     drop_shapes cave stream
   | _ -> ()