AoC 2020 Day 8

day8.swift

import Foundation

enum Operation: String {
    case noOperation = "nop"
    case accumulator = "acc"
    case jump = "jmp"
}

struct Instruction {
    let operation: Operation
    let argument: Int
}

extension Instruction {
    init(_ instructionString: String) {
        let separated = instructionString.split(separator: " ")
        operation = Operation(rawValue: String(separated[0]))!
        argument = Int(separated[1])!
    }
}

struct InstructionIterator: IteratorProtocol {
    let instructions: [Instruction]
    var position: Int

    var accumulator: Int
    var isInfinite: Bool

    private var iteratedPositions: Set<Int>
    private var shouldReplaceInstruction: Bool

    mutating func next() -> Instruction? {
        // Finish if the program reaches last instruction + 1
        if position == instructions.count { return nil }

        let instruction = instructions[position]

        if shouldReplaceInstruction {
            switch instruction.operation {
            case .noOperation, .jump:
                if let accumulator = terminationAccumulator(whenReplacing: instruction) {
                    self.accumulator = accumulator
                    return nil
                }
            case .accumulator:
                break
            }
        }

        isInfinite = !iteratedPositions.insert(position).inserted
        guard !isInfinite else { return nil }

        switch instruction.operation {
        case .noOperation:
            position = position + 1
        case .accumulator:
            accumulator += instruction.argument
            position = position + 1
        case .jump:
            position += instruction.argument
        }

        return instruction
    }

    private func terminationAccumulator(whenReplacing instruction: Instruction) -> Int? {
        var iterator = makeIterator(replacing: instruction)
        while iterator.next() != nil { }

        if !iterator.isInfinite {
            return iterator.accumulator
        }

        return nil
    }

    private func makeIterator(replacing instruction: Instruction) -> InstructionIterator {
        var newInstructions = instructions
        newInstructions[position] = Instruction(
            operation: instruction.operation == .noOperation ? .jump : .noOperation,
            argument: instruction.argument
        )

        return InstructionIterator(
            instructions: newInstructions,
            position: position,
            accumulator: accumulator,
            isInfinite: isInfinite,
            iteratedPositions: iteratedPositions,
            shouldReplaceInstruction: false
        )
    }
}

extension InstructionIterator {
    init(_ instructions: [Instruction], shouldReplaceInstruction: Bool = false) {
        self.instructions = instructions
        position = instructions.startIndex
        accumulator = 0
        isInfinite = false
        iteratedPositions = []
        self.shouldReplaceInstruction = shouldReplaceInstruction
    }
}

final class Day820: Day {
    override func perform() {
        let input = String.input(forDay: 8, year: 2020)
        let instructions = input
            .split(separator: "\n")
            .map(String.init).map(Instruction.init)

        // part 1
        var iterator = InstructionIterator(instructions)
        while iterator.next() != nil { }
        print(iterator.accumulator)

        // part 2
        var iterator2 = InstructionIterator(instructions, shouldReplaceInstruction: true)
        while iterator2.next() != nil { }
        print(iterator2.accumulator)
    }
}