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

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.