From a8f3d96abf304e24e9aa31ed0578f6d3d4c95e03 Mon Sep 17 00:00:00 2001
From: Matthew Gretton-Dann <matt@gretton-dann.org.uk>
Date: Tue, 17 Dec 2024 10:29:59 +0000
Subject: [PATCH] 2024 day 17 part 2.

---
 bin/day2417.ml | 66 ++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 58 insertions(+), 8 deletions(-)

diff --git a/bin/day2417.ml b/bin/day2417.ml
index 0d156e0..cd84e35 100644
--- a/bin/day2417.ml
+++ b/bin/day2417.ml
@@ -6,7 +6,13 @@ type vm = {
   code : int array;
   out : int list;
 }
+(** Virtual Machine definition
 
+    Consists of three general-purpose registers: A, B, C, an instruction pointer
+    (ip), and array of code, and a list of output. The output list is in reverse
+    order of output. *)
+
+(** Generate a VM from a list of strings *)
 let vm_of_strings lst =
   let re_val = Str.regexp {|Register [ABC]: \([0-9]+\)|} in
   let re_prog = Str.regexp {|Program: \([0-9,]+\)|} in
@@ -39,10 +45,18 @@ let vm_of_strings lst =
   assert (List.is_empty lst);
   { a; b; c; ip; code; out }
 
+(** Generate a VM from a file *)
 let vm_of_file fname = Aoc.strings_of_file fname |> vm_of_strings
+
+(** [is_halted vm] returns true if the VM is halted. *)
+let is_halted vm = vm.ip >= Array.length vm.code || vm.ip < 0
+
+(** [get_literal vm] returns the Literal value from the current IP + 1 in [vm].
+*)
 let get_literal vm = vm.code.(vm.ip + 1)
 
-let print_combo vm =
+(** [print_combo vm] prints the combo operand at the current IP + 1 in [vm]. *)
+let[@warning "-32"] print_combo vm =
   match vm.code.(vm.ip + 1) with
   | 0 | 1 | 2 | 3 -> print_int vm.code.(vm.ip + 1)
   | 4 -> Printf.printf "A (=%d)" vm.a
@@ -51,6 +65,8 @@ let print_combo vm =
   | 7 -> failwith "print_combo reserved"
   | _ -> failwith "print_combo not 3-bit"
 
+(** [get_combo vm] returns the value of interpreting the combo operant at IP + 1
+    in [vm]. *)
 let get_combo vm =
   match vm.code.(vm.ip + 1) with
   | 0 | 1 | 2 | 3 -> vm.code.(vm.ip + 1)
@@ -60,9 +76,8 @@ let get_combo vm =
   | 7 -> failwith "get_combo reserved"
   | _ -> failwith "get_combo not 3-bit"
 
-let is_halted vm = vm.ip >= Array.length vm.code || vm.ip < 0
-
-let print_insn vm =
+(** [print_insn vm] prints the current instruction at VM. *)
+let[@warning "-32"] print_insn vm =
   Printf.printf "%d: " vm.ip;
   if is_halted vm then print_endline "Halted"
   else Printf.printf "%d " vm.code.(vm.ip);
@@ -88,10 +103,12 @@ let print_insn vm =
   | _ -> failwith "print_insn");
   Printf.printf " A=%d B=%d C=%d\n" vm.a vm.b vm.c
 
+(** [execute_insn vm] executes the current instruction in [vm] returns an
+    updated VM. *)
 let execute_insn vm =
   if is_halted vm then vm
-  else (
-    print_insn vm;
+  else
+    (*print_insn vm;*)
     match vm.code.(vm.ip) with
     | 0 -> { vm with a = vm.a / (1 lsl get_combo vm); ip = vm.ip + 2 }
     | 1 -> { vm with b = vm.b lxor get_literal vm; ip = vm.ip + 2 }
@@ -103,14 +120,47 @@ let execute_insn vm =
     | 5 -> { vm with out = (get_combo vm land 7) :: vm.out; ip = vm.ip + 2 }
     | 6 -> { vm with b = vm.a / (1 lsl get_combo vm); ip = vm.ip + 2 }
     | 7 -> { vm with c = vm.a / (1 lsl get_combo vm); ip = vm.ip + 2 }
-    | _ -> failwith "execute_insn")
+    | _ -> failwith "execute_insn"
 
+(** [execute_until_halted vm] executes a VM until it runs out of steam. *)
 let rec execute_until_halted vm =
   match is_halted vm with
   | true -> vm
   | false -> execute_until_halted (execute_insn vm)
 
+(** [string_of_ouput vm] gives the output of [vm]. *)
 let string_of_output vm =
   List.rev vm.out |> List.map string_of_int |> String.concat ","
 
-let _ = Aoc.main vm_of_file [ (string_of_output, execute_until_halted) ]
+(** [scan_digit acc ip vm] updates the acc for A so that the output of running
+    VM agrees in the last [ip] digits with the input program. *)
+let scan_digit acc ip vm =
+  let rec impl v =
+    let acc = acc + (v lsl (3 * ip)) in
+    let vm = { vm with a = acc } in
+    let vm = execute_until_halted vm in
+    let out = List.rev vm.out in
+    if List.length out <= ip then impl (v + 1)
+    else if
+      List.of_seq (Seq.drop ip (List.to_seq out))
+      = List.of_seq (Seq.drop ip (Array.to_seq vm.code))
+    then acc
+    else impl (v + 1)
+  in
+  impl 0
+
+(** [scan_all vm] generates the input A which means the output of executing [vm]
+    is the same as the input program. *)
+let scan_all vm =
+  let rec impl acc ip =
+    if ip < 0 || ip >= Array.length vm.code then { vm with a = acc }
+    else impl (scan_digit acc ip vm) (ip - 1)
+  in
+  impl 0 (Array.length vm.code - 1)
+
+(** [string_of_a vm] returns the A register of [vm]. *)
+let string_of_a vm = string_of_int vm.a
+
+let _ =
+  Aoc.main vm_of_file
+    [ (string_of_output, execute_until_halted); (string_of_a, scan_all) ]