Haskell でスタックを(2)

Haskell でスタックを書く、続き。スタックのサイズは固定のまま、スタックの要素を変更できるようにした。

  1. スタックは Foreign.Marshal.Array.mallocArray で確保する。
  2. 確保した領域へのポインタは(グローバル変数の替わりに)状態モナドで管理する。
-- Stack
import Control.Monad.State
import Foreign.Ptr
import Foreign.Marshal.Alloc
import Foreign.Marshal.Array
import Foreign.Storable

stackSize :: Int
stackSize = 20

initStack :: Storable a => IO (Ptr a, Int)
initStack = do
  p <- mallocArray stackSize
  return (p, 0)
-- https://www.haskell.org/haskellwiki/Scoped_type_variables より
sizeOfPtr :: (Storable a) => Ptr a -> Int
sizeOfPtr = sizeOf . (undefined :: Ptr a -> a)

push :: (Storable a, Num a) => a -> StateT (Ptr a, Int) IO ()
push x = do
  (stack, sp) <- get
  if sp >= stackSize then
    error "oops. Stack overflow"
    else
    do
      lift $ pokeElemOff stack sp x
      put (stack, sp + 1)
      return ()

pop :: (Storable a, Num a) => StateT (Ptr a, Int) IO a
pop = do
  (stack, sp) <- get
  if sp > 0 then
    do
      x <- lift $ peekElemOff stack (sp - 1)
      put (stack, sp - 1)
      return x
    else error "oops. Stack undeflow"

calc :: (Storable a, Num a) => (a -> a -> a) -> StateT (Ptr a, Int) IO ()
calc f = do
  x <- pop
  y <- pop
  push (f x y)

add :: (Storable a, Num a) => StateT (Ptr a, Int) IO ()
add = calc (+)
sub :: (Storable a, Num a) => StateT (Ptr a, Int) IO ()
sub = calc (-)

以下のように使う。

-- initStack >>= runStateT test0

test0 :: StateT (Ptr Int, Int) IO Int
test0 = do
  push 9
  pop
  
test1 :: StateT (Ptr Int, Int) IO ()
test1 = do
  push 3
  push 99
  x <- pop
  lift $ putStrLn $ "x: " ++ show x
  y <- pop
  lift $ putStrLn $ "y: " ++ show y

test2 :: StateT (Ptr Int, Int) IO Int
test2 = do
  push 8
  push 9
  pop
  pop  
  push 11
  push 13  
  pop
  pop

test4 :: StateT (Ptr Float, Int) IO Float
test4 = do
  push 5.0
  push 7.0
  push 8.0
  add
  sub
  pop

参考