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