module Hasklet4 where

import Control.Monad


-- * 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))

-- type StackM a = MaybeT (State Stack) a   -- 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 x   = SM (\s -> Just (x,s))

  -- (>>=) :: StateM a -> (a -> StateM b) -> StateM b
  -- SM f >>= g = SM (\s -> case f s of
  --                          Nothing      -> Nothing
  --                          Just (x, s') -> let SM h = g x
  --                                          in h s')

  SM f >>= g = SM (\s -> f s >>= (\(a,s') -> runWith s' (g a)))


-- ** Primitive operations

-- | Push a value onto the stack.
push :: Int -> StackM ()
push i = SM $ \s -> Just ((), i:s)

-- | Pop a value off the stack and return it.
pop :: StackM Int
pop = SM $ \s -> case s of
                   i:s -> Just (i,s)
                   _   -> Nothing

-- | Peek at the value on top of the stack without popping it.
peek :: StackM Int
peek = SM $ \s -> case s of
                    i:s -> Just (i, i:s)
                    _   -> Nothing

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

-- | Check if a value is equivalent to true.
isTrue :: Int -> Bool
isTrue 0 = False
isTrue _ = True

-- | Monadic semantics of commands.
cmd :: Cmd -> StackM ()
cmd (Push i)  = push i
cmd Drop      = pop >> return ()
cmd Dig       = pop >>= dig
cmd Dup       = peek >>= push
cmd Neg       = pop >>= push . negate
cmd Add       = liftM2 (+) pop pop >>= push
cmd Mul       = liftM2 (*) pop pop >>= push
cmd LEq       = liftM2 (<=) pop pop >>= \b -> if b then push 1 else push 0
cmd (If t e)  = do c <- pop
                   if isTrue c
                     then prog t
                     else prog e
cmd (While b) = do c <- peek
                   if isTrue c
                     then prog b >> cmd (While b)
                     else return ()

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

-- mapM_ :: Monad m => (a -> m b) -> [a] -> m ()
-- mapM_ f []     = return ()
-- mapM_ f (x:xs) = f x >> mapM_ f xs

-- prog []     = return ()
-- prog (c:cs) = cmd c >> prog cs


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