Extract 2024 day 17 part 2 solutions

src/main/scala/eu/sim642/adventofcode2024/Day17.scala

@@ -66,56 +66,66 @@ object Day17 {
 
   def runOutput(input: Input): String = runOutput0(input).mkString(",")
 
-  def findQuineA(input: Input): Int = {
-    Iterator.from(0)
-      .find(newA => runOutput0(Input(input.registers.copy(a = newA), input.program)) == input.program)
-      .get
+  trait Part2Solution {
+    def findQuineA(input: Input): Long
   }
 
-  def myProg(initialA: Int, expectedOutputs: Iterator[Int]): Boolean = {
-    var a: Int = initialA
-    var b: Int = 0
-    var c: Int = 0
-    while (a != 0) {
-      b = a & 0b111
-      b = b ^ 1
-      c = a >> b
-      b = b ^ 5
-      b = b ^ c
-      a = a >> 3
-      if ((b & 0b111) != expectedOutputs.next())
-        return false
+  object NaivePart2Solution extends Part2Solution {
+    override def findQuineA(input: Input): Long = {
+      Iterator.from(0)
+        .find(newA => runOutput0(Input(input.registers.copy(a = newA), input.program)) == input.program)
+        .get
     }
-    !expectedOutputs.hasNext
   }
 
-  def findQuineA2(input: Input): Int = {
-    Iterator.from(0)
-      .find(newA => myProg(newA, input.program.iterator))
-      .get
-  }
+  object SemiNaivePart2Solution extends Part2Solution {
+    def myProg(initialA: Int, expectedOutputs: Iterator[Int]): Boolean = {
+      var a: Int = initialA
+      var b: Int = 0
+      var c: Int = 0
+      while (a != 0) {
+        b = a & 0b111
+        b = b ^ 1
+        c = a >> b
+        b = b ^ 5
+        b = b ^ c
+        a = a >> 3
+        if ((b & 0b111) != expectedOutputs.next())
+          return false
+      }
+      !expectedOutputs.hasNext
+    }
 
-  def findQuineA3(input: Input): Long = {
-    val ctx = new Context(Map("model" -> "true").asJava)
-    import ctx._
-    val s = mkSolver()
-
-    val bits = input.program.size * 3
-    val initialA = mkBVConst("initialA", bits)
-
-    for ((instruction, i) <- input.program.zipWithIndex) {
-      val a = mkBVLSHR(initialA, mkBV(i * 3, bits))
-      var b = mkBVAND(a, mkBV(7, bits))
-      b = mkBVXOR(b, mkBV(1, bits))
-      val c = mkBVLSHR(a, b)
-      b = mkBVXOR(b, mkBV(5, bits))
-      b = mkBVXOR(b, c)
-      val out = mkBVAND(b, mkBV(7, bits))
-      s.add(mkEq(out, mkBV(instruction, bits)))
+    override def findQuineA(input: Input): Long = {
+      Iterator.from(0)
+        .find(newA => myProg(newA, input.program.iterator))
+        .get
     }
+  }
 
-    assert(s.check() == Status.SATISFIABLE)
-    s.getModel.evaluate(initialA, false).toString.toLong
+  object Z3Part2Solution extends Part2Solution {
+    override def findQuineA(input: Input): Long = {
+      val ctx = new Context(Map("model" -> "true").asJava)
+      import ctx._
+      val s = mkSolver()
+
+      val bits = input.program.size * 3
+      val initialA = mkBVConst("initialA", bits)
+
+      for ((instruction, i) <- input.program.zipWithIndex) {
+        val a = mkBVLSHR(initialA, mkBV(i * 3, bits))
+        var b = mkBVAND(a, mkBV(7, bits))
+        b = mkBVXOR(b, mkBV(1, bits))
+        val c = mkBVLSHR(a, b)
+        b = mkBVXOR(b, mkBV(5, bits))
+        b = mkBVXOR(b, c)
+        val out = mkBVAND(b, mkBV(7, bits))
+        s.add(mkEq(out, mkBV(instruction, bits)))
+      }
+
+      assert(s.check() == Status.SATISFIABLE)
+      s.getModel.evaluate(initialA, false).toString.toLong
+    }
   }
 
   def parseInput(input: String): Input = input match {
@@ -128,10 +138,9 @@ object Day17 {
   lazy val input: String = scala.io.Source.fromInputStream(getClass.getResourceAsStream("day17.txt")).mkString.trim
 
   def main(args: Array[String]): Unit = {
+    import Z3Part2Solution._
     println(runOutput(parseInput(input)))
-    println(findQuineA3(parseInput(input)))
-
-    // part 2: 164540892147389 - correct
+    println(findQuineA(parseInput(input)))
 
     // part 1: 4,5,0,4,7,4,3,0,0 - wrong (bst used literal not combo operand)
   }

src/test/scala/eu/sim642/adventofcode2024/Day17Test.scala

@@ -1,9 +1,17 @@
 package eu.sim642.adventofcode2024
 
-import Day17._
+import Day17.*
+import Day17Test.*
+import org.scalatest.Suites
 import org.scalatest.funsuite.AnyFunSuite
 
-class Day17Test extends AnyFunSuite {
+class Day17Test extends Suites(
+  new Part1Test,
+  new NaivePart2SolutionTest,
+  new Z3Part2SolutionTest,
+)
+
+object Day17Test {
 
   // TODO: small examples
 
@@ -35,17 +43,31 @@ class Day17Test extends AnyFunSuite {
       |
       |Program: 0,3,5,4,3,0""".stripMargin
 
-  test("Part 1 examples") {
-    assert(runOutput(parseInput(exampleInput)) == "4,6,3,5,6,3,5,2,1,0")
-    assert(runOutput(parseInput(exampleInput1)) == "0,1,2")
-    assert(runOutput(parseInput(exampleInput2)) == "4,2,5,6,7,7,7,7,3,1,0")
+  class Part1Test extends AnyFunSuite {
+    test("Part 1 examples") {
+      assert(runOutput(parseInput(exampleInput)) == "4,6,3,5,6,3,5,2,1,0")
+      assert(runOutput(parseInput(exampleInput1)) == "0,1,2")
+      assert(runOutput(parseInput(exampleInput2)) == "4,2,5,6,7,7,7,7,3,1,0")
+    }
+
+    test("Part 1 input answer") {
+      assert(runOutput(parseInput(input)) == "4,3,2,6,4,5,3,2,4")
+    }
   }
 
-  test("Part 1 input answer") {
-    assert(runOutput(parseInput(input)) == "4,3,2,6,4,5,3,2,4")
+  abstract class Part2SolutionExampleTest(part2Solution: Part2Solution) extends AnyFunSuite {
+    test("Part 2 examples") {
+      assert(part2Solution.findQuineA(parseInput(exampleInput3)) == 117440)
+    }
   }
 
-  test("Part 2 examples") {
-    assert(findQuineA(parseInput(exampleInput3)) == 117440)
+  abstract class Part2SolutionInputTest(part2Solution: Part2Solution) extends AnyFunSuite {
+    test("Part 2 input answer") {
+      assert(part2Solution.findQuineA(parseInput(input)) == 164540892147389L)
+    }
   }
+
+  class NaivePart2SolutionTest extends Part2SolutionExampleTest(NaivePart2Solution)
+
+  class Z3Part2SolutionTest extends Part2SolutionInputTest(Z3Part2Solution)
 }