Day8.scala

package sungkmi.aoc2020.day8

import scala.collection.immutable.{BitSet, LazyList}

enum Instruction:
  case Acc(arg: Int) extends Instruction
  case Jmp(arg: Int) extends Instruction
  case Nop(arg: Int) extends Instruction

type Instructions = IndexedSeq[Instruction]

case class ProgramState(pc: Int, acc: Int, runSet: BitSet)

object ProgramState:
  val init: ProgramState = ProgramState(0, 0, BitSet.empty)

extension (s: ProgramState):
  def run(instruction: Instruction): ProgramState = instruction match
    case Acc(arg) => ProgramState(s.pc + 1, s.acc + arg, s.runSet + s.pc)
    case Jmp(arg) => ProgramState(s.pc + arg, s.acc, s.runSet + s.pc)
    case Nop(_) => ProgramState(s.pc + 1, s.acc, s.runSet + s.pc)

import Instruction._

def parseInstruction(s: String): Instruction =
  val splitted = s `split` " "
  val arg = splitted.tail.head
  splitted.head match
    case "acc" => Instruction.Acc(arg.toInt)
    case "jmp" => Jmp(arg.toInt)
    case "nop" => Nop(arg.toInt)

def getResult(instructions: Instructions): Int =
  @annotation.tailrec
  def loop(state: ProgramState): Int =
   if state.runSet(state.pc) then state.acc else loop(state.run(instructions(state.pc)))
  loop(ProgramState.init)

def getFiniteResult(instructions: Instructions): Option[Int] =
  @annotation.tailrec
  def loop(state: ProgramState): Option[Int] =
    if state.runSet(state.pc) then None
    else if state.pc == instructions.size then Some(state.acc)
    else loop(state.run(instructions(state.pc)))
  loop(ProgramState.init)

def modifiedInstructions(instructions: Instructions): Seq[Instructions] =
  def modifiedAt(index: Int): Option[Instructions] =
    val (front, end) = instructions `splitAt` index
    end.head match
      case Jmp(arg) => Some(front ++ (Nop(arg) +: end.tail))
      case Nop(arg) => Some(front ++ (Jmp(arg) +: end.tail))
      case _ => None
  0 until instructions.size flatMap modifiedAt

def run(input: String): Int =
  val instructions = input.split("\n").map(parseInstruction).toIndexedSeq
  getResult(instructions)

def run2(input: String): Int =
  val instructions = input.split("\n").map(parseInstruction).toIndexedSeq
  modifiedInstructions(instructions).to(LazyList).flatMap(getFiniteResult).head

@main def part1: Unit =
  val ans = run(input)
  println(ans)

@main def part2: Unit =
  val ans = run2(input)
  println(ans)

//lazy val input: String = """acc -1

Day8Test.scala

package sungkmi.aoc2020.day8

class Day8Test extends munit.FunSuite {
  val testInput = """nop +0
acc +1
jmp +4
acc +3
jmp -3
acc -99
acc +1
jmp -4
acc +6"""

  test("run") {
    assertEquals(run(testInput), 5)
  }

  test("run2") {
    assertEquals(run2(testInput), 8)
  }
}