ghc/lib/exts/ST.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1996
   3 %
   4 \section[module_ST]{The State Transformer Monad, @ST@}
   5
   6 \begin{code}
   7 {-# OPTIONS -fno-implicit-prelude #-}
   8
   9 module ST (
  10
  11         ST,
  12
  13         runST,                          -- :: (forall s. ST s a) -> a
  14         fixST,                          -- :: (a -> ST s a) -> ST s a
  15
  16         unsafeInterleaveST,
  17
  18         -- ST is one, so you'll likely need some Monad bits
  19         module Monad,
  20
  21         STRef,
  22         newSTRef, readSTRef, writeSTRef,
  23
  24         STArray,
  25         newSTArray, readSTArray, writeSTArray, boundsSTArray,
  26         thawSTArray, freezeSTArray, unsafeFreezeSTArray,
  27         Ix
  28
  29     ) where
  30
  31 #ifdef __HUGS__
  32 import PreludeBuiltin
  33 #define MutableVar Ref
  34 #define readVar    primReadRef
  35 #define writeVar   primWriteRef
  36 #define newVar     primNewRef
  37 #else
  38 import PrelArr
  39 import PrelST
  40 import PrelBase ( Eq(..), Int, Bool, ($), ()(..) )
  41 #endif
  42 import Monad
  43 import Ix
  44
  45 \end{code}
  46
  47 %*********************************************************
  48 %*                                                      *
  49 \subsection{Variables}
  50 %*                                                      *
  51 %*********************************************************
  52
  53 \begin{code}
  54 newtype STRef s a = STRef (MutableVar s a)
  55     deriving Eq
  56
  57 newSTRef :: a -> ST s (STRef s a)
  58 newSTRef v = newVar v >>= \ var -> return (STRef var)
  59
  60 readSTRef :: STRef s a -> ST s a
  61 readSTRef (STRef var) = readVar var
  62
  63 writeSTRef :: STRef s a -> a -> ST s ()
  64 writeSTRef (STRef var) v = writeVar var v
  65 \end{code}
  66
  67 %*********************************************************
  68 %*                                                      *
  69 \subsection{Arrays}
  70 %*                                                      *
  71 %*********************************************************
  72
  73 \begin{code}
  74 newSTArray              :: Ix ix => (ix,ix) -> elt -> ST s (STArray s ix elt)
  75 writeSTArray            :: Ix ix => STArray s ix elt -> ix -> elt -> ST s ()
  76 readSTArray             :: Ix ix => STArray s ix elt -> ix -> ST s elt
  77 boundsSTArray           :: Ix ix => STArray s ix elt -> (ix, ix)
  78 thawSTArray             :: Ix ix => Array ix elt -> ST s (STArray s ix elt)
  79 freezeSTArray           :: Ix ix => STArray s ix elt -> ST s (Array ix elt)
  80 unsafeFreezeSTArray     :: Ix ix => STArray s ix elt -> ST s (Array ix elt)
  81
  82 #ifdef __HUGS__
  83 data STArray s ix elt = STArray (ix,ix) (PrimMutableArray s elt)
  84   deriving Eq
  85
  86 newSTArray ixs elt = do
  87   { arr <- primNewArray (rangeSize ixs) elt
  88   ; return (STArray ixs arr)
  89   }
  90
  91 boundsSTArray (STArray ixs arr)        = ixs
  92 readSTArray   (STArray ixs arr) ix     = primReadArray arr (index ixs ix)
  93 writeSTArray  (STArray ixs arr) ix elt = primWriteArray arr (index ixs ix) elt
  94 freezeSTArray (STArray ixs arr)        = do
  95   { arr' <- primFreezeArray arr
  96   ; return (Array ixs arr')
  97   }
  98
  99 unsafeFreezeSTArray (STArray ixs arr)  = do
 100   { arr' <- primUnsafeFreezeArray arr
 101   ; return (Array ixs arr')
 102   }
 103
 104 thawSTArray (Array ixs arr) = do
 105   { arr' <- primThawArray arr
 106   ; return (STArray ixs arr')
 107   }
 108
 109 primFreezeArray :: PrimMutableArray s a -> ST s (PrimArray a)
 110 primFreezeArray arr = do
 111   { let n = primSizeMutableArray arr
 112   ; arr' <- primNewArray n arrEleBottom
 113   ; mapM_ (copy arr arr') [0..n-1]
 114   ; primUnsafeFreezeArray arr'
 115   }
 116  where
 117   copy arr arr' i = do { x <- primReadArray arr i; primWriteArray arr' i x }
 118   arrEleBottom = error "primFreezeArray: panic"
 119
 120 primThawArray :: PrimArray a -> ST s (PrimMutableArray s a)
 121 primThawArray arr = do
 122   { let n = primSizeArray arr
 123   ; arr' <- primNewArray n arrEleBottom
 124   ; mapM_ (copy arr arr') [0..n-1]
 125   ; return arr'
 126   }
 127  where
 128   copy arr arr' i = primWriteArray arr' i (primIndexArray arr i)
 129   arrEleBottom = error "primFreezeArray: panic"
 130 #else
 131 newtype STArray s ix elt = STArray (MutableArray s ix elt)
 132     deriving Eq
 133
 134 newSTArray ixs elt =
 135     newArray ixs elt >>= \arr ->
 136     return (STArray arr)
 137
 138 boundsSTArray (STArray arr) = boundsOfArray arr
 139
 140 readSTArray (STArray arr) ix = readArray arr ix
 141
 142 writeSTArray (STArray arr) ix elt = writeArray arr ix elt
 143
 144 thawSTArray arr = thawArray arr >>= \starr -> return (STArray starr)
 145
 146 freezeSTArray (STArray arr) = freezeArray arr
 147
 148 unsafeFreezeSTArray (STArray arr) = unsafeFreezeArray arr
 149 #endif
 150 \end{code}
 151