{-# 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