Brainfuck interpreter

brainfuck.hs

{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE TemplateHaskell            #-}

module Brainfuck where

import Control.Lens
import Control.Monad.Except
import Control.Monad.Reader
import Control.Monad.State
import Control.Monad.Trans.Either
import Data.Char
import System.IO
import Text.Printf

-- -----------------------------------------------------------------------------
-- Simple infinite zipper and mini API

infixr 5 :~
data Stream a = a :~ Stream a
  deriving Functor

instance Applicative Stream where
    pure x = x :~ pure x
    (f :~ fs) <*> (x :~ xs) = f x :~ (fs <*> xs)

data Zipper a = Zipper (Stream a) a (Stream a)
  deriving Functor

instance Show a => Show (Zipper a) where
    show (Zipper xs y zs) =
        let x3 :~ x2 :~ x1 :~ _ = xs
            z1 :~ z2 :~ z3 :~ _ = zs
        in printf "[...%s,%s,%s,>%s<,%s,%s,%s...]"
                  (show x1) (show x2) (show x3) (show y)
                  (show z1) (show z2) (show z3)


instance Applicative Zipper where
    pure x = Zipper (pure x) x (pure x)
    Zipper fs g hs <*> Zipper xs y zs = Zipper (fs <*> xs) (g y) (hs <*> zs)

-- | Get the focus of a zipper.
focus :: Zipper a -> a
focus (Zipper _ x _) = x

-- | Modify the focus of a zipper.
modFocus :: (a -> a) -> Zipper a -> Zipper a
modFocus f (Zipper xs y zs) = Zipper xs (f y) zs

-- | Set the focus of a zipper.
setFocus :: a -> Zipper a -> Zipper a
setFocus = modFocus . const

-- | Move left.
moveLeft :: Zipper a -> Zipper a
moveLeft (Zipper (x :~ xs) y zs) = Zipper xs x (y :~ zs)

-- | Move right.
moveRight :: Zipper a -> Zipper a
moveRight (Zipper xs y (z :~ zs)) = Zipper (y :~ xs) z zs

-- -----------------------------------------------------------------------------
-- Brainfuck datatypes and API

data Command
    = GoLeft
    | GoRight
    | Incr
    | Decr
    | Output
    | Input
    | LeftBracket
    | RightBracket
    deriving Show

readCommand :: Char -> Maybe Command
readCommand '<' = Just GoLeft
readCommand '>' = Just GoRight
readCommand '+' = Just Incr
readCommand '-' = Just Decr
readCommand '.' = Just Output
readCommand ',' = Just Input
readCommand '[' = Just LeftBracket
readCommand ']' = Just RightBracket
readCommand _   = Nothing

type Program = [Command]

-- The state of a Brainfuck interpreter:
--
-- - The unconsumed input
--
-- - The tape (an infinite array of integers, with a focus)
--
-- - A stack of programs, the head of which is resumed to upon looping back via
--   a ']' command.
data BFState = BFState
    { _bfInput :: Program
    , _bfTape  :: Zipper Int
    , _bfStack :: [Program]
    } deriving Show
makeLenses ''BFState

data BFError
    = MissingRightBracket
    | UnexpectedRightBracket
    deriving Show

-- Brainfuck interpreter
newtype Brainfuck a = Brainfuck { unBrainfuck :: StateT BFState (EitherT BFError IO) a }
  deriving (Functor, Applicative, Monad, MonadIO, MonadState BFState, MonadError BFError)

-- | Run the Brainfuck interpreter.
runBrainfuck :: String -> IO (Either BFError Int)
runBrainfuck = fmap (fmap fst) . debugBrainfuck

-- | Run the Brainfuck interpreter and return the resulting state.
debugBrainfuck :: String -> IO (Either BFError (Int, BFState))
debugBrainfuck prog =
    let initialState = BFState (parseProgram prog) (pure 0) []
    in runEitherT (runStateT (unBrainfuck interpreter) initialState)

-- | Get the value at the currently focused address.
currentFocus :: Brainfuck Int
currentFocus = uses bfTape focus

-- | Pop the top command of the input and return it.
popCommand :: Brainfuck (Maybe Command)
popCommand = use bfInput >>= \case
    []     -> return Nothing
    (c:cs) -> bfInput .= cs >> return (Just c)

-- | Push the current input onto the program stack.
pushInput :: Brainfuck ()
pushInput = use bfInput >>= \input -> bfStack %= (input :)

-- | Drop everything up to and including an unmatched right bracket.
dropUntilRB :: Brainfuck ()
dropUntilRB = go 0
  where
    go :: Int -> Brainfuck ()
    go n = popCommand >>= \case
        Nothing -> throwError MissingRightBracket
        Just RightBracket ->
            case n of
                0 -> return ()
                _ -> go (n-1)
        Just LeftBracket -> go (n+1)
        Just _ -> go n

-- -----------------------------------------------------------------------------
-- Brainfuck parsing and interpreting

-- | Parse a Brainfuck program from a string. All strings are legal, because
-- non-command characters are ignored.
parseProgram :: String -> Program
parseProgram = foldr (maybe id (:) . readCommand) []

-- | Interpret a Brainfuck program in IO. Return the focused element at the end
-- of the computation.
interpreter :: Brainfuck Int
interpreter = popCommand >>= \case
    -- No more commands; return the focused element.
    Nothing -> currentFocus

    -- Easy cases that hardly need any explanation - just do something with
    -- the current character, and then interpret the rest.

    Just GoRight -> bfTape %= moveRight                             >> interpreter
    Just GoLeft  -> bfTape %= moveLeft                              >> interpreter
    Just Incr    -> bfTape %= modFocus (+1)                         >> interpreter
    Just Decr    -> bfTape %= modFocus (subtract 1)                 >> interpreter
    Just Output  -> currentFocus >>= liftIO . putChar . chr         >> interpreter
    Just Input   -> liftIO getChar >>= (bfTape %=) . setFocus . ord >> interpreter

    -- And here, it gets slightly more interesting.

    Just LeftBracket -> currentFocus >>= \case
        -- Focus is zero - jump to the command after the matching ']' and
        -- resume execution.
        0 -> dropUntilRB >> interpreter

        -- Focus is non-zero - push the current program onto the stack
        -- (to be popped and resumed from if ']' loops back), and resume
        -- execution.
        _ -> pushInput >> interpreter

    Just RightBracket -> do
        use bfStack >>= \case
            -- If we encounter a ']', we should have at some point in the past
            -- pushed a program to resume from onto the stack.
            [] -> throwError UnexpectedRightBracket
            (p:ps) -> currentFocus >>= \case
                -- Focus is zero - toss the last saved program off the stack
                -- (we will never loop back to it), and resume execution.
                0 -> bfStack .= ps >> interpreter
                -- Focus is non-zero - resume execution from the saved program
                -- on top of the stack.
                _ -> bfInput .= p >> interpreter

-- -----------------------------------------------------------------------------
-- Sample programs

-- Echoes each input character.
echo :: String
echo = "+[,.]"

-- Echoes N input characters.
echoN :: Int -> String
echoN n = replicate n '+' ++ "[->,.<]"

echoUppercase :: String
echoUppercase = "+[,--------------------------------.]"

-- "Hello, World!" stolen from Wikipedia.
helloWorld :: String
helloWorld = "\
\   [ This program prints 'Hello World!' and a newline to the screen, its\
\     length is 106 active command characters [it is not the shortest.]\

\     This loop is a 'comment loop', it's a simple way of adding a comment\
\     to a BF program such that you don't have to worry about any command\
\     characters. Any '.', ',', '+', '-', '<' and '>' characters are simply\
\     ignored, the '[' and ']' characters just have to be balanced.\
\   ]\
\   +++++ +++               Set Cell #0 to 8\
\   [\
\       >++++               Add 4 to Cell #1; this will always set Cell #1 to 4\
\       [                   as the cell will be cleared by the loop\
\           >++             Add 2 to Cell #2\
\           >+++            Add 3 to Cell #3\
\           >+++            Add 3 to Cell #4\
\           >+              Add 1 to Cell #5\
\           <<<<-           Decrement the loop counter in Cell #1\
\       ]                   Loop till Cell #1 is zero; number of iterations is 4\
\       >+                  Add 1 to Cell #2\
\       >+                  Add 1 to Cell #3\
\       >-                  Subtract 1 from Cell #4\
\       >>+                 Add 1 to Cell #6\
\       [<]                 Move back to the first zero cell you find; this will\
\                           be Cell #1 which was cleared by the previous loop\
\       <-                  Decrement the loop Counter in Cell #0\
\   ]                       Loop till Cell #0 is zero; number of iterations is 8\

\   The result of this is:\
\   Cell No :   0   1   2   3   4   5   6\
\   Contents:   0   0  72 104  88  32   8\
\   Pointer :   ^\

\   >>.                     Cell #2 has value 72 which is 'H'\
\   >---.                   Subtract 3 from Cell #3 to get 101 which is 'e'\
\   +++++++..+++.           Likewise for 'llo' from Cell #3\
\   >>.                     Cell #5 is 32 for the space\
\   <-.                     Subtract 1 from Cell #4 for 87 to give a 'W'\
\   <.                      Cell #3 was set to 'o' from the end of 'Hello'\
\   +++.------.--------.    Cell #3 for 'rl' and 'd'\
\   >>+.                    Add 1 to Cell #5 gives us an exclamation point\
\   >++.                    And finally a newline from Cell #6"