Stack-based programming using monad transformers

Stack.hs

import Control.Monad
import Control.Monad.Identity (runIdentity)
import Control.Monad.Error
import Control.Monad.State.Lazy

import Data.Maybe (listToMaybe)


-- |Monad transformer which stores a stack internally
type StackT s m = StateT [s] m

-- |Stack of monad transformers representing our computation
type StackCompT e s m = ErrorT e (StackT s m)

push :: (Error e, Monad m) => s -> StackCompT e s m ()
push n = modify (n:)

pop :: (Error e, Monad m) => StackCompT e s m s
pop = do
    m <- liftM listToMaybe get
    modify tail
    case m of
      Nothing -> throwError . strMsg $ "stack underflow in pop"
      Just n -> return n

unaryS :: (Error e, Monad m) => (s -> s) -> StackCompT e s m ()
unaryS f = push =<< liftM f pop

binaryS :: (Error e, Monad m) => (s -> s -> s) -> StackCompT e s m ()
binaryS f = push =<< liftM2 f pop pop

-- Numeric stack ops
neg :: (Error e, Monad m, Num s) => StackCompT e s m ()
neg = unaryS negate

add :: (Error e, Monad m, Num s) => StackCompT e s m ()
add = binaryS (+)

sub :: (Error e, Monad m, Num s) => StackCompT e s m ()
sub = neg >> add

mul :: (Error e, Monad m, Num s) => StackCompT e s m ()
mul = binaryS (*)

divS :: (Error e, Monad m, Integral s) => StackCompT e s m ()
divS = binaryS div

modS :: (Error e, Monad m, Integral s) => StackCompT e s m ()
modS = binaryS mod


-- Utility function to throw an exception if we get wrong number of args
require :: (Error e, Monad m) => Int -> String -> StackCompT e s m ()
require n name = do
    stack <- get
    when (length stack < n) $ throwError . strMsg $ "stack underflow in " ++ name

-- Stack manipulation ops
nip :: (Error e, Monad m) => StackCompT e s m ()
nip = do
    require 2 "nip"
    modify $ \(a:b:xs) -> a:xs

dup :: (Error e, Monad m) => StackCompT e s m ()
dup = do
    require 1 "dup"
    modify $ \(a:xs) -> a:a:xs


over :: (Error e, Monad m) => StackCompT e s m ()
over = do
    require 2 "over"
    modify $ \(a:b:xs) -> b:a:b:xs

tuck :: (Error e, Monad m) => StackCompT e s m ()
tuck = do
    require 2 "tuck"
    modify $ \(a:b:xs) -> a:b:a:xs

swap :: (Error e, Monad m) => StackCompT e s m ()
swap = do
    require 2 "swap"
    modify $ \(a:b:xs) -> b:a:xs

rot :: (Error e, Monad m) => StackCompT e s m ()
rot = do
    require 3 "rot"
    modify $ \(a:b:c:xs) -> c:a:b:xs

rot' :: (Error e, Monad m) => StackCompT e s m ()
rot' = do
    require 3 "rot'"
    modify $ \(a:b:c:xs) -> b:c:a:xs

runStackCompT :: (Error e, Monad m) => StackCompT e s m a -> m (Either e [s])
runStackCompT comp = do
    (a, s) <- runStateT (runErrorT comp) []
    return $ case a of
       (Left err) -> Left err
       (Right _) -> Right s

-- An example of arithmetic in this monad
example :: (Error e, Monad m, Num s) => StackCompT e s m ()
example = do
    -- Compute 2*(3-5)^2
    let square = dup >> mul
    push 2
    push 3
    push 5
    sub
    square
    mul


main = case (runIdentity $ runStackCompT example :: Either String [Int]) of
    (Left err) -> putStrLn $ "Error while running stack computation: " ++ err
    (Right stack) -> putStr "Resulting stack: " >> print stack