Brainfuck Interpreter in Haskell

Brainfuck.hs

{-# LANGUAGE LambdaCase     #-}
{-# LANGUAGE NamedFieldPuns #-}

module Brainfuck
    ( evalString
    ) where

import           Control.Monad
import           Data.Char     (chr, ord)
import           Data.Maybe

data BFAst
    = ShiftRight
    | ShiftLeft
    | Increment
    | Decrement
    | Print
    | Read
    | Loop [BFAst]
    deriving (Show)

newtype Cell =
    Cell Int

data DataPointer = DataPointer
    { previous :: [Cell]
    , current  :: Cell
    , next     :: [Cell]
    }

data BFState = BFState
    { dataPointer :: DataPointer
    , input       :: String
    , output      :: Maybe String
    }

newtype BFStateM a =
    BFStateM (BFState -> (BFState, a))

instance Functor BFStateM where
    fmap = (<*>) . pure

instance Applicative BFStateM where
    pure = return
    (<*>) = ap

instance Monad BFStateM where
    return a = BFStateM $ \state -> (state, a)
    (BFStateM m) >>= f =
        BFStateM $ \state ->
            let (state', a) = m state
                (BFStateM m') = f a
            in m' state'

evalString :: String -> String -> Maybe String
evalString source input =
    let program = parseProgram source
    in snd $ runBF input (eval program)

runBF :: String -> BFStateM a -> (BFState, a)
runBF input (BFStateM m) = m $ initState input

eval :: [BFAst] -> BFStateM (Maybe String)
eval [] = stateReader output
eval (x:xs) = do
    out <- stateReader output
    if isNothing out
        then return Nothing
        else evalStmt x >> eval xs

evalStmt :: BFAst -> BFStateM ()
evalStmt =
    \case
        ShiftRight -> shiftRight
        ShiftLeft -> shiftLeft
        Increment -> increment
        Decrement -> decrement
        Print -> currentCell >>= printChar . cellToChar
        Read ->
            readChar >>= \case
                Nothing -> return ()
                Just a -> modifyCurrentCell $ const (charToCell a)
        Loop body -> evalLoop body

evalLoop :: [BFAst] -> BFStateM ()
evalLoop body = do
    c <- currentCell
    unless (isZero c) $ eval body >> evalLoop body

initDataPointer :: DataPointer
initDataPointer = DataPointer {previous = [], current = newCell, next = []}

initState :: String -> BFState
initState input =
    BFState
    {dataPointer = initDataPointer, input = input, output = Just ""}

stateReader :: (BFState -> a) -> BFStateM a
stateReader f = BFStateM (\state -> (state, f state))

modifyState :: (BFState -> BFState) -> BFStateM ()
modifyState f = BFStateM $ \state -> (f state, ())

modifyDataPointer :: (DataPointer -> DataPointer) -> BFStateM ()
modifyDataPointer f =
    modifyState $ \state -> state {dataPointer = f $ dataPointer state}

currentCell :: BFStateM Cell
currentCell = stateReader $ current . dataPointer

modifyCurrentCell :: (Cell -> Cell) -> BFStateM ()
modifyCurrentCell f =
    modifyDataPointer $ \dp -> dp {current = f $ current dp}

increment :: BFStateM ()
increment = modifyCurrentCell incrementCell

decrement :: BFStateM ()
decrement = modifyCurrentCell decrementCell

readChar :: BFStateM (Maybe Char)
readChar =
    BFStateM $ \state ->
        case input state of
            []     -> (state {output = Nothing}, Nothing)
            (x:xs) -> (state {input = xs}, Just x)

printChar :: Char -> BFStateM ()
printChar c =
    modifyState $ \case
        state@BFState {output = Just output} ->
            state {output = Just $ output ++ [c]}
        state -> state

newCell :: Cell
newCell = Cell 0

isZero :: Cell -> Bool
isZero (Cell 0) = True
isZero _        = False

incrementCell :: Cell -> Cell
incrementCell (Cell n)
    | n == 255 = Cell 0
    | otherwise = Cell (succ n)

decrementCell :: Cell -> Cell
decrementCell (Cell n)
    | n == 0 = Cell 255
    | otherwise = Cell (pred n)

charToCell :: Char -> Cell
charToCell c = Cell $ ord c

cellToChar :: Cell -> Char
cellToChar (Cell c) = chr c

shiftRight :: BFStateM ()
shiftRight = modifyDataPointer shiftRight'
  where
    shiftRight' :: DataPointer -> DataPointer
    shiftRight' DataPointer {previous, current, next = []} =
        DataPointer (current : previous) newCell []
    shiftRight' DataPointer {previous, current, next = (x:xs)} =
        DataPointer (current : previous) x xs

shiftLeft :: BFStateM ()
shiftLeft = modifyDataPointer shiftLeft'
  where
    shiftLeft' :: DataPointer -> DataPointer
    shiftLeft' DataPointer {previous = [], current, next} =
        DataPointer [] newCell (current : next)
    shiftLeft' DataPointer {previous = (x:xs), current, next} =
        DataPointer xs x (current : next)

parseProgram :: String -> [BFAst]
parseProgram [] = []
parseProgram ('[':xs) =
    let (loop, rest) = parseLoop xs
    in loop : parseProgram rest
parseProgram (x:xs) = parseChar x : parseProgram xs

parseLoop :: String -> (BFAst, String)
parseLoop [] = undefined
parseLoop (']':xs) = (Loop [], xs)
parseLoop ('[':xs) =
    let (loop, xs') = parseLoop xs
        (Loop ys, xs'') = parseLoop xs'
    in (Loop $ loop : ys, xs'')
parseLoop (x:xs) =
    let (Loop ys, xs') = parseLoop xs
    in (Loop $ parseChar x : ys, xs')

parseChar :: Char -> BFAst
parseChar =
    \case
        '>' -> ShiftRight
        '<' -> ShiftLeft
        '+' -> Increment
        '-' -> Decrement
        '.' -> Print
        ',' -> Read
        c -> error $ "invalid input: " ++ [c]