bf.hs

module Brainfuck
    ( executeString
    ) where

import Data.Map (Map)
import qualified Data.Map as Map
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import Control.Monad.Trans.State
import Control.Monad.Trans.Maybe
import Data.Functor.Identity
import Control.Monad
import Control.Applicative
import Data.Char

data Parse = Parse
           { index :: Int
           , startStack :: [Int]
           , jumpTable :: Map Int Int }

parse :: String -> Map Int Int
parse input = jumpTable $ execState (s input) $ Parse 0 [] Map.empty where
    s (c : xs) = do 
        st <- get
        let i = index st
            j = head $ startStack st
        case c of '[' -> put $ st { startStack = index st : startStack st }
                  ']' -> put $ st { startStack = tail $ startStack st
                                  , jumpTable = Map.insert i j $ Map.insert j i $ jumpTable st }
                  _   -> return () 
        modify $ \st -> st { index = index st + 1 }
        s xs
    s [] = return ()

data Exec = Exec
          { pc :: Int
          , pointer :: Int
          , output :: String
          , input :: String
          , memory :: Map Int Char }

-- | Interprets the Brainfuck source code from the first argument, while
-- supplying it with input from the second. May fail on insufficient input.
executeString :: String -> String -> Maybe String
executeString source inp = reverse <$> output <$> runIdentity (runMaybeT $ execStateT s (Exec 0 0 "" inp $ Map.singleton 0 $ chr 0))
    where jt = parse source
          code = Vector.fromList source 
          end = length source
          inc c = chr $ (ord c + 1) `mod` 256
          dec '\0' = chr 255
          dec c = chr $ ord c - 1
          init st = st { memory = Map.alter (<|> Just (chr 0)) (pointer st) $ memory st }
          s :: StateT Exec (MaybeT Identity) ()
          s = do st <- get
                 let c = code Vector.! pc st
                     val = memory st Map.! pointer st
                 guard (c /= ',' || not (null $ input st))
                 put $ case c of
                     '>' -> init $ st { pointer = pointer st + 1 }
                     '<' -> init $ st { pointer = pointer st - 1 }
                     '+' -> st { memory = Map.adjust inc (pointer st) $ memory st }
                     '-' -> st { memory = Map.adjust dec (pointer st) $ memory st }
                     ',' -> st { memory = Map.insert (pointer st) (head $ input st) $ memory st
                               , input = tail $ input st }
                     '.' -> st { output = val : output st }
                     '[' -> if val == chr 0 then st { pc = jt Map.! pc st } else st
                     ']' -> if val /= chr 0 then st { pc = jt Map.! pc st } else st
                 st <- get
                 put $ st { pc = pc st + 1 }
                 unless (pc st + 1 == end) s