Attempt to make lazy ST thread safe

Use `noDuplicate#` to prevent lazy `ST` thunks from
being evaluated in multiple GHC threads.

Some lazy `ST` functions added laziness that did not seem to be useful
(e.g.,
creating lazy pairs that will never be matched unless one of their
components
is demanded). Stripped that out.

Reviewers: ekmett, simonpj, bgamari, simonmar, austin, hvr

Subscribers: rwbarton, thomie

Differential Revision: https://phabricator.haskell.org/D3038

GHC Trac Issues: #11760

8 files changed, 306 insertions, 24 deletions

diff --git a/libraries/base/Control/Monad/ST/Lazy/Imp.hs b/libraries/base/Control/Monad/ST/Lazy/Imp.hs
index 414c06c8c3..9883def001 100644
--- a/libraries/base/Control/Monad/ST/Lazy/Imp.hs
+++ b/libraries/base/Control/Monad/ST/Lazy/Imp.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE Unsafe #-}
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE MagicHash, UnboxedTuples, RankNTypes #-}
 {-# OPTIONS_HADDOCK hide #-}
 
@@ -43,6 +44,7 @@ import qualified Control.Monad.ST.Unsafe as ST
 
 import qualified GHC.ST as GHC.ST
 import GHC.Base
+import qualified Control.Monad.Fail as Fail
 
 -- | The lazy state-transformer monad.
 -- A computation of type @'ST' s a@ transforms an internal state indexed
@@ -59,35 +61,127 @@ import GHC.Base
 -- The '>>=' and '>>' operations are not strict in the state.  For example,
 --
 -- @'runST' (writeSTRef _|_ v >>= readSTRef _|_ >> return 2) = 2@
-newtype ST s a = ST (State s -> (a, State s))
+newtype ST s a = ST { unST :: State s -> (a, State s) }
+
+-- A lifted state token. This can be imagined as a moment in the timeline
+-- of a lazy state thread. Forcing the token forces all delayed actions in
+-- the thread up until that moment to be performed.
 data State s = S# (State# s)
 
+{- Note [Lazy ST and multithreading]
+
+We used to imagine that passing a polymorphic state token was all that we
+needed to keep state threads separate (see Launchbury and Peyton Jones, 1994:
+https://www.microsoft.com/en-us/research/publication/lazy-functional-state-threads/).
+But this breaks down in the face of concurrency (see #11760). Whereas a strict
+ST computation runs to completion before producing anything, a value produced
+by running a lazy ST computation may contain a thunk that, when forced, will
+lead to further stateful computations. If such a thunk is entered by more than
+one thread, then they may both read from and write to the same references and
+arrays, interfering with each other. To work around this, any time we lazily
+suspend execution of a lazy ST computation, we bind the result pair to a
+NOINLINE binding (ensuring that it is not duplicated) and calculate that
+pair using (unsafePerformIO . evaluate), ensuring that only one thread will
+enter the thunk. We still use lifted state tokens to actually drive execution,
+so in these cases we effectively deal with *two* state tokens: the lifted
+one we get from the previous computation, and the unlifted one we pull out of
+thin air. -}
+
+{- Note [Lazy ST: not producing lazy pairs]
+
+The fixST and strictToLazyST functions used to construct functions that
+produced lazy pairs. Why don't we need that laziness? The ST type is kept
+abstract, so no one outside this module can ever get their hands on a (result,
+State s) pair. We ourselves never match on such pairs when performing ST
+computations unless we also force one of their components. So no one should be
+able to detect the change. By refraining from producing such thunks (which
+reference delayed ST computations), we avoid having to ask whether we have to
+wrap them up with unsafePerformIO. See Note [Lazy ST and multithreading]. -}
+
+-- | This is a terrible hack to prevent a thunk from being entered twice.
+-- Simon Peyton Jones would very much like to be rid of it.
+noDup :: a -> a
+noDup a = runRW# (\s ->
+  case noDuplicate# s of
+    _ -> a)
+
 -- | @since 2.01
 instance Functor (ST s) where
     fmap f m = ST $ \ s ->
       let
-       ST m_a = m
-       (r,new_s) = m_a s
+        -- See Note [Lazy ST and multithreading]
+        {-# NOINLINE res #-}
+        res = noDup (unST m s)
+        (r,new_s) = res
       in
-      (f r,new_s)
+        (f r,new_s)
+
+    x <$ m = ST $ \ s ->
+      let
+        {-# NOINLINE s' #-}
+        -- See Note [Lazy ST and multithreading]
+        s' = noDup (snd (unST m s))
+      in (x, s')
 
 -- | @since 2.01
 instance Applicative (ST s) where
     pure a = ST $ \ s -> (a,s)
-    (<*>) = ap
+
+    fm <*> xm = ST $ \ s ->
+       let
+         {-# NOINLINE res1 #-}
+         !res1 = unST fm s
+         !(f, s') = res1
+
+         {-# NOINLINE res2 #-}
+         -- See Note [Lazy ST and multithreading]
+         res2 = noDup (unST xm s')
+         (x, s'') = res2
+       in (f x, s'')
+    -- Why can we use a strict binding for res1? If someone
+    -- forces the (f x, s'') pair, then they must need
+    -- f or s''. To get s'', they need s'.
+
+    m *> n = ST $ \s ->
+       let
+         {-# NOINLINE s' #-}
+         -- See Note [Lazy ST and multithreading]
+         s' = noDup (snd (unST m s))
+       in unST n s'
+
+    m <* n = ST $ \s ->
+       let
+         {-# NOINLINE res1 #-}
+         !res1 = unST m s
+         !(mr, s') = res1
+
+         {-# NOINLINE s'' #-}
+         -- See Note [Lazy ST and multithreading]
+         s'' = noDup (snd (unST n s'))
+       in (mr, s'')
+    -- Why can we use a strict binding for res1? The same reason as
+    -- in <*>. If someone demands the (mr, s'') pair, then they will
+    -- force mr or s''. To get s'', they need s'.
 
 -- | @since 2.01
 instance Monad (ST s) where
 
-        fail s   = errorWithoutStackTrace s
+    fail s   = errorWithoutStackTrace s
+
+    (>>) = (*>)
+
+    m >>= k = ST $ \ s ->
+       let
+         -- See Note [Lazy ST and multithreading]
+         {-# NOINLINE res #-}
+         res = noDup (unST m s)
+         (r,new_s) = res
+       in
+         unST (k r) new_s
 
-        (ST m) >>= k
-         = ST $ \ s ->
-           let
-             (r,new_s) = m s
-             ST k_a = k r
-           in
-           k_a new_s
+-- | @since 4.10
+instance Fail.MonadFail (ST s) where
+    fail s = errorWithoutStackTrace s
 
 -- | Return the value computed by a state transformer computation.
 -- The @forall@ ensures that the internal state used by the 'ST'
@@ -101,10 +195,13 @@ runST (ST st) = runRW# (\s -> case st (S# s) of (r, _) -> r)
 fixST :: (a -> ST s a) -> ST s a
 fixST m = ST (\ s -> 
                 let 
-                   ST m_r = m r
-                   (r,s') = m_r s
-                in
-                   (r,s'))
+                   q@(r,_s') = unST (m r) s
+                in q)
+-- Why don't we need unsafePerformIO in fixST? We create a thunk, q,
+-- to perform a lazy state computation, and we pass a reference to that
+-- thunk, r, to m. Uh oh? No, I think it should be fine, because that thunk
+-- itself is demanded directly in the `let` body. See also
+-- Note [Lazy ST: not producing lazy pairs].
 
 -- | @since 2.01
 instance MonadFix (ST s) where
@@ -119,13 +216,10 @@ thread passed to 'strictToLazyST' is not performed until the result of
 the lazy state thread it returns is demanded.
 -}
 strictToLazyST :: ST.ST s a -> ST s a
-strictToLazyST m = ST $ \s ->
-        let 
-           pr = case s of { S# s# -> GHC.ST.liftST m s# }
-           r  = case pr of { GHC.ST.STret _ v -> v }
-           s' = case pr of { GHC.ST.STret s2# _ -> S# s2# }
-        in
-        (r, s')
+strictToLazyST (GHC.ST.ST m) = ST $ \(S# s) ->
+  case m s of
+    (# s', a #) -> (a, S# s')
+-- See Note [Lazy ST: not producing lazy pairs]
 
 {-| 
 Convert a lazy 'ST' computation into a strict one.
diff --git a/libraries/base/tests/T11760.hs b/libraries/base/tests/T11760.hs
new file mode 100644
index 0000000000..875c15916d
--- /dev/null
+++ b/libraries/base/tests/T11760.hs
@@ -0,0 +1,51 @@
+-- Written by Bertram Felgenhauer
+--
+-- https://ghc.haskell.org/trac/ghc/ticket/11760#comment:14
+--
+-- Compile with -threaded -with-rtsopts=-N2
+
+{-# LANGUAGE BangPatterns #-}
+import Control.Concurrent
+import Control.Monad
+import Control.Monad.ST.Lazy
+import Control.Exception
+import Data.STRef
+import Data.IORef
+import Control.Concurrent.MVar
+import Data.List
+
+-- evil ST action that tries to synchronize (by busy waiting on the
+-- shared STRef) with a concurrent evaluation
+evil :: ST s [Int]
+evil = do
+    r <- strictToLazyST $ newSTRef 0
+    replicateM 100 $ do
+        i <- strictToLazyST $ readSTRef r
+        let !j = i + 1
+        strictToLazyST $ writeSTRef r j
+        let go 0 = return ()
+            go n = do
+                i' <- strictToLazyST $ readSTRef r
+                when (j == i') $ go (n-1)
+        go 100
+        return j
+
+main = do
+    let res = runST evil
+    s0 <- newIORef (map pred (0 : res))
+    s1 <- newIORef (map pred (1 : res))
+    m0 <- newMVar ()
+    m1 <- newMVar ()
+    forkIO $ do
+        putMVar m0 ()
+        readIORef s0 >>= evaluate . foldl' (+) 0
+        putMVar m0 ()
+    forkIO $ do
+        putMVar m1 ()
+        readIORef s1 >>= evaluate . foldl' (+) 0
+        putMVar m1 ()
+    threadDelay 10000
+    replicateM 3 $ takeMVar m0 >> takeMVar m1
+    v0 <- tail <$> readIORef s0
+    v1 <- tail <$> readIORef s1
+    print (v0 == v1)
diff --git a/libraries/base/tests/T11760.stdout b/libraries/base/tests/T11760.stdout
new file mode 100644
index 0000000000..0ca95142bb
--- /dev/null
+++ b/libraries/base/tests/T11760.stdout
@@ -0,0 +1 @@
+True
diff --git a/libraries/base/tests/all.T b/libraries/base/tests/all.T
index 3be05afa8b..7ce6a81385 100644
--- a/libraries/base/tests/all.T
+++ b/libraries/base/tests/all.T
@@ -196,3 +196,5 @@ test('T10149', normal, compile_and_run, [''])
 test('T11334a', normal, compile_and_run, [''])
 test('T11555', normal, compile_and_run, [''])
 test('T12852', when(opsys('mingw32'), skip), compile_and_run, [''])
+test('lazySTexamples', normal, compile_and_run, [''])
+test('T11760', normal, compile_and_run, ['-threaded -with-rtsopts=-N2'])
diff --git a/libraries/base/tests/fixUsingLazyST.stdout b/libraries/base/tests/fixUsingLazyST.stdout
new file mode 100644
index 0000000000..740182381e
--- /dev/null
+++ b/libraries/base/tests/fixUsingLazyST.stdout
@@ -0,0 +1,2 @@
+120
+720
diff --git a/libraries/base/tests/lazySTexamples.hs b/libraries/base/tests/lazySTexamples.hs
new file mode 100644
index 0000000000..6319aeab9e
--- /dev/null
+++ b/libraries/base/tests/lazySTexamples.hs
@@ -0,0 +1,110 @@
+import Data.STRef.Lazy
+import Control.Monad.ST.Lazy as L
+import Control.Monad.ST.Strict 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 go where
+  go = 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 main where
+  main = 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)
+  hPutStrLn stderr $ show (take 5 (L.runST example5))
diff --git a/libraries/base/tests/lazySTexamples.stderr b/libraries/base/tests/lazySTexamples.stderr
new file mode 100644
index 0000000000..4cc7747038
--- /dev/null
+++ b/libraries/base/tests/lazySTexamples.stderr
@@ -0,0 +1,16 @@
+[<A0>
+<init end0>
+<init begin0>
+0,<A1>
+<C0>
+<B0>
+1,<A2>
+<C1>
+<B1>
+2,<A3>
+<C2>
+<B2>
+3,<A4>
+<C3>
+<B3>
+4]
diff --git a/libraries/base/tests/lazySTexamples.stdout b/libraries/base/tests/lazySTexamples.stdout
new file mode 100644
index 0000000000..9c6dcc7459
--- /dev/null
+++ b/libraries/base/tests/lazySTexamples.stdout
@@ -0,0 +1,6 @@
+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]