2019 day 2 in Scala - tail recursive, more functional

2019/scala/02/day02/src/main/scala/Main.scala

@@ -1,57 +1,55 @@
-import scala.collection.mutable.ArrayBuffer
+import scala.collection.immutable.Vector
 import scala.io.Source
+import scala.annotation.tailrec
 
 case class UnknownInstruction() extends Exception
 case class SolutionNotFound() extends Exception
 
 @main def hello(): Unit =
-  val input: ArrayBuffer[Int] = Source.fromFile("input").getLines().next().split(",").map(_.toInt).to(ArrayBuffer)
+  val input: Vector[Int] = Source.fromFile("input").getLines().next().split(",").map(_.toInt).to(Vector)
 
-  printf("Part 1: %d\n", solve1(input.clone()))
+  printf("Part 1: %d\n", solve1(input))
   printf("Part 2: %d\n", solve2(input))
 
-def step(op: Int, memory : ArrayBuffer[Int]): (Option[Int], ArrayBuffer[Int]) =
+def step(op: Int, memory : Vector[Int]): (Option[Int], Vector[Int]) =
   memory(op) match
-    case 1  => (Some(op + 4), add(op, memory))
-    case 2  => (Some(op + 4), mul(op, memory))
+    case 1 | 2 => (Some(op + 4), binary_op(op, memory))
     case 99 => (None, memory)
     case _  => throw new UnknownInstruction
 
-def run(memory: ArrayBuffer[Int]): ArrayBuffer[Int] =
-  var op: Option[Int] = Some(0)
-  var mem: ArrayBuffer[Int] = memory
-  while(op.nonEmpty) {
-    val (nop, nmemory) = step(op.get, memory)
-    op = nop
-    mem = nmemory
+def run(memory: Vector[Int]): Vector[Int] =
+  @tailrec
+  def inner_run(state: (Option[Int], Vector[Int])): Vector[Int] = state match {
+    case (None, memory) => memory
+    case (Some(ip), memory) => inner_run(step(ip, memory))
   }
-  mem
 
-def add(op: Int, memory: ArrayBuffer[Int]): ArrayBuffer[Int] =
-  memory.update(memory(op + 3), indirect(op + 1, memory) + indirect(op + 2, memory))
-  memory
+  inner_run(Some(0), memory)
 
-def mul(op: Int, memory: ArrayBuffer[Int]): ArrayBuffer[Int] =
-  memory.update(memory(op + 3), indirect(op + 1, memory) * indirect(op + 2, memory))
-  memory
+def binary_op(op: Int, memory: Vector[Int]): Vector[Int] =
+  val o1 = indirect(op + 1, memory)
+  val o2 = indirect(op + 2, memory)
+  val res = memory(op) match {
+    case 1 => o1 + o2
+    case 2 => o1 * o2
+    case _  => throw new UnknownInstruction
+  }
+  memory.updated(memory(op + 3), res)
 
-def indirect(address: Int, memory: ArrayBuffer[Int]): Int =
+def indirect(address: Int, memory: Vector[Int]): Int =
   memory(memory(address))
 
-def run_modified(m1: Int, m2: Int, memory: ArrayBuffer[Int]): Int =
-  memory(1) = m1
-  memory(2) = m2
-  var done = run(memory)
-  done(0)
+def run_modified(m1: Int, m2: Int, memory: Vector[Int]): Int =
+  run(memory.updated(1, m1).updated(2, m2))(0)
 
-def solve1(memory: ArrayBuffer[Int]): Int =
+def solve1(memory: Vector[Int]): Int =
   run_modified(12, 2, memory)
 
-def solve2(memory: ArrayBuffer[Int]): Int =
+def solve2(memory: Vector[Int]): Int =
   val options = for i <- 0 to 99
                     j <- 0 to 99
                 yield (i, j)
-  options.find((i, j) => run_modified(i, j, memory.clone()) == 19690720) match
+  options.find((i, j) => run_modified(i, j, memory) == 19690720) match
     case Some(i, j) => 100 * i + j
     case None => throw new SolutionNotFound
     

2019/scala/02/day02/src/test/scala/MySuite.scala

@@ -1,42 +1,42 @@
-import scala.collection.mutable.ArrayBuffer
+import scala.collection.immutable.Vector
 
 // For more information on writing tests, see
 // https://scalameta.org/munit/docs/getting-started.html
 class MySuite extends munit.FunSuite {
   test("step") {
-    val in = ArrayBuffer(1,0,0,0,99)
+    val in = Vector(1,0,0,0,99)
     val (next, out) = step(0, in)
     assertEquals(next, Some(4))
-    assertEquals(out, ArrayBuffer(2,0,0,0,99))
+    assertEquals(out, Vector(2,0,0,0,99))
   }
 
   test("step") {
-    val in = ArrayBuffer(2,3,0,3,99)
+    val in = Vector(2,3,0,3,99)
     val (next, out) = step(0, in)
     assertEquals(next, Some(4))
-    assertEquals(out, ArrayBuffer(2,3,0,6,99))
+    assertEquals(out, Vector(2,3,0,6,99))
   }
 
   test("step") {
-    val in = ArrayBuffer(2,4,4,5,99,0)
+    val in = Vector(2,4,4,5,99,0)
     val (next, out) = step(0, in)
     assertEquals(next, Some(4))
-    assertEquals(out, ArrayBuffer(2,4,4,5,99,9801))
+    assertEquals(out, Vector(2,4,4,5,99,9801))
   }
 
   test("step") {
-    val in = ArrayBuffer(1,1,1,4,99,5,6,0,99)
+    val in = Vector(1,1,1,4,99,5,6,0,99)
     val (next, out) = step(0, in)
     assertEquals(next, Some(4))
 
     val (next2, out2) = step(4, out)
     assertEquals(next2, Some(8))
-    assertEquals(out2, ArrayBuffer(30,1,1,4,2,5,6,0,99))
+    assertEquals(out2, Vector(30,1,1,4,2,5,6,0,99))
   }
 
   test("run") {
-    val in = ArrayBuffer(1,1,1,4,99,5,6,0,99)
+    val in = Vector(1,1,1,4,99,5,6,0,99)
     val out = run(in)
-    assertEquals(out, ArrayBuffer(30,1,1,4,2,5,6,0,99))
+    assertEquals(out, Vector(30,1,1,4,2,5,6,0,99))
   }
 }