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         ST.unsafeIOToST, ST.stToIO,
  30
  31         strictToLazyST, lazyToStrictST
  32     ) where
  33
  34 import qualified ST
  35 import qualified PrelST
  36 import PrelArr
  37 import PrelBase ( Eq(..), Int, Bool, ($), ()(..) )
  38 import Monad
  39 import Ix
  40 import PrelGHC
  41
  42 newtype ST s a = ST (State s -> (a, State s))
  43
  44 data State s = S# (State# s)
  45
  46 instance Functor (ST s) where
  47     fmap f m = ST $ \ s ->
  48       let
  49        ST m_a = m
  50        (r,new_s) = m_a s
  51       in
  52       (f r,new_s)
  53
  54 instance Monad (ST s) where
  55
  56         return a = ST $ \ s -> (a,s)
  57         m >> k   =  m >>= \ _ -> k
  58         fail s   = error s
  59
  60         (ST m) >>= k
  61          = ST $ \ s ->
  62            let
  63              (r,new_s) = m s
  64              ST k_a = k r
  65            in
  66            k_a new_s
  67
  68 {-# NOINLINE runST #-}
  69 runST :: (forall s. ST s a) -> a
  70 runST st = case st of ST the_st -> let (r,_) = the_st (S# realWorld#) in r
  71 \end{code}
  72
  73 %*********************************************************
  74 %*                                                      *
  75 \subsection{Variables}
  76 %*                                                      *
  77 %*********************************************************
  78
  79 \begin{code}
  80 newSTRef   :: a -> ST s (ST.STRef s a)
  81 readSTRef  :: ST.STRef s a -> ST s a
  82 writeSTRef :: ST.STRef s a -> a -> ST s ()
  83
  84 newSTRef   = strictToLazyST . ST.newSTRef
  85 readSTRef  = strictToLazyST . ST.readSTRef
  86 writeSTRef r a = strictToLazyST (ST.writeSTRef r a)
  87
  88 \end{code}
  89
  90 %*********************************************************
  91 %*                                                      *
  92 \subsection{Arrays}
  93 %*                                                      *
  94 %*********************************************************
  95
  96 \begin{code}
  97 newtype STArray s ix elt = STArray (MutableArray s ix elt)
  98
  99 newSTArray          :: Ix ix => (ix,ix) -> elt -> ST s (STArray s ix elt)
 100 readSTArray         :: Ix ix => STArray s ix elt -> ix -> ST s elt
 101 writeSTArray        :: Ix ix => STArray s ix elt -> ix -> elt -> ST s ()
 102 boundsSTArray       :: Ix ix => STArray s ix elt -> (ix, ix)
 103 thawSTArray         :: Ix ix => Array ix elt -> ST s (STArray s ix elt)
 104 freezeSTArray       :: Ix ix => STArray s ix elt -> ST s (Array ix elt)
 105 unsafeFreezeSTArray :: Ix ix => STArray s ix elt -> ST s (Array ix elt)
 106
 107 newSTArray ixs init     =
 108            strictToLazyST (newArray ixs init) >>= \arr ->
 109            return (STArray arr)
 110
 111 readSTArray (STArray arr) ix = strictToLazyST (readArray arr ix)
 112 writeSTArray (STArray arr) ix v = strictToLazyST (writeArray arr ix v)
 113 boundsSTArray (STArray arr) = boundsOfArray arr
 114 thawSTArray arr =
 115             strictToLazyST (thawArray arr) >>= \ marr ->
 116             return (STArray marr)
 117
 118 freezeSTArray (STArray arr) = strictToLazyST (freezeArray arr)
 119 unsafeFreezeSTArray (STArray arr) = strictToLazyST (unsafeFreezeArray arr)
 120
 121 strictToLazyST :: PrelST.ST s a -> ST s a
 122 strictToLazyST m = ST $ \s ->
 123         let
 124            pr = case s of { S# s# -> PrelST.liftST m s# }
 125            r  = case pr of { PrelST.STret _ v -> v }
 126            s' = case pr of { PrelST.STret s2# _ -> S# s2# }
 127         in
 128         (r, s')
 129
 130 lazyToStrictST :: ST s a -> PrelST.ST s a
 131 lazyToStrictST (ST m) = PrelST.ST $ \s ->
 132         case (m (S# s)) of (a, S# s') -> (# s', a #)
 133
 134 unsafeInterleaveST :: ST s a -> ST s a
 135 unsafeInterleaveST = strictToLazyST . ST.unsafeInterleaveST . lazyToStrictST
 136
 137 \end{code}