module Control.Monad.BrainFuck where

import qualified Control.Monad as M
import qualified Data.Char as Char

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

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
	return ret = BrainFuck $ \loc -> ([], loc, ret)
	a >>= b = BrainFuck $ \start -> let
		(left, mid, val) = func a start
		(right, end, ret) = func (b val) mid
		in (left ++ right, end, ret)

next, prev, incr, decr, output, input, open, close :: 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 ']'

-- | 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 data
-- pointer points to 0.
loopUnless0 :: BrainFuck a -> BrainFuck a
loopUnless0 a = do
	open
	r <- a
	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 ()

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

add, sub :: Int -> BrainFuck ()
-- adds an Int to the byte at the data pointer
add = multi incr
-- subtracts an Int from the byte at the data pointer
sub = multi decr

-- | 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 :: Int -> 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

-- | 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 (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 :: Int -> (DataPointer -> BrainFuck ()) -> BrainFuck ()
loopFrom n a = alloc $ do
	i <- addr
	add n
	loopUnless0 $ do
		a i
		setAddr i

-- | Runs an action in an infinite loop. The action should avoid
-- touching the current memory cell.
forever :: BrainFuck a -> BrainFuck ()
forever a = loopFrom 1 $ \_ -> do
	_ <- a
	return ()

-- | Runs the action unless the data pointer points to 0.
unless0 :: BrainFuck () -> BrainFuck ()
unless0 a = do
	start <- addr
	loopUnless0 $ do
		a
		next
		zero
	prev

-- | Adds the current and next data cells. The next cell is zeroed
-- and the sum is left in the current cell.
sum :: BrainFuck ()
sum = 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 :: Int -> BrainFuck ()
mult y = do
	x <- addr
	alloc $ do
		c1 <- addr

		-- Copy x to c1, and zero x.
		setAddr x
		loopUnless0 $ do
			setAddr c1
			incr
			setAddr x
			decr
	
		setAddr c1
		loopUnless0 $ do
			loopFrom y $ \c2 -> do
				decr
				setAddr x
				incr
			setAddr c1
			decr

-- | 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 c
			output
			go False cs c
	  where
		delta = c - n

-- | Prints out the alphabet, repeatedly.
demo :: String
demo = brainfuck $ forever $ do
	alloc $ do
		c <- addr
		set start
		
		loopFrom numchars $ \i -> do
			decr
			setAddr c
			output
			incr
		withChar ' ' output
  where
	start = Char.ord 'a'
	end = Char.ord 'z'
	numchars = end - start + 1

-- | Copy input to output.
cat :: String
cat = brainfuck $ forever $ 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!"

-- | Simple optimiser for brainfuck code.
optimize :: String -> String
optimize s = 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
	go orig [] ('[':'-':']':l) = go orig [] l -- [-] at start is noop
	go orig [] ('>':'[':'-':']':l) = go orig [] l -- >[-] at start is noop
	go orig r (c:l) = go orig (c:r) l