--- a/2023/day17.ml	Sun Dec 31 09:21:55 2023 +0100
+++ b/2023/day17.ml	Sun Dec 31 09:37:40 2023 +0100
@@ -159,9 +159,13 @@
         {
           r = r';
           c = c';
+          (* Direction this position is facing when entering its tile *)
           dir = dir';
+          (* Allow tracking of path by remembering previous tile *)
           prev = (r, c);
+          (* Number of straight tiles in a row *)
           straight =
+            (* careful, tricky! (not sure if I got this right, either *)
             (if Int.(straight = 1) || equal_direction dir dir' then straight + 1
              else 1);
           heatloss = pos.heatloss + (field_at field r' c').value;
@@ -169,6 +173,40 @@
     in
     List.map neighbors' ~f:pos_of_neighbor
 
+  (* Apply Dijkstra's algorithm (or something like that...)
+     to find the shortest path according to restrictions. *)
+  let solve field =
+    let dstr, dstc = dst field in
+    let rec loop (q : Pospq.t) =
+      (* No path could be found, options exhausted *)
+      if Pospq.is_empty q then None
+      else
+        let q, pos = Pospq.remove_min_elt_exn q in
+        (* Arrived at destination *)
+        if Int.equal pos.r dstr && Int.equal pos.c dstc then (
+          field_update field pos.r pos.c (fun v ->
+              { v with min_so_far = pos.heatloss; prev = pos.prev });
+          Some pos.heatloss)
+        else
+          (* Check if tile has been visited before *)
+          let min_so_far = (field_at field pos.r pos.c).min_so_far in
+          if min_so_far = Int.max_value then (
+            (* We are first; found minimal path to tile; update cost and previous tile *)
+            field_update field pos.r pos.c (fun v ->
+                { v with min_so_far = pos.heatloss; prev = pos.prev });
+            let next = next_options field pos in
+            (* Add all next options to the queue *)
+            let q =
+              List.fold_left next ~init:q ~f:(fun q' opt ->
+                  Pospq.add q' ~prio:opt.heatloss opt)
+            in
+            loop q
+            (* Already visited this tile, skip *))
+          else loop q
+    in
+    loop (List.fold_left initial ~init:Pospq.empty ~f:(Pospq.add ~prio:0))
+
+  (* Recursively follow `prev` links in the field array. *)
   let rec trace_path ?(acc = []) field r c =
     match (r, c) with
     | 0, 0 -> (0, 0) :: acc
@@ -177,16 +215,19 @@
         let r', c' = entry.prev in
         trace_path ~acc:((r, c) :: acc) field r' c'
 
+  (* Trace back the path from the destination to the start *)
   let start_trace_path field =
     let r, c = dst field in
     trace_path field r c
 
+  (* Create a 2D array mapping visited tiles *)
   let visualizer_field field path =
     let a = Array.map field.field ~f:(fun _ -> 0) in
     let f (r, c) = a.(rc_to_ix field r c) <- 1 in
     List.iter path ~f;
     a
 
+  (* Print the 2D array *)
   let print_visualizer_field field r c =
     let rows = Sequence.range 0 r and cols = Sequence.range 0 c in
     let f r' =
@@ -195,36 +236,10 @@
     in
     Sequence.iter rows ~f
 
+  (* Create and print the path on a 2D map *)
   let visualize_field field path =
     let f = visualizer_field field path in
     print_visualizer_field f field.r field.c
-
-  (* Apply Dijkstra's algorithm (or something like that...)
-     to find the shortest path according to restrictions. *)
-  let solve field =
-    let dstr, dstc = dst field in
-    let rec loop (q : Pospq.t) =
-      if Pospq.is_empty q then None
-      else
-        let q, pos = Pospq.remove_min_elt_exn q in
-        if Int.equal pos.r dstr && Int.equal pos.c dstc then (
-          field_update field pos.r pos.c (fun v ->
-              { v with min_so_far = pos.heatloss; prev = pos.prev });
-          Some pos.heatloss)
-        else
-          let min_so_far = (field_at field pos.r pos.c).min_so_far in
-          if min_so_far = Int.max_value then (
-            field_update field pos.r pos.c (fun v ->
-                { v with min_so_far = pos.heatloss; prev = pos.prev });
-            let next = next_options field pos in
-            let q =
-              List.fold_left next ~init:q ~f:(fun q' opt ->
-                  Pospq.add q' ~prio:opt.heatloss opt)
-            in
-            loop q)
-          else loop q
-    in
-    loop (List.fold_left initial ~init:Pospq.empty ~f:(Pospq.add ~prio:0))
 end
 
 let () =
@@ -237,5 +252,5 @@
     (Sexp.to_string_hum
        (List.sexp_of_t
           (fun (r, c) -> Sexp.List [ Int.sexp_of_t r; Int.sexp_of_t c ])
-        trace));
+          trace));
   Part1.visualize_field field trace