{-# LANGUAGE LambdaCase #-}
module Asm where

import qualified Data.Vector.Unboxed as U
import Data.List
import Data.Word


data Ins
    = NOP Int
    | ACC Int
    | JMP Int
  deriving (Show)

parseIns :: String -> Ins
parseIns line
  | Just idx <- findIndex (== ' ') line
  = case parseArg . tail <$> splitAt idx line of
      ("nop", arg) -> NOP arg
      ("acc", arg) -> ACC arg
      ("jmp", arg) -> JMP arg
      _ -> error ("Invalid instruction: '" ++ line ++ "'")
  where
    parseArg :: String -> Int
    parseArg ('+':s) = read s
    parseArg s = read s
parseIns line = error ("Invalid instruction (no space): '" ++ line ++ "'")

data State
    = State { sAcc :: Int
            , sIP  :: Int }
  deriving (Show)

initState :: State
initState = State 0 0

evalIns :: Ins -> State -> State
evalIns = \case
    NOP _ -> jmprel 1
    ACC i -> jmprel 1 . (\s -> s { sAcc = sAcc s + i })
    JMP i -> jmprel i
  where
    jmprel :: Int -> State -> State
    jmprel off state = state { sIP = sIP state + off }

data Program = Program (U.Vector (Word8, Int)) (Int -> Maybe Ins)

programLength :: Program -> Int
programLength (Program vec _) = U.length vec

programGet :: Program -> Int -> Ins
programGet (Program _ overrider) i | Just ins <- overrider i = ins
programGet (Program vec _) i =
    case vec U.! i of
      (0, a) -> NOP a
      (1, a) -> ACC a
      (2, a) -> JMP a
      _ -> error "Invalid program element"

programFrom :: [Ins] -> Program
programFrom l =
    Program (U.fromList [case ins of
                           NOP a -> (0, a)
                           ACC a -> (1, a)
                           JMP a -> (2, a)
                        | ins <- l])
            (const Nothing)

withOverrider :: Program -> (Int -> Maybe Ins) -> Program
withOverrider (Program vec _) overrider = Program vec overrider

evalUntil :: (s -> State -> (s, Bool)) -> Program -> s -> State -> State
evalUntil predicate prog aux state =
    let (aux', cont) = predicate aux state
        state' = evalIns (programGet prog (sIP state)) state
    in if cont
           then state
           else evalUntil predicate prog aux' state'