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                  -- abstract, instance of Functor, Monad.
  12       , runST               -- :: (forall s. ST s a) -> a
  13       , fixST               -- :: (a -> ST s a) -> ST s a
  14       , unsafeInterleaveST  -- :: ST s a -> ST s a
  15
  16       , STRef
  17       , newSTRef
  18       , readSTRef
  19       , writeSTRef
  20
  21       , unsafeIOToST
  22       , stToIO
  23
  24       , STArray
  25       , newSTArray
  26       , readSTArray
  27       , writeSTArray
  28       , boundsSTArray
  29       , thawSTArray
  30       , freezeSTArray
  31       , unsafeFreezeSTArray
  32 #ifndef __HUGS__
  33 -- no 'good' reason, just doesn't support it right now.
  34       , unsafeThawSTArray
  35 #endif
  36
  37       ) where
  38
  39 #ifdef __HUGS__
  40 import PreludeBuiltin
  41 #define MutableVar Ref
  42 #define readVar    primReadRef
  43 #define writeVar   primWriteRef
  44 #define newVar     primNewRef
  45 #else
  46 import PrelArr
  47 import PrelST
  48 import PrelBase ( Eq(..), Int, Bool, ($), ()(..), unsafeCoerce# )
  49 import PrelIOBase ( IO(..), stToIO )
  50 #endif
  51 import Monad
  52 import Ix
  53
  54 \end{code}
  55
  56 %*********************************************************
  57 %*                                                      *
  58 \subsection{Variables}
  59 %*                                                      *
  60 %*********************************************************
  61
  62 \begin{code}
  63 newtype STRef s a = STRef (MutableVar s a)
  64     deriving Eq
  65
  66 newSTRef :: a -> ST s (STRef s a)
  67 newSTRef v = newVar v >>= \ var -> return (STRef var)
  68
  69 readSTRef :: STRef s a -> ST s a
  70 readSTRef (STRef var) = readVar var
  71
  72 writeSTRef :: STRef s a -> a -> ST s ()
  73 writeSTRef (STRef var) v = writeVar var v
  74 \end{code}
  75
  76 %*********************************************************
  77 %*                                                      *
  78 \subsection{Arrays}
  79 %*                                                      *
  80 %*********************************************************
  81
  82 \begin{code}
  83 newSTArray              :: Ix ix => (ix,ix) -> elt -> ST s (STArray s ix elt)
  84 writeSTArray            :: Ix ix => STArray s ix elt -> ix -> elt -> ST s ()
  85 readSTArray             :: Ix ix => STArray s ix elt -> ix -> ST s elt
  86 boundsSTArray           :: Ix ix => STArray s ix elt -> (ix, ix)
  87 thawSTArray             :: Ix ix => Array ix elt -> ST s (STArray s ix elt)
  88 freezeSTArray           :: Ix ix => STArray s ix elt -> ST s (Array ix elt)
  89 unsafeFreezeSTArray     :: Ix ix => STArray s ix elt -> ST s (Array ix elt)
  90
  91 #ifndef __HUGS__
  92 -- see export list comment..
  93 unsafeThawSTArray       :: Ix ix => Array ix elt -> ST s (STArray s ix elt)
  94 #endif
  95
  96 #ifdef __HUGS__
  97 data STArray s ix elt = STArray (ix,ix) (PrimMutableArray s elt)
  98   deriving Eq
  99
 100 newSTArray ixs elt = do
 101   { arr <- primNewArray (rangeSize ixs) elt
 102   ; return (STArray ixs arr)
 103   }
 104
 105 boundsSTArray (STArray ixs arr)        = ixs
 106 readSTArray   (STArray ixs arr) ix     = primReadArray arr (index ixs ix)
 107 writeSTArray  (STArray ixs arr) ix elt = primWriteArray arr (index ixs ix) elt
 108 freezeSTArray (STArray ixs arr)        = do
 109   { arr' <- primFreezeArray arr
 110   ; return (Array ixs arr')
 111   }
 112
 113 unsafeFreezeSTArray (STArray ixs arr)  = do
 114   { arr' <- primUnsafeFreezeArray arr
 115   ; return (Array ixs arr')
 116   }
 117
 118 thawSTArray (Array ixs arr) = do
 119   { arr' <- primThawArray arr
 120   ; return (STArray ixs arr')
 121   }
 122
 123 primFreezeArray :: PrimMutableArray s a -> ST s (PrimArray a)
 124 primFreezeArray arr = do
 125   { let n = primSizeMutableArray arr
 126   ; arr' <- primNewArray n arrEleBottom
 127   ; mapM_ (copy arr arr') [0..n-1]
 128   ; primUnsafeFreezeArray arr'
 129   }
 130  where
 131   copy arr arr' i = do { x <- primReadArray arr i; primWriteArray arr' i x }
 132   arrEleBottom = error "primFreezeArray: panic"
 133
 134 primThawArray :: PrimArray a -> ST s (PrimMutableArray s a)
 135 primThawArray arr = do
 136   { let n = primSizeArray arr
 137   ; arr' <- primNewArray n arrEleBottom
 138   ; mapM_ (copy arr arr') [0..n-1]
 139   ; return arr'
 140   }
 141  where
 142   copy arr arr' i = primWriteArray arr' i (primIndexArray arr i)
 143   arrEleBottom = error "primFreezeArray: panic"
 144 #else
 145 newtype STArray s ix elt = STArray (MutableArray s ix elt)
 146     deriving Eq
 147
 148 newSTArray ixs elt =
 149     newArray ixs elt >>= \arr ->
 150     return (STArray arr)
 151
 152 boundsSTArray (STArray arr) = boundsOfArray arr
 153
 154 readSTArray (STArray arr) ix = readArray arr ix
 155
 156 writeSTArray (STArray arr) ix elt = writeArray arr ix elt
 157
 158 thawSTArray arr = thawArray arr >>= \starr -> return (STArray starr)
 159
 160 freezeSTArray (STArray arr) = freezeArray arr
 161
 162 unsafeFreezeSTArray (STArray arr) = unsafeFreezeArray arr
 163 unsafeThawSTArray arr = unsafeThawArray arr >>= \ marr -> return (STArray marr)
 164
 165 #endif
 166 \end{code}
 167
 168
 169 \begin{code}
 170 unsafeIOToST       :: IO a -> ST s a
 171 #ifdef __HUGS__
 172 unsafeIOToST = primUnsafeCoerce
 173 #else
 174 unsafeIOToST (IO io) = ST $ \ s ->
 175     case ((unsafeCoerce# io) s) of
 176       (#  new_s, a #) -> unsafeCoerce# (STret new_s a)
 177 --      IOfail new_s e -> error ("I/O Error (unsafeIOToST): " ++ showsPrec 0 e "\n")
 178 #endif
 179 \end{code}