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