bf.hs

{-# LANGUAGE MultiWayIf, LambdaCase #-}

import qualified Data.Vector.Unboxed.Mutable as VM
import qualified Data.Vector.Unboxed as V
import Control.Monad
import Control.Monad.ST
import Control.Applicative
import Data.Functor
import Data.List
import Data.Char
import System.IO
import Text.Printf

shrink :: String -> String
shrink (s:ss)
  | s `elem` "<>+-.,[]@" = s : shrink ss
  | otherwise = shrink ss
shrink [] = []

compile :: String -> (V.Vector Char, V.Vector Int)
compile source' =
  let
    source = shrink source'
    len = length source

    codeMem = V.fromList source
    jumpMem = runST $ do
      jumpMemBuild <- VM.replicate len (-1)
      let
        build cursor stack ('[':srcs) = build (cursor + 1) (cursor : stack) srcs
        build cursor (top:stack) (']':srcs) = do
          VM.write jumpMemBuild cursor top
          VM.write jumpMemBuild top cursor
          build (cursor + 1) stack srcs
        build cursor stack (_:srcs) = build (cursor + 1) stack srcs
        build _ _ _ = pure ()
      build 0 [] source

      V.unsafeFreeze jumpMemBuild

  in
    (codeMem, jumpMem)

desugar = go 0 where
  go count ('#':_) = ""
  go count (c:cs)
    | isDigit c = go (count * 10 + ord c - 48) cs
    | c `elem` "<>+-,.[]@" = (if count >= 1 then take count (repeat c) else [c]) ++ go 0 cs
    | otherwise = go count cs
  go _ cs = cs

getSource = go [] where
  go acc = do
    eof <- isEOF
    if eof then
      return $ join $ reverse acc
    else do
      line <- desugar <$> getLine
      go (line : acc)

modifyVM vec f i = do
  v <- VM.unsafeRead vec i
  VM.unsafeWrite vec i (f v)

main = do
  memN <- read <$> getLine
  input <- getLine
  source <- getSource

  let (codeMem, jumpMem) = compile source

  putStrLn $ filter (/= '@') source
  --putStrLn (show codeMem)
  --putStrLn (show jumpMem)

  dataMem <- VM.replicate memN 0
  let
    run ip ptr input
      | ptr < 0 = putStrLn ("\nERROR: ptr = " ++ show ptr ++ " at " ++ take ip source ++ "!@!" ++ drop ip source) >> pure ptr
      | ptr >= memN = putStrLn ("\nERROR: ptr = " ++ show ptr ++ " at " ++ take ip source ++ "!@!" ++ drop ip source) >> pure ptr
      | ip >= V.length codeMem = pure ptr
      | otherwise = do
        case codeMem V.! ip of
          '>' -> run (ip + 1) (ptr + 1) input
          '<' -> run (ip + 1) (ptr - 1) input
          '+' -> do
            flip (modifyVM dataMem) ptr $ \case
              255 -> 0
              d -> d + 1
            run (ip + 1) ptr input
          '-' -> do
            flip (modifyVM dataMem) ptr $ \case
              0 -> 255
              d -> d - 1
            run (ip + 1) ptr input
          '.' -> do
            d <- VM.read dataMem ptr
            putChar $ chr d
            run (ip + 1) ptr input
          ',' -> do
            VM.write dataMem ptr (ord (head input))
            run (ip + 1) ptr (tail input)
          '[' -> do
            d <- VM.read dataMem ptr
            if d == 0 then
              run (jumpMem V.! ip) ptr input
            else
              run (ip + 1) ptr input
          ']' -> do
            d <- VM.read dataMem ptr
            if d == 0 then
              run (ip + 1) ptr input
            else
              run (jumpMem V.! ip) ptr input
          '@' -> do
            putStrLn ""
            putStrLn $ "ptr: " ++ show ptr
            forM_ [0..memN-1] $ \i -> do
              d <- VM.read dataMem i
              if i == ptr then
                printf "\x1b[1;35m%02x\x1b[m " d
              else
                printf "%02x " d
            putStrLn ""
            run (ip + 1) ptr input

  run 0 0 input