From fcd0a9585dd4673c9972f20925009ac71d9ef5d8 Mon Sep 17 00:00:00 2001
From: Matt Russell <MattRussellUK@gmail.com>
Date: Sat, 7 Dec 2024 14:05:54 +0000
Subject: [PATCH] wip

---
 Aoc2024/Day07/Solve.lean | 35 +++++++++++------------------------
 Aoc2024/Day07/Types.lean |  7 +++++++
 Aoc2024/Utils.lean       |  9 +++++++++
 3 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/Aoc2024/Day07/Solve.lean b/Aoc2024/Day07/Solve.lean
index b2d8390..584db91 100644
--- a/Aoc2024/Day07/Solve.lean
+++ b/Aoc2024/Day07/Solve.lean
@@ -5,35 +5,23 @@ import Aoc2024.Day07.Types
 
 -- Part 1
 
-private def concatenate (a : Int) (b : Int) : Int := (reprStr a ++ reprStr b).toInt!
-
 private def evalNumbers (numbers : List Int) (operators : List Operator): Int :=
-  match numbers with
-    | [] => 0
-    | h :: t =>
-      t.zip operators |>.foldl (fun acc (n, op) => match op with
-        | .add => acc + n
-        | .multiply => acc * n
-        | .concatenation => concatenate acc n
-      ) h
+  numbers.zip (.add :: operators) |>.foldl (λ acc (n, op) => op.fn acc n) 0
 
 #guard evalNumbers [11, 6, 16, 20] [.add, .multiply, .add] == 292
 
-private def replicateM (n : Nat) (options : List α) : List (List α) :=
-  let rec loop (n : Nat) (acc : List (List α)) : List (List α) :=
-    match n with
-    | 0 => acc
-    | n + 1 => loop n (acc.bind λ l => options.map (λ o => o :: l))
-  loop n [[]]
-#guard replicateM 0 [1, 2] == [[]]
-#guard replicateM 3 [1, 2] == [[1, 1, 1], [2, 1, 1], [1, 2, 1], [2, 2, 1], [1, 1, 2], [2, 1, 2], [1, 2, 2], [2, 2, 2]]
+private def Equation.isTrue (equation : Equation) (operators : List Operator) : Bool :=
+  evalNumbers equation.numbers operators == equation.testValue
+
+private def Equation.operatorPositions (equation : Equation) : Nat := equation.numbers.length - 1
 
 private def canEquationPossiblyBeTrue (operators : List Operator) (equation : Equation) : Bool :=
-  let operators := replicateM equation.numbers.length operators
-  operators.any λ ops => evalNumbers equation.numbers ops == equation.testValue
+  replicateM equation.operatorPositions operators |>.any equation.isTrue
+
+private def solveGeneric (operators : List Operator) (equations : List Equation) : Int :=
+  equations.filter (canEquationPossiblyBeTrue operators) |>.sumBy (·.testValue)
 
-private def solvePart1 (equations : List Equation) : Int :=
-  equations.filter (canEquationPossiblyBeTrue [.add, .multiply]) |>.sumBy (·.testValue)
+private def solvePart1 : List Equation -> Int := solveGeneric [.add, .multiply]
 
 def parseAndSolvePart1 (s : String): Except String Int := parseEquations s |>.map solvePart1
 
@@ -41,8 +29,7 @@ def parseAndSolvePart1 (s : String): Except String Int := parseEquations s |>.ma
 
 -- Part 2
 
-private def solvePart2 (equations : List Equation) : Int :=
-  equations.filter (canEquationPossiblyBeTrue [.add, .multiply, .concatenation]) |>.sumBy (·.testValue)
+private def solvePart2 : List Equation -> Int := solveGeneric [.add, .multiply, .concatenation]
 
 def parseAndSolvePart2 (s : String): Except String Int := parseEquations s |>.map solvePart2
 
diff --git a/Aoc2024/Day07/Types.lean b/Aoc2024/Day07/Types.lean
index beb80b5..b0233c3 100644
--- a/Aoc2024/Day07/Types.lean
+++ b/Aoc2024/Day07/Types.lean
@@ -8,3 +8,10 @@ inductive Operator where
   | multiply
   | concatenation
 deriving Repr, BEq, Hashable, Inhabited
+
+private def concatenate (a : Int) (b : Int) : Int := (reprStr a ++ reprStr b).toInt!
+
+def Operator.fn : Operator -> Int -> Int -> Int
+  | .add, a, b => a + b
+  | .multiply, a, b => a * b
+  | .concatenation, a, b => concatenate a b
diff --git a/Aoc2024/Utils.lean b/Aoc2024/Utils.lean
index f888abc..1a39f7b 100644
--- a/Aoc2024/Utils.lean
+++ b/Aoc2024/Utils.lean
@@ -163,3 +163,12 @@ namespace Option
   #guard (some 42).getOrThrow "Error" == Except.ok 42
   #guard (none : Option Int).getOrThrow "Error" == Except.error "Error"
 end Option
+
+def replicateM (n : Nat) (options : List α) : List (List α) :=
+  let rec loop (n : Nat) (acc : List (List α)) : List (List α) :=
+    match n with
+    | 0 => acc
+    | n + 1 => loop n (acc.bind λ l => options.map (λ o => o :: l))
+  loop n [[]]
+#guard replicateM 0 [1, 2] == [[]]
+#guard replicateM 3 [1, 2] == [[1, 1, 1], [2, 1, 1], [1, 2, 1], [2, 2, 1], [1, 1, 2], [2, 1, 2], [1, 2, 2], [2, 2, 2]]