trivial optimizing brainfuck interpreter

C.hs

{-# LANGUAGE OverloadedStrings #-}
import Data.Array.IO
import Data.Array.Base
import Control.Applicative
import Control.Monad
import Data.Attoparsec.Char8
import qualified Data.ByteString.Char8 as B
import System.Environment
import Debug.Trace

data Primitive =
    Move      {-# UNPACK #-} !Int
  | Add       {-# UNPACK #-} !Int
  | Output
  | Input
  | Loop [Primitive]
  -- optimizations
  | Zero
  | MoveSet   {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  | Assign    {-# UNPACK #-} !Int {-# UNPACK #-} !Int
  deriving Show

bf :: Parser [Primitive]
bf = many instruction

instruction :: Parser Primitive
instruction = primitive <|> loop

primitive :: Parser Primitive
primitive =
      Move . negate . B.length <$> takeWhile1 (== '<')
  <|> Move          . B.length <$> takeWhile1 (== '>')
  <|> Add           . B.length <$> takeWhile1 (== '+')
  <|> Add  . negate . B.length <$> takeWhile1 (== '-')
  <|> char8 '.' *> pure Output
  <|> char8 ',' *> pure Input

loop :: Parser Primitive
loop = char8 '[' *> (Loop <$> bf) <* char8 ']'

optimize :: [Primitive] -> [Primitive]
optimize = second_pass . first_pass
 where
  first_pass [] = []
  first_pass (Loop [Add (-1)]:as) = Zero : first_pass as
  first_pass (a@(Loop [Add (-1), Move m1, Add n2, Move m2]):as)
    | m1 == -m2 = Assign m1 n2 : optimize as
    | otherwise = a : optimize as
  first_pass (Loop a:as) = Loop (optimize a) : first_pass as
  first_pass (a:as) = a : first_pass as

  second_pass [] = []
  second_pass (Move n : Add m : as) = MoveSet n m : second_pass as
  second_pass (a:as) = a : second_pass as

type Tape = IOUArray Int Int
type Ptr = Int

eval :: [Primitive] -> IO ()
eval ps = do
  tape <- newArray (1,30000) 0 :: IO Tape
  --print $ optimize ps
  void $ go (optimize ps) tape 1
 where
  go :: [Primitive] -> Tape -> Ptr -> IO Ptr

  go [] _ sp = return sp

  go (Move n : ps) tape sp = go ps tape (sp + n)

  go (Add n : ps) tape sp = do
    v <- unsafeRead tape sp
    unsafeWrite tape sp (v + n)
    go ps tape sp

  go (Zero : ps) tape sp = do
    unsafeWrite tape sp 0
    go ps tape sp

  go (MoveSet n m : ps) tape sp0 = do
    let sp = sp0 + n
    v <- unsafeRead tape sp
    unsafeWrite tape sp (v + m)
    go ps tape sp

  go (Output:ps) tape sp = do
    putChar . toEnum =<< unsafeRead tape sp
    go ps tape sp

  go (Input:ps) tape sp = do
    getChar >>= unsafeWrite tape sp . fromEnum
    go ps tape sp

  go (Loop p:ps) tape sp0 = lgo sp0 >>= go ps tape
   where
    lgo sp = do
      v <- unsafeRead tape sp
      if v /= 0
        then go p tape sp >>= lgo
        else return sp

  go (Assign to x : ps) tape sp = do
    v0 <- unsafeRead tape sp
    v1 <- unsafeRead tape (sp + to)
    unsafeWrite tape (sp + to) (v1 + v0*x)
    unsafeWrite tape sp 0
    go ps tape sp

main :: IO ()
main = do
  program <- B.readFile =<< head <$> getArgs
  let sanitized = B.filter (`B.elem` "<>+-.,[]") program
  either error eval (parseOnly bf sanitized)