From 5d99f413de37e8248c8ea263741dc5baa2e04e2c Mon Sep 17 00:00:00 2001
From: Matthew Gretton-Dann <matt@gretton-dann.org.uk>
Date: Tue, 3 Dec 2024 08:28:34 +0000
Subject: [PATCH] 2024 day 3 part 2.

---
 bin/day2403.ml | 39 +++++++++++++++++++++++++++++++--------
 1 file changed, 31 insertions(+), 8 deletions(-)

diff --git a/bin/day2403.ml b/bin/day2403.ml
index b82fb6d..3d8a96f 100644
--- a/bin/day2403.ml
+++ b/bin/day2403.ml
@@ -1,7 +1,13 @@
+(** [matched_group_opt n s] returns [Some (Str.matched_group n s)] if the [n]th
+    group matched, or [None] if the group did not match. *)
 let matched_group_opt n s = try Some (Str.matched_group n s) with _ -> None
 
-let find_nums s =
-  let r = Str.regexp {|mul(\([0-9][0-9]?[0-9]?\),\([0-9][0-9]?[0-9]?\))|} in
+(** [find_nums s] returns the list of instructions given in [s]. *)
+let instrs_of_string s =
+  let r =
+    Str.regexp
+      {|do()\|don't()\|mul(\([0-9][0-9]?[0-9]?\),\([0-9][0-9]?[0-9]?\))|}
+  in
   let rec impl acc pos =
     try
       let _ = Str.search_forward r s pos in
@@ -13,11 +19,28 @@ let find_nums s =
         (Str.match_end ())
     with Not_found -> acc
   in
-  impl [] 0
+  List.rev (impl [] 0)
 
-let nums_from_file fname =
-  Aoc.strings_from_file fname |> List.map find_nums |> List.concat
+let instrs_of_file fname =
+  Aoc.strings_from_file fname |> List.map instrs_of_string |> List.concat
 
-let mac acc a b = acc + (Option.value a ~default:0 * Option.value b ~default:0)
-let day2403a = List.fold_left (fun acc (_, a, b) -> mac acc a b) 0
-let _ = Aoc.main nums_from_file [ (string_of_int, day2403a) ]
+(** [mac_opt acc a b] returns [acc + a' * b'] if [a = Some a'] and 
+    [b = Some b'].  If either [a] or [b] are [None] then the result is [acc]. *)
+let mac_opt acc a b =
+  acc + (Option.value a ~default:0 * Option.value b ~default:0)
+
+let day2403a = List.fold_left (fun acc (_, a, b) -> mac_opt acc a b) 0
+
+let day2403b lst =
+  let rec impl acc enabled = function
+    | [] -> acc
+    | ("do()", _, _) :: t -> impl acc true t
+    | ("don't()", _, _) :: t -> impl acc false t
+    | (_, a, b) :: t ->
+        if enabled then impl (mac_opt acc a b) true t else impl acc false t
+  in
+  impl 0 true lst
+
+let _ =
+  Aoc.main instrs_of_file
+    [ (string_of_int, day2403a); (string_of_int, day2403b) ]