How can I make my stack monad faster?

Question

I made a stack monad that lets one manipulate and control values on a stack.

I want to know if it's correct (it seems correct), how to make it faster, and how I could detect stack overflows without carrying around the stack size and checking against it (I've heard about guard pages but I have no idea how I'd use them with Haskell).

{-# LANGUAGE RankNTypes #-}
module Main where
import Data.IORef
import Foreign.Ptr
import Foreign.Storable
import Foreign.Marshal.Alloc
import Control.Exception ( bracket )
import System.IO.Unsafe
newtype StackRef s a = StackRef (Ptr a)
data Stack = Stack (Ptr ())
newtype StackMonad s a = StackMonad (Stack -> IO (Stack, a))
instance Monad (StackMonad s) where
 return x = ioToStackMonad (return x)
 (StackMonad m) >>= f = StackMonad $ \stack -> do
 (newStack, x) <- m stack
 let (StackMonad g) = f x
 g newStack
runStackWithSize :: Int -> (forall s. StackMonad s a) -> a
runStackWithSize stackSize (StackMonad f) = unsafePerformIO $
 bracket (mallocBytes stackSize) free $ \theStack -> do
 (_, a) <- f (Stack theStack)
 return a
-- | 1024 bytes is the usual size for a stack
runStack :: (forall s. StackMonad s a) -> a
runStack = runStackWithSize 1024
newStackRef :: Storable a => a -> StackMonad s (StackRef s a)
newStackRef value = StackMonad $ \(Stack topOfStack) -> do
 let ptr = castPtr (alignPtr topOfStack (alignment value))
 poke ptr value
 return (Stack (plusPtr ptr (sizeOf value)), StackRef ptr)
ioToStackMonad :: IO a -> StackMonad s a
ioToStackMonad action = StackMonad $ \ptr -> do
 a <- action
 return (ptr, a)
writeStackRef :: Storable a => StackRef s a -> a -> StackMonad s ()
writeStackRef (StackRef p) val = ioToStackMonad (poke p val)
readStackRef :: Storable a => StackRef s a -> StackMonad s a
readStackRef (StackRef p) = ioToStackMonad (peek p)
{- |
 The type hackery means that pointers to the old stuff isn't reachable,
 therefore it's okay to pop the stack pointer.
-}
stack :: (forall s. (forall b. StackRef t b -> StackRef s b) -> StackMonad s a) -> StackMonad t a
stack f = let (StackMonad s) = f (\(StackRef p) -> StackRef p) in StackMonad $ \old_ptr -> do
 (_, state) <- s old_ptr
 return (old_ptr, state)
stack_ :: (forall s. (forall b. StackRef t b -> StackRef s b) -> StackMonad s ()) -> StackMonad t ()
stack_ f = let (StackMonad s) = f (\(StackRef p) -> StackRef p) in StackMonad $ \old_ptr -> do
 _ <- s old_ptr
 return (old_ptr, ())
x :: Int
x = runStack $ do
 a <- newStackRef 5
 c <- stack $ \lift1 -> do
 b <- newStackRef 1
 let liftedA = lift1 a
 aValue <- readStackRef liftedA
 writeStackRef liftedA 6
 bValue <- readStackRef b
 return (aValue + bValue)
 newAValue <- readStackRef a
 return (c + newAValue)
main = print x

• FUZxxl: Did you tried out an implementation where the stack is just a plain list? Letting the runtime do the work for you is usually faster.
• jpaugh: Actually, a list is a stack: it's biased toward accessing the first element. And, [] is a Monad... and likely heavily optimized, although it is not as fast as, say a Vector.
• user12742: Performance is only a side point. The point is semi-deterministic allocation of memory (which will also be neat for debugging.) For that reason just using a list is incorrect.
• PressingOnAlways: I don't know Haskell, but I know if I were approaching this problem in a language I'm familiar with, I would write tests whether it be unit tests or integration tests to determine if it works as expected. Run it through its paces and assert that you don't get overflows and the stack responds properly.

Answer 1

This code is simply incorrect. It is even not a stack (because it lacks pop operation and it is impossible to implement it in a type-safe way). Here is an example:

runStack $ do
 newStackRef (5:: Int)
 newStackRef False
 newStackRef 'x'

according to the definition of the newStackRef it just pokes Storable representation of the value into the memory block and advances the pointer. Any information about the type of that value is lost.

To pop a value from such a stack you need to specify correct type and compiler (or RTS) is not able to warn you if the type is not correct.

To the monad part of the question. From the instance Monad definition it is easy to see that it almost the same as the State monad (here is how it defined in Control.Monad.State.Trans.State.Lazy)

instance (Monad m) => Monad (StateT s m) where
 return a = state $ \s -> (a, s)
 m >>= k = StateT $ \s -> do
 ~(a, s') <- runStateT m s
 runStateT (k a) s'

The only difference is that it operates on top of IO an has fixed type of state. So it is actually StateT (Ptr ()) IO.

It is even not clear what it is for a stack to be an instance of Monad. E.g. list monad allows nondeterministic computations (in prolog-like style), Maybe monad allows computations that fail.

To enable computation to access and operate stack it is easier to use existing State monad with some stack implementation as a state. I see no reason to implement your own Stack monad.