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