ghc/lib/exts/LazyST.lhs

   1 %
   2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1997
   3 %
   4
   5 \section[LazyST]{The Lazy State Transformer Monad, @LazyST@}
   6
   7 This module presents an identical interface to ST, but the underlying
   8 implementation of the state thread is lazy.
   9
  10 \begin{code}
  11 module LazyST (
  12
  13         ST,
  14
  15         runST,
  16         unsafeInterleaveST,
  17
  18         -- ST is one, so you'll likely need some Monad bits
  19         module Monad,
  20
  21         ST.STRef,
  22         newSTRef, readSTRef, writeSTRef,
  23
  24         STArray,
  25         newSTArray, readSTArray, writeSTArray, boundsSTArray,
  26         thawSTArray, freezeSTArray, unsafeFreezeSTArray,
  27         Ix,
  28
  29         strictToLazyST, lazyToStrictST
  30     ) where
  31
  32 import qualified ST
  33 import qualified PrelST
  34 import PrelArr
  35 import PrelBase ( Eq(..), Int, Bool, ($), ()(..) )
  36 import Monad
  37 import Ix
  38 import PrelGHC
  39
  40 newtype ST s a = ST (PrelST.State s -> (a,PrelST.State s))
  41
  42 instance Functor (ST s) where
  43     map f m = ST $ \ s ->
  44       let
  45        ST m_a = m
  46        (r,new_s) = m_a s
  47       in
  48       (f r,new_s)
  49
  50 instance Monad (ST s) where
  51
  52         return a = ST $ \ s -> (a,s)
  53         m >> k   =  m >>= \ _ -> k
  54
  55         (ST m) >>= k
  56          = ST $ \ s ->
  57            let
  58              (r,new_s) = m s
  59              ST k_a = k r
  60            in
  61            k_a new_s
  62
  63 {-# NOINLINE runST #-}
  64 runST :: (forall s. ST s a) -> a
  65 runST st = case st of ST st -> let (r,_) = st (PrelST.S# realWorld#) in r
  66 \end{code}
  67
  68 %*********************************************************
  69 %*                                                      *
  70 \subsection{Variables}
  71 %*                                                      *
  72 %*********************************************************
  73
  74 \begin{code}
  75 newSTRef   :: a -> ST s (ST.STRef s a)
  76 readSTRef  :: ST.STRef s a -> ST s a
  77 writeSTRef :: ST.STRef s a -> a -> ST s ()
  78
  79 newSTRef   = strictToLazyST . ST.newSTRef
  80 readSTRef  = strictToLazyST . ST.readSTRef
  81 writeSTRef r a = strictToLazyST (ST.writeSTRef r a)
  82 \end{code}
  83
  84 %*********************************************************
  85 %*                                                      *
  86 \subsection{Arrays}
  87 %*                                                      *
  88 %*********************************************************
  89
  90 \begin{code}
  91 newtype STArray s ix elt = STArray (MutableArray s ix elt)
  92
  93 newSTArray          :: Ix ix => (ix,ix) -> elt -> ST s (STArray s ix elt)
  94 readSTArray         :: Ix ix => STArray s ix elt -> ix -> ST s elt
  95 writeSTArray        :: Ix ix => STArray s ix elt -> ix -> elt -> ST s ()
  96 boundsSTArray       :: Ix ix => STArray s ix elt -> (ix, ix)
  97 thawSTArray         :: Ix ix => Array ix elt -> ST s (STArray s ix elt)
  98 freezeSTArray       :: Ix ix => STArray s ix elt -> ST s (Array ix elt)
  99 unsafeFreezeSTArray :: Ix ix => STArray s ix elt -> ST s (Array ix elt)
 100
 101 newSTArray ixs init     =
 102            strictToLazyST (newArray ixs init) >>= \arr ->
 103            return (STArray arr)
 104
 105 readSTArray (STArray arr) ix = strictToLazyST (readArray arr ix)
 106 writeSTArray (STArray arr) ix v = strictToLazyST (writeArray arr ix v)
 107 boundsSTArray (STArray arr) = boundsOfArray arr
 108 thawSTArray arr =
 109             strictToLazyST (thawArray arr) >>= \arr ->
 110             return (STArray arr)
 111 freezeSTArray (STArray arr) = strictToLazyST (freezeArray arr)
 112 unsafeFreezeSTArray (STArray arr) = strictToLazyST (unsafeFreezeArray arr)
 113
 114 strictToLazyST :: PrelST.ST s a -> ST s a
 115 strictToLazyST m = ST $ \s ->
 116         let
 117            pr = case s of { PrelST.S# s# -> PrelST.liftST m s# }
 118            r  = case pr of { PrelST.STret s2# r -> r }
 119            s' = case pr of { PrelST.STret s2# r -> PrelST.S# s2# }
 120         in
 121         (r, s')
 122
 123 lazyToStrictST :: ST s a -> PrelST.ST s a
 124 lazyToStrictST (ST m) = PrelST.ST $ \s ->
 125         case (m (PrelST.S# s)) of (a, PrelST.S# s') -> (# s', a #)
 126
 127 unsafeInterleaveST :: ST s a -> ST s a
 128 unsafeInterleaveST = strictToLazyST . ST.unsafeInterleaveST . lazyToStrictST
 129
 130 \end{code}