2024 Day 17

DaySeventeen.hs

module DaySeventeen (wins, program, go, runSingleInput, workBackwards, partOne, evaluateOperand, runProgram, buildComputer, execute, Instruction(..), Operand(..), Combo(..), parsePair) where
import Debug.Trace
import Data.Bits (xor)
import Data.Maybe (isJust)

import Data.List.Split (chunksOf)
import qualified Data.Sequence as S

data Computer = Computer { instructionPointer :: Int, registerA :: Int, registerB :: Int, registerC :: Int } deriving (Ord, Eq, Show)

data Combo = LiteralZero | LiteralOne | LiteralTwo | LiteralThree | GetRegisterA | GetRegisterB | GetRegisterC deriving (Ord, Eq, Show, Enum)

data Operand = ComboOperand Combo | IntOperand Int deriving (Eq, Ord, Show)

data Instruction = Adv | Bxl | Bst | Jnz | Bxc | Out | Bdv | Cdv deriving (Ord, Eq, Show, Enum)
type Executable = (Instruction, Operand)

type RawInput = ((Int, Int, Int), [Int])

buildComputer :: RawInput -> (Computer, S.Seq Int)
buildComputer ((ra,rb,rc), commands) = (Computer 0 ra rb rc, S.fromList commands)

parseOperand :: Instruction -> Int -> Operand
parseOperand Adv n = ComboOperand (toEnum n)
parseOperand Bdv n = ComboOperand (toEnum n)
parseOperand Cdv n = ComboOperand (toEnum n)
parseOperand Bst n = ComboOperand (toEnum n)
parseOperand Bxl n = IntOperand n
parseOperand Jnz n = IntOperand n
parseOperand Bxc _ = IntOperand 0
parseOperand Out n = ComboOperand (toEnum n)

evaluateOperand :: Computer -> Operand -> Int
evaluateOperand computer (IntOperand i) = i
evaluateOperand computer (ComboOperand LiteralZero) = 0
evaluateOperand computer (ComboOperand LiteralOne) = 1
evaluateOperand computer (ComboOperand LiteralTwo) = 2
evaluateOperand computer (ComboOperand LiteralThree) = 3
evaluateOperand computer (ComboOperand GetRegisterA) = (registerA computer)
evaluateOperand computer (ComboOperand GetRegisterB) = (registerB computer)
evaluateOperand computer (ComboOperand GetRegisterC) = (registerC computer)


execute :: Computer -> Executable -> (Computer, Maybe Int)
execute computer (Adv, op) = (computer { instructionPointer=instructionPointer computer + 2, registerA=(registerA computer) `div` (2 ^ (evaluateOperand computer op))}, Nothing)
execute computer (Bdv, op) = (computer { instructionPointer=instructionPointer computer + 2, registerB=(registerA computer) `div` (2 ^ (evaluateOperand computer op))}, Nothing)
execute computer (Cdv, op) = (computer { instructionPointer=instructionPointer computer + 2, registerC=(registerA computer) `div` (2 ^ (evaluateOperand computer op))}, Nothing)

execute computer (Bxl, op) = (computer { instructionPointer=instructionPointer computer + 2, registerB=(registerB computer) `xor` (evaluateOperand computer op)}, Nothing)
execute computer (Bst, op) = (computer { instructionPointer=instructionPointer computer + 2, registerB=(evaluateOperand computer op `mod` 8)}, Nothing)
execute computer (Jnz, op) 
  | registerA computer == 0 = (computer {instructionPointer=(instructionPointer computer) + 2}, Nothing)
  | otherwise = (computer {instructionPointer=(evaluateOperand computer op)}, Nothing)
execute computer (Bxc, op) = (computer { instructionPointer=(instructionPointer computer) + 2, registerB=(registerB computer `xor` registerC computer)}, Nothing)
execute computer (Out, op) = (computer { instructionPointer=(instructionPointer computer) + 2}, Just (evaluateOperand computer op `mod` 8))


parsePair :: Int -> Int -> Executable
parsePair opcode argcode = (instruction, arg)
  where instruction = toEnum opcode
        arg = parseOperand instruction argcode


runProgram :: Computer -> S.Seq Int -> [Int]
runProgram computer instructions
  | instructionPointer computer >= S.length instructions = []
  | otherwise = prependMaybe newOutput $ runProgram newComputer instructions
  where (newComputer, newOutput) = execute computer executable
        opCode = instructions `S.index` (instructionPointer computer)
        argCode = instructions `S.index` (instructionPointer computer + 1)
        executable = parsePair opCode argCode
    
prependMaybe :: Maybe t -> [t] -> [t]
prependMaybe Nothing items = items
prependMaybe (Just e) items = e:items

sampleInput :: RawInput
sampleInput = ((729, 0, 0), [0,1,5,4,3,0])

partOne = buildComputer sampleInput

-- part 2
-- what is the program doing?

-- * Take register A mod 8, write to register B
-- * Register B bitwise OR with 3, write to register B
-- * Read register A, divide by (2^Bregister), write to register C
-- * Divide register A by 8
-- * Bitwise XOR registers B,C, write to register B
-- * Bitwise XOR register B with 5, write to register B
-- * Output register B
-- * Jump to register 0 (as long as A is not 0)


program = [2,4,1,3,7,5,0,3,4,3,1,5,5,5,3,0]

-- this is a simplified version of the program, which shows what is output according to the value in register A
go :: Int -> Int
go registerA = (((registerA `mod` 8) `xor` 3) `xor` (registerA `div` (2 ^ ((registerA `mod` 8) `xor` 3))) `xor` 5) `mod` 8

-- given an input which generates the last N digits of the program,
-- this will give an input which generates the last N+1 digits of the program
workBackwards :: Int -> Int -> [Int]
workBackwards dropN current = [(current * 8) + x | x <- [0..7], runSingleInput (current * 8 + x) == (drop dropN program)]

runSingleInput x = (uncurry runProgram) $ buildComputer ((x, 0, 0), program)

-- this can/should be refactored to work recursively ("programming by wishful thinking") but this is how i actually got the answer
wins = concatMap (workBackwards 0) $ concatMap (workBackwards 1) $ concatMap (workBackwards 2) $ concatMap (workBackwards 3) $ concatMap (workBackwards 4) $ concatMap (workBackwards 5) $ concatMap (workBackwards 6) $ concatMap (workBackwards 7) $ concatMap (workBackwards 8) $ concatMap (workBackwards 9) $ concatMap (workBackwards 10) $ concatMap (workBackwards 11) $ concatMap (workBackwards 12) seeds


seeds = take 5 $ [x | x <- [0..], runSingleInput x == [5,3,0]]