Add lazy ST monad

lib/AllOfLib.hs

@@ -18,7 +18,11 @@ import Control.Monad
 import Control.Monad.Fail
 import Control.Monad.IO.Class
 import Control.Monad.ST
+import Control.Monad.ST.Lazy
+import Control.Monad.ST.Lazy.Unsafe
+import Control.Monad.ST.Strict
 import Control.Monad.ST_Type
+import Control.Monad.ST.Unsafe
 import Data.Bits
 import Data.Bool
 import Data.Bool_Type
@@ -94,6 +98,8 @@ import Data.RealFrac
 import Data.Records
 import Data.Semigroup
 import Data.STRef
+import Data.STRef.Lazy
+import Data.STRef.Strict
 import Data.String
 import Data.Text
 import Data.Text.Encoding

lib/Control/Monad/ST.hs

@@ -1,6 +1,7 @@
 module Control.Monad.ST(
   ST,
   runST,
+  fixST,
   --
   RealWorld, stToIO,  -- GHC compat
   ) where
@@ -13,6 +14,9 @@ import System.IO(fixIO)
 runST :: (forall s . ST s a) -> a
 runST (ST ioa) = primPerformIO ioa
 
+fixST :: (a -> ST s a) -> ST s a
+fixST f = ST (fixIO (unST . f))
+
 instance Functor (ST s) where
   fmap f (ST x) = ST (fmap f x)
 
@@ -24,13 +28,13 @@ instance Monad (ST s) where
   ST x >>= f = ST (x >>= (unST . f))
 
 instance MonadFix (ST s) where
-  mfix f = ST (fixIO (unST . f))
+  mfix = fixST
 
 ---------------------------------
 -- This does not belong here since it's GHC specific,
 -- but to be compatible we do it the same way.
 
 data RealWorld  -- Just to be compatible with GHC.  We don't use it.
 
-stToIO :: forall a . ST RealWorld a -> IO a
+stToIO :: ST RealWorld a -> IO a
 stToIO = unST

lib/Control/Monad/ST/Lazy.hs

@@ -0,0 +1,60 @@
+module Control.Monad.ST.Lazy(
+  ST,
+  runST,
+  fixST,
+  strictToLazyST,
+  lazyToStrictST,
+  --
+  RealWorld, stToIO,  -- GHC compat
+  ) where
+import qualified Prelude(); import MiniPrelude
+import Control.Monad.Fix
+import Control.Monad.ST(RealWorld)
+import qualified Control.Monad.ST_Type as S
+import System.IO.Unsafe(unsafeInterleaveIO, unsafePerformIO)
+
+newtype ST s a = ST { unST :: State s -> (a, State s) }
+
+data State s = S
+
+runST :: (forall s . ST s a) -> a
+runST (ST st) = case st S of (r, _) -> r
+
+fixST :: (a -> ST s a) -> ST s a
+fixST f = ST $ \s -> let q@(r, _s') = unST (f r) s in q
+
+strictToLazyST :: S.ST s a -> ST s a
+strictToLazyST (S.ST io) = ST $ \s -> s `seq`
+  case unsafePerformIO (fmap MkSolo io) of
+    MkSolo a -> (a, s)
+
+lazyToStrictST :: ST s a -> S.ST s a
+lazyToStrictST (ST st) = S.ST $ fmap id (case st S of (x, _) -> pure x)
+
+instance Functor (ST s) where
+  fmap f m = ST $ \s ->
+    let (r, s') = unST m s
+    in (f r, s')
+
+instance Applicative (ST s) where
+  pure x = ST $ \s -> (x, s)
+  liftA2 f m n = ST $ \s ->
+    let
+      (x, s') = unST m s
+      (y, s'') = unST n s'
+    in (f x y, s'')
+
+instance Monad (ST s) where
+  m >>= k = ST $ \s ->
+    let (r, s') = unST m s
+    in unST (k r) s'
+
+instance MonadFix (ST s) where
+  mfix = fixST
+
+---------------------------------
+-- This does not belong here since it's GHC specific,
+-- but to be compatible we do it the same way.
+
+stToIO :: ST RealWorld a -> IO a
+stToIO = S.unST . lazyToStrictST

lib/Control/Monad/ST/Lazy/Unsafe.hs

@@ -0,0 +1,15 @@
+module Control.Monad.ST.Lazy.Unsafe(
+  unsafeInterleaveST,
+  unsafeIOToST,
+  ) where
+import qualified Prelude();
+import Primitives(IO)
+import qualified Control.Monad.ST.Unsafe as S
+import Control.Monad.ST.Lazy
+import Data.Function
+
+unsafeInterleaveST :: ST s a -> ST s a
+unsafeInterleaveST = strictToLazyST . S.unsafeInterleaveST . lazyToStrictST
+
+unsafeIOToST :: IO a -> ST s a
+unsafeIOToST = strictToLazyST . S.unsafeIOToST

lib/Control/Monad/ST/Strict.hs

@@ -0,0 +1,3 @@
+module Control.Monad.ST.Strict (module Control.Monad.ST) where
+
+import Control.Monad.ST

lib/Control/Monad/ST/Unsafe.hs

@@ -4,6 +4,8 @@ module Control.Monad.ST.Unsafe(
   unsafeIOToST,
   unsafeSTToIO,
   ) where
+import qualified Prelude()
+import Primitives(IO)
 import Control.Monad.ST_Type
 import System.IO.Unsafe
 

lib/Control/Monad/ST_Type.hs

@@ -9,5 +9,5 @@ import {-# SOURCE #-} Data.Typeable
 -- The ST monad is implemented with the IO monad.
 newtype ST s a = ST (IO a)
 
-unST :: forall s a . ST s a -> IO a
+unST :: ST s a -> IO a
 unST (ST io) = io

lib/Data/STRef/Lazy.hs

@@ -0,0 +1,36 @@
+-----------------------------------------------------------------------------
+-- |
+-- Module      :  Data.STRef.Lazy
+-- Copyright   :  (c) The University of Glasgow 2001
+-- License     :  BSD-style (see the file libraries/base/LICENSE)
+--
+-- Maintainer  :  [email protected]
+-- Stability   :  stable
+-- Portability :  non-portable (uses Control.Monad.ST.Lazy)
+--
+-- Mutable references in the lazy ST monad.
+--
+-----------------------------------------------------------------------------
+
+module Data.STRef.Lazy (
+        -- * STRefs
+        ST.STRef,       -- abstract
+        newSTRef,
+        readSTRef,
+        writeSTRef,
+        modifySTRef
+ ) where
+
+import qualified Prelude (); import MiniPrelude
+import Control.Monad.ST.Lazy
+import qualified Data.STRef as ST
+
+newSTRef    :: a -> ST s (ST.STRef s a)
+readSTRef   :: ST.STRef s a -> ST s a
+writeSTRef  :: ST.STRef s a -> a -> ST s ()
+modifySTRef :: ST.STRef s a -> (a -> a) -> ST s ()
+
+newSTRef        = strictToLazyST . ST.newSTRef
+readSTRef       = strictToLazyST . ST.readSTRef
+writeSTRef  r a = strictToLazyST (ST.writeSTRef r a)
+modifySTRef r f = strictToLazyST (ST.modifySTRef r f)

lib/Data/STRef/Strict.hs

@@ -0,0 +1,3 @@
+module Data.STRef.Strict (module Data.STRef) where
+
+import Data.STRef

lib/base.cabal

@@ -45,6 +45,9 @@ library base
         Control.Monad.Fix
         Control.Monad.IO.Class
         Control.Monad.ST
+        Control.Monad.ST.Lazy
+        Control.Monad.ST.Lazy.Unsafe
+        Control.Monad.ST.Strict
         Control.Monad.ST.Unsafe
         Control.Monad.STM
         Control.Monad.Zip
@@ -122,6 +125,8 @@ library base
         Data.RealFrac
         Data.Records
         Data.STRef
+        Data.STRef.Lazy
+        Data.STRef.Strict
         Data.Semigroup
         Data.String
         Data.Text

tests/LazyST.hs

@@ -0,0 +1,110 @@
+-- taken from https://gitlab.haskell.org/ghc/ghc/-/blob/c94aaacd4c4e31a2fe2cb8dadcdd14c7621d27c5/libraries/base/tests/lazySTexamples.hs
+
+import Data.STRef.Lazy
+import Control.Monad.ST.Lazy as L
+import Control.Monad.ST as S
+import qualified Data.STRef as S
+import Data.Function (fix)
+import System.IO (hPutStrLn, stderr)
+import Debug.Trace (trace)
+
+-- The following implements `fix` using lazy `ST`. It is based on code
+-- by Oleg Kiselyov (source: http://okmij.org/ftp/Haskell/Fix.hs) which is
+-- in the public domain according to the main page (http://okmij.org/ftp/).
+
+fact :: (Int -> Int) -> Int -> Int
+fact self 0 = 1
+fact self n = n * self (pred n)
+
+-- Test liftM style (Oleg's original style)
+fix1 :: (a -> a) -> a
+fix1 f = L.runST $ do
+  wrap <- newSTRef (error "black hole")
+  let aux = readSTRef wrap >>= (\x -> x >>= pure . f)
+  writeSTRef wrap aux
+  aux
+
+-- Test fmap style
+fix2 :: (a -> a) -> a
+fix2 f = L.runST $ do
+  wrap <- newSTRef (error "black hole")
+  let aux = readSTRef wrap >>= \x -> f <$> x
+  writeSTRef wrap aux
+  aux
+
+-- The following examples are by Albert Y. C. Lai, and included (under the
+-- GHC license) with his permission:
+-- https://mail.haskell.org/pipermail/haskell-cafe/2017-January/126182.html
+
+example1 :: [Int]
+example1 = L.runST $ do
+    v <- strictToLazyST (S.newSTRef 0)
+    fix (\loop -> do
+            n <- strictToLazyST (do n <- S.readSTRef v
+                                    S.writeSTRef v (n+1)
+                                    return n
+                                )
+            ns <- loop
+            return (n : ns))
+
+example2 :: [Int]
+example2 = L.runST $ do
+    v <- strictToLazyST (S.newSTRef 0)
+    sequence (repeat (strictToLazyST (do n <- S.readSTRef v
+                                         S.writeSTRef v (n+1)
+                                         return n
+                                     )))
+
+example3 :: L.ST s [Integer]
+example3 = do
+    r <- newSTRef 0
+    let loop = do
+            x <- readSTRef r
+            writeSTRef r $ x + 1
+            xs <- loop
+            writeSTRef r $ x + 2
+            return $ x : xs
+    loop
+
+example4 :: L.ST s [Integer]
+example4 = do
+    r <- newSTRef 0
+    let loop = do
+            x <- readSTRef r
+            writeSTRef r $ x + 1
+            xs <- loop
+            error "this line is dead code"
+            return $ x : xs
+    loop
+
+star n s = trace ("<" ++ s ++ show n ++ ">") (return ())
+
+-- Albert called this "Sprinkle sprinkle little stars, how
+-- I wonder when you are"
+example5 :: L.ST s [Integer]
+example5 = do
+    star 0 "init begin"
+    r <- newSTRef 0
+    star 0 "init end"
+    let loop n = do
+            star n "A"
+            x <- readSTRef r
+            star n "B"
+            writeSTRef r $ x + 1
+            star n "C"
+            xs <- loop (n+1)
+            star n "D"
+            writeSTRef r $ x + 2
+            star n "E"
+            return $ x : xs
+    loop 0
+
+main :: IO ()
+main = do
+  print $ fix1 fact 5
+  print $ fix2 fact 6
+  print $ take 5 example1
+  print $ take 5 example2
+  print $ take 10 (L.runST example3)
+  print $ take 10 (L.runST example4)
+  print $ take 5 (L.runST example5)

tests/LazyST.ref

@@ -0,0 +1,22 @@
+120
+720
+[0,1,2,3,4]
+[0,1,2,3,4]
+[0,1,2,3,4,5,6,7,8,9]
+[0,1,2,3,4,5,6,7,8,9]
+[<A0>
+<init end0>
+<init begin0>
+0,<A1>
+<C0>
+<B0>
+1,<A2>
+<C1>
+<B1>
+2,<A3>
+<C2>
+<B2>
+3,<A4>
+<C3>
+<B3>
+4]

tests/Makefile

@@ -119,6 +119,7 @@ test:
 	$(TMHS) Logarithms && $(EVAL) > Logarithms.out && diff Logarithms.ref Logarithms.out
 	$(TMHS) Builder    && $(EVAL) > Builder.out    && diff Builder.ref Builder.out
 	$(TMHS) Unboxed    && $(EVAL) > Unboxed.out    && diff Unboxed.ref Unboxed.out
+	$(TMHS) LazyST     && $(EVAL) &> LazyST.out    && diff LazyST.ref LazyST.out
 
 testforimp:
 	MHSDIR=.. $(TMHS) ForeignC -of.exe && ./f.exe > ForeignC.out && diff ForeignC.ref ForeignC.out