{-# LANGUAGE TupleSections #-}
module Hasklet4 where

import Control.Monad
import Data.Bifunctor
import Data.Bool


-- * Stack language syntax

-- | Stack programs.
type Prog = [Cmd]

-- | Commands for working with stacks of integers. 0 is treated as 'false',
--   all other values as 'true'. The examples below help illustrate the
--   behavior of some of the less obvious commands.
data Cmd
   = Push Int      -- ^ push an integer onto the stack
   | Drop          -- ^ drop the top element on the stack
   | Dig           -- ^ moves the ith element down to the top of the stack
   | Dup           -- ^ duplicate the top element on the stack
   | Neg           -- ^ negate the number on top of the stack
   | Add           -- ^ add the top two numbers on the stack
   | Mul           -- ^ multiply the top two numbers on the stack
   | LEq           -- ^ check whether the top element is less-than-or-equal to the second
   | If Prog Prog  -- ^ if the value on top is true, run the first program, else the second (consumes the test element)
   | While Prog    -- ^ loop as long as the top element is true (does not consume the test element)
  deriving (Eq,Show)


-- ** Example programs and expected results

-- Note that the expected results are written with the top element of the
-- stack on the *right*, which is the convention for stack-based languages.
-- However, since we're encoding stacks with Haskell lists, the resulting
-- Haskell values will be in the reverse order.

-- | Result: 4 5 5
p1 = [Push 4, Push 5, Push 6, Drop, Dup]

-- | Result: 10 11 13 14 12
p2 = [Push 10, Push 11, Push 12, Push 13, Push 14, Push 3, Dig]

-- | Result: 27 -5
p3 = [Push 3, Push 4, Push 5, Add, Mul, Push 5, Neg]

-- | Result: 0 1
p4 = [Push 3, Push 4, LEq, Push 4, Push 3, LEq]

-- | Result: 22
p5 = [Push 2, Push 3, Push 4, LEq, If [Push 10, Add] [Push 20, Add]]

-- | Compute the factorial of the top element of the stack.
fac = [
  Push 1,             -- acc = 1
  Push 2, Dig,        -- move i to top
  While [             -- while i /= 0
    Dup,              --   duplicate i
    Push 3, Dig,      --   move accumulator to top
    Mul,              --   acc * i
    Push 2, Dig,      --   move i back to top
    Push 1, Neg, Add  --   decrement i
    ],
  Drop                -- drop i to leave only acc
  ]

-- | Several programs that cause errors if run on an empty stack.
bads = [
  -- stack underflow errors
  [Neg],
  [Push 2, Add],
  [Push 3, Mul],
  [Push 4, Drop, Drop],
  [Dup],
  [Push 5, Neg, Dig],
  [If [] []],
  [While []],
  -- digging too deep and too greedily, or trying to dig up
  [Push 6, Push 2, Dig],
  [Push 7, Push 8, Push 2, Neg, Dig]
  ]


-- * Stack language semantics

-- ** Stack-tracking monad

-- | A stack of integers.
type Stack = [Int]

-- | A monad that maintains a stack as state and may also fail.
--   (A combination of the State and Maybe monads.)
data StackM a = SM (Stack -> Maybe (a, Stack))

-- | Run a computation with the given initial stack.
runWith :: Stack -> StackM a -> Maybe (a, Stack)
runWith s (SM f) = f s

instance Functor StackM where
  fmap = liftM

instance Applicative StackM where
  pure  = return
  (<*>) = ap

instance Monad StackM where
  return a = SM $ \s -> Just (a, s)
  (SM f) >>= g = SM $ \s -> f s >>= \(a, s') -> runWith s' $ g a

modify :: (Stack -> Stack) -> StackM ()
modify f = SM $ Just . ((),) . f

gets :: (Stack -> a) -> StackM a
gets f = SM $ \s -> Just (f s, s)

fail_ :: StackM a
fail_ = SM $ const Nothing

-- ** Primitive operations

-- | Push a value onto the stack.
push :: Int -> StackM ()
push i = modify (i:)

-- | Pop a value off the stack and return it.
pop :: StackM Int
pop = peek <* modify tail

popBool :: StackM Bool
popBool = (/= 0) <$> pop

-- | Peek at the value on top of the stack without popping it.
peek :: StackM Int
peek = gets safeHead >>= maybe fail_ return
    where
        safeHead [] = Nothing
        safeHead (x:xs) = Just x

peekBool :: StackM Bool
peekBool = (/= 0) <$> peek

fromBool :: Bool -> Int
fromBool = bool 0 1

-- | Move the ith element from the top of the stack to the top.
dig :: Int -> StackM ()
dig i = SM $ \s -> if i > 0 && i <= length s
                   then let (xs, y:ys) = splitAt (i-1) s
                        in Just ((), y : xs ++ ys)
                   else Nothing


-- ** Stack language semantics

binop :: (Int -> Int -> Int) -> StackM ()
binop f = liftM2 f pop pop >>= push

-- | Monadic semantics of commands.
cmd :: Cmd -> StackM ()
cmd (Push i)  = push i
cmd Drop      = void pop
cmd Dig       = pop >>= dig
cmd Dup       = peek >>= push
cmd Neg       = pop >>= (push . negate)
cmd Add       = binop (+)
cmd Mul       = binop (*)
cmd LEq       = binop ((fromBool.) . (<=))
cmd (If t e)  = popBool >>= bool (prog e) (prog t)
cmd (While b) = peekBool >>= bool (return ()) (prog $ b ++ [While b])

-- | Monadic semantics of programs.
prog :: Prog -> StackM ()
prog = mapM_ cmd


-- | Run a stack program with an initially empty stack, returning the
--   resulting stack or an error.
--
--   >>> runProg p1
--   Just [5,5,4]
--
--   >>> runProg p2
--   Just [12,14,13,11,10]
--
--   >>> runProg p3
--   Just [-5,27]
--
--   >>> runProg p4
--   Just [1,0]
-- 
--   >>> runProg p5
--   Just [22]
--
--   >>> runProg (Push 10 : fac)
--   Just [3628800]
--
--   >>> all (== Nothing) (map runProg bads)
--   True
--
runProg :: Prog -> Maybe Stack
runProg = fmap snd . runWith [] . prog