{-# LANGUAGE DeriveFunctor #-}

module Control.Monad.BrainFuck where

import qualified Control.Monad as M
import Control.Monad (ap)
import qualified Data.Char as Char
import Data.Word (Word8)
import System.Process
import System.IO
import System.Exit
import System.Directory

newtype BrainFuck a = BrainFuck (DataPointer -> ([Char], DataPointer, a))
        deriving (Functor)

type DataPointer = Integer

-- | Retrieve the inner function
func :: BrainFuck f -> DataPointer -> ([Char], DataPointer, f)
func (BrainFuck f) = f

-- | Evaluate the monad and get a brainfuck program
brainfuck :: BrainFuck f -> String
brainfuck (BrainFuck f) = bytes where (bytes, _, _) = f 0

instance Monad BrainFuck where
        a >>= b = BrainFuck $ \start -> let
                (left, mid, val) = func a start
                (right, end, ret) = func (b val) mid
                in (left ++ right, end, ret)

instance Applicative BrainFuck where
        pure v = BrainFuck $ \loc -> ([], loc, v)
        (<*>) = ap

next, prev, incr, decr, output, input, open, close, debug :: BrainFuck ()
-- | move data pointer right
next = opcode' succ '>'
-- | move data pointer left
prev = opcode' pred '<'
-- | increment data
incr = opcode '+'
-- | decrement data
decr = opcode '-'
-- | output byte at data pointer
output = opcode '.'
-- | input byte, storing at data pointer
input = opcode ','
-- | if byte at data pointer is zero, jump to opcode after close
open = opcode '['
-- | if byte at data pointer is nonzero, jump to optoce after matching open
close = opcode ']'

-- | ignored if brainfuck machine does not support debugging
debug = opcode '#'

-- | Adds an arbitrary character to the program.
-- Should not be used directly.
opcode :: Char -> BrainFuck ()
opcode = opcode' id

-- | Adds an arbitrary character to the program,
-- and updates the data pointer.
-- Should not be used directly.
opcode' :: (DataPointer -> DataPointer) -> Char -> BrainFuck ()
opcode' f x =  BrainFuck $ \loc -> ([x], f loc, ())

-- | The loop is only run if the data pointer doesn't point to 0.
--
-- On entry, the loop body is run, and then it loops, until the address
-- the data pointer originally pointed at has a value of 0.
--
-- Any change that the loop body makes to the address of the data pointer
-- is reset each time through the loop (and at the end). This is necessary
-- to keep the BrainFuck monad's DataPointer consistent no matter what
-- happens when running the loop.
loopUnless0 :: BrainFuck a -> BrainFuck a
loopUnless0 a = do
        start <- addr
        open
        r <- a
        setAddr start
        close
        return r

-- | Gets the current address of the data pointer.
addr :: BrainFuck DataPointer
addr = BrainFuck $ \loc -> ([], loc, loc)

-- | Moves the data pointer to a specific address.
setAddr :: Integer -> BrainFuck ()
setAddr n = do
        a <- addr
        if a > n
                then prev >> setAddr n
                else if a < n
                        then next >> setAddr n
                        else return ()

-- | Runs an action with the data pointer temporarily set to an address,
-- then restores the data pointer.
withAddr :: Integer -> BrainFuck a -> BrainFuck a
withAddr n a = do
        old <- addr
        setAddr n
        r <- a
        setAddr old
        return r

-- | Run an action multiple times.
multi :: BrainFuck () -> Int -> BrainFuck ()
multi c n = do
        _ <- sequence (replicate n c)
        return ()

add, sub :: Word8 -> BrainFuck ()
-- adds an byte to the byte at the data pointer (wraps on overflow)
add = multi incr . fromIntegral
-- subtracts an byte from the byte at the data pointer (wraps on underflow)
sub = multi decr . fromIntegral

-- | Zeros the current data cell.
zero :: BrainFuck ()
zero = loopUnless0 decr

-- | Changes the current data cell to contain a specific value.
-- (It can start at any value).
set :: Word8 -> BrainFuck ()
set n = do
        zero
        add n

-- | For higher-level programming in brainfuck, it's useful to have a way
-- to run a function, while allocating a memory cell, which is initialized
-- to contain 0.
--
-- This and many of the functions below assume that
-- cells to the left are in use, while cells to the right
-- are unused and may contain any data. Higher-level functions should
-- generally avoid changing the current cell, and should instead alloc
-- a new one to use.
alloc :: BrainFuck a -> BrainFuck a
alloc a = do
        next
        zero
        cell <- addr
        r <- a
        setAddr cell
        prev
        return r

alloc' :: (DataPointer -> BrainFuck a) -> BrainFuck a
alloc' a = alloc $ addr >>= a

-- | Allocates a new memory cell, populates it with a Char, and runs
-- the action.
withChar :: Char.Char -> BrainFuck a -> BrainFuck a
withChar c a = alloc $ do
        set $ fromIntegral $ Char.ord c
        a

-- | Allocates a cell and uses it as the loop counter, starting from 
-- the provided value. The action will continue running in a loop until
-- it decrements the counter to 0.
loopFrom :: Word8 -> (DataPointer -> BrainFuck ()) -> BrainFuck ()
loopFrom n a =
        alloc' $ \i -> do
                add n
                loopUnless0 $
                        a i

-- | Runs an action in an infinite loop. The action can modify its
-- current memory cell, or allocate and use new ones, and will not
-- exit the loop.
forever :: BrainFuck a -> BrainFuck ()
forever a = loopFrom 1 $ \_i ->
        alloc' $ \_scratch -> do
                _ <- a
                return ()

-- | Copies the value of a cell.
copy :: DataPointer -> DataPointer -> BrainFuck ()
copy src dest = do
        withAddr dest zero
        alloc' $ \tmp -> do
                -- copy src to new and tmp, but this clobbers src
                setAddr src
                loopUnless0 $ do
                        decr
                        setAddr dest
                        incr
                        setAddr tmp
                        incr
                -- put src back how it was
                setAddr tmp
                loopUnless0 $ do
                        decr
                        setAddr src
                        incr

-- | Runs the action with a new cell that is logically NOT
-- the value of the passed pointer.
withNot :: DataPointer -> BrainFuck () -> BrainFuck ()
withNot p a = alloc' $ \tmp1 -> do
        copy p tmp1
        alloc' $ \tmp2 -> do
                incr
                setAddr tmp1
                loopUnless0 $ do
                        zero
                        setAddr tmp2
                        decr
                setAddr tmp2
                a

-- | Runs the action unless the data pointer points to 0.
unless0 :: DataPointer -> BrainFuck () -> BrainFuck ()
unless0 p a = do
        start <- addr
        alloc' $ \tmp -> do
                copy p tmp
                loopUnless0 $ do
                        setAddr start
                        a
                        setAddr tmp
                        zero -- don't loop

-- | Runs the action when the data pointer points to 0.
when0 :: DataPointer -> BrainFuck () -> BrainFuck ()
when0 p a = withNot p $
        addr >>= flip unless0 a

-- | Monadic if; the first action is run if the data pointer points to 0,
-- else the second action is run.
if0 :: DataPointer -> (BrainFuck (),  BrainFuck ()) -> BrainFuck ()
if0 p (y, n) = do
        when0 p y
        unless0 p n

-- | Adds the current and next data cells. The next cell is zeroed
-- and the sum is left in the current cell.
sumNext :: BrainFuck ()
sumNext = do
        next
        loopUnless0 $ do
                prev
                incr
                next
                decr
        prev

-- | Multiplies the current data cell by some value. Uses and zeros some
-- of the following cells.
mult :: Word8 -> BrainFuck ()
mult y = do
        x <- addr
        alloc' $ \c1 -> do
                -- Copy x to c1, and zero x.
                setAddr x
                loopUnless0 $ do
                        decr
                        setAddr c1
                        incr

                setAddr c1
                loopUnless0 $ do
                        decr
                        loopFrom y $ \_c2 -> do
                                decr
                                setAddr x
                                incr

-- | Displays a string. Tries to generate a fairly small brainfuck program,
-- using a few encoding tricks. The current cell is modified, and not
-- cleaned up at the end, so run using alloc if necessary.
display :: String -> BrainFuck ()
display s = start >>= go True (map Char.ord s)
  where
        -- Get to letter 104 ('a' is 97) quickly by multiplication.
        start = do
                zero
                _x <- addr
                add 13
                mult 8
                return (13 * 8)

        go _ [] _ = return ()
        go started (c:cs) n
                | not started && delta > 13 = start >>= go True (c:cs)
                | otherwise = do
                        if abs delta < c
                                then multi (if delta > 0 then incr else decr) (abs delta)
                                else set (fromIntegral c)
                        output
                        go False cs c
          where
                delta = c - n

-- | Prints out the alphabet, repeatedly, with some whitespace fun
-- to make it more interesting.
demo :: String
demo = brainfuck demo'

demo' :: BrainFuck ()
demo' = alloc' $ \tilt -> do
        settilt
        alloc' $ \tc -> do
                decr -- set to 255
                alloc' $ \wc -> do
                        copy tilt wc
                        forever $ alloc' $ \c -> do
                                set (fromIntegral start)
                                loopFrom (fromIntegral numchars) $ \_i -> do
                                        decr

                                        withAddr wc decr
                                        when0 wc $ do
                                                loopFrom 3 $ \_ -> do
                                                        decr
                                                        withChar ' ' output
                                                copy tilt wc
                                                withAddr tc decr
                                                when0 tc $ do
                                                        withAddr tc decr
                                                        withAddr tilt decr
                                                        when0 tilt $ do
                                                                withAddr tilt
                                                                        settilt

                                        setAddr c
                                        output
                                        incr
  where
        start = Char.ord 'a'
        end = Char.ord 'z'
        numchars = end - start + 1

        settilt = do
                set 4
                mult 5

-- | Copy input to output.
cat :: String
cat = brainfuck $ forever $ alloc $ input >> output

-- | Simple hello world.
helloworld :: String
helloworld = brainfuck $
        M.forM_ "hello, world!" $
                flip withChar output

-- | Optimized to use less space.
helloworld' :: String
helloworld' = optimize $ brainfuck $ display "hello, world!"

{-
euler1 :: BrainFuck
euler1 = alloc' $ \result -> alloc' $ \scratch -> do
	let accum i = do
		setAddr scratch
		dup i
		setAddr result
		sumNext
		
	let checkmod i n = alloc $ do
		dup i
		mod n
		when0 $
			accum i

	loopFrom 1000 $ \i -> do
		decr
		alloc $ do
			checkmod i 5
			print0	
			unless0 $
				checkmod 1 3

	setAddr result
	decimalOutput
-}

-- | Simple optimiser for brainfuck code.
optimize :: String -> String
optimize s = trimend $ go s [] s
  where
        go orig r [] =
                let new = reverse r in
                if new /= orig
                        then go new [] new
                        else new
        go orig r ('>':'<':l) = go orig r l -- <> is a noop
        go orig r ('<':'>':l) = go orig r l -- >< is a noop
        go orig r ('+':'-':l) = go orig r l -- +- is a noop
        go orig r ('-':'+':l) = go orig r l -- -+ is a noop
        go orig r ('[':']':l) = go orig r l -- [] is a noop
        -- [-] before , is a noop, because the read overwrites the value
        go orig r ('[':'-':']':',':l) = go orig r (',':l)
        -- so is + or - before ,
        go orig r ('+':',':l) = go orig r (',':l)
        go orig r ('-':',':l) = go orig r (',':l)
        -- [-][-] is the same as [-]
        go orig r ('[':'-':']':'[':'-':']':l) = go orig r ('[':'-':']':l)
        -- [-] at start is noop, because memory starts empty
        go orig [] ('[':'-':']':l) = go orig [] l
        -- >[-] at start is generally a noop (for programs using alloc)
        go orig [] ('>':'[':'-':']':l) = go orig [] l
        go orig r (c:l) = go orig (c:r) l

        -- Any sequence of > and < at the end of a program is a noop.
        trimend = reverse . dropWhile (`elem` "><") . reverse

-- | Uses the MOVfuscator v1 to build a Linux executable from a BrainFuck
-- monad action. The executable is implemented entirely using the
-- MOV instruction. It may be a little slow, especially when loops are
-- involved.
--
-- Prereq: Clone https://github.com/xoreaxeaxeax/movfuscator, build the v1
-- brainfuck to MOV compiler (in the poc/) directory, and put it in PATH.
-- Also apt-get install nasm.
compile :: FilePath -> BrainFuck () -> IO ()
compile dest b = movfuscate b >>= nasm dest >>= ld dest

type ASM = String

movfuscate :: BrainFuck () -> IO ASM
movfuscate b = do
        (Just hin, Just hout, Nothing, h) <- createProcess $
                (proc "movfuscator" [])
                        { std_out = CreatePipe
                        , std_in = CreatePipe
                        }
        hPutStr hin $ optimize $ brainfuck b
        hClose hin
        asm <- hGetContents hout
        writeFile "asm" asm
        c <- waitForProcess h
        case c of
                ExitSuccess -> return ()
                _ -> error ("movfuscator exited nonzero")
        return asm

nasm :: FilePath -> ASM -> IO FilePath
nasm dest a = do
        -- nasm needs a real file, can't use stdin
        writeFile asmfile a
        callProcess "nasm" ["-felf", asmfile, "-o", objfile]
        removeFile asmfile
        return objfile
  where
        asmfile = dest ++ ".asm"
        objfile = dest ++ ".o"

ld :: FilePath -> FilePath -> IO ()
ld dest objfile = do
        callProcess "ld" ["-melf_i386", objfile, "-o", dest]
        removeFile objfile