Harden fixST

Trac #15349 reveals that lazy blackholing can cause trouble for
`fixST` much like it can for `fixIO`. Make `fixST` work just
like `fixIO`.

Reviewers: simonmar, hvr, bgamari

Reviewed By: simonmar

Subscribers: rwbarton, thomie, carter

GHC Trac Issues: #15349

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

2 files changed, 54 insertions, 4 deletions

diff --git a/libraries/base/Control/Monad/Fix.hs b/libraries/base/Control/Monad/Fix.hs
index bb269847b6..a58e2828f3 100644
--- a/libraries/base/Control/Monad/Fix.hs
+++ b/libraries/base/Control/Monad/Fix.hs
@@ -33,7 +33,7 @@ import Data.Ord ( Down(..) )
 import GHC.Base ( Monad, NonEmpty(..), errorWithoutStackTrace, (.) )
 import GHC.Generics
 import GHC.List ( head, tail )
-import GHC.ST
+import Control.Monad.ST.Imp
 import System.IO
 
 -- | Monads having fixed points with a \'knot-tying\' semantics.
diff --git a/libraries/base/Control/Monad/ST/Imp.hs b/libraries/base/Control/Monad/ST/Imp.hs
index c053dcc64d..4d6b12c119 100644
--- a/libraries/base/Control/Monad/ST/Imp.hs
+++ b/libraries/base/Control/Monad/ST/Imp.hs
@@ -1,3 +1,4 @@
+{-# LANGUAGE NoImplicitPrelude #-}
 {-# LANGUAGE Unsafe #-}
 {-# OPTIONS_HADDOCK hide #-}
 
@@ -34,7 +35,56 @@ module Control.Monad.ST.Imp (
         unsafeSTToIO
     ) where
 
-import GHC.ST           ( ST, runST, fixST, unsafeInterleaveST
+import GHC.ST           ( ST, runST, unsafeInterleaveST
                         , unsafeDupableInterleaveST )
-import GHC.Base         ( RealWorld )
-import GHC.IO           ( stToIO, unsafeIOToST, unsafeSTToIO )
+import GHC.Base         ( RealWorld, ($), return )
+import GHC.IO           ( stToIO, unsafeIOToST, unsafeSTToIO
+                        , unsafeDupableInterleaveIO )
+import GHC.MVar         ( readMVar, putMVar, newEmptyMVar )
+import Control.Exception.Base
+                        ( catch, throwIO, NonTermination (..)
+                        , BlockedIndefinitelyOnMVar (..) )
+
+-- | Allow the result of a state transformer computation to be used (lazily)
+-- inside the computation.
+--
+-- Note that if @f@ is strict, @'fixST' f = _|_@.
+fixST :: (a -> ST s a) -> ST s a
+-- See Note [fixST]
+fixST k = unsafeIOToST $ do
+    m <- newEmptyMVar
+    ans <- unsafeDupableInterleaveIO
+             (readMVar m `catch` \BlockedIndefinitelyOnMVar ->
+                                    throwIO NonTermination)
+    result <- unsafeSTToIO (k ans)
+    putMVar m result
+    return result
+
+{- Note [fixST]
+   ~~~~~~~~~~~~
+
+For many years, we implemented fixST much like a pure fixpoint,
+using liftST:
+
+  fixST :: (a -> ST s a) -> ST s a
+  fixST k = ST $ \ s ->
+      let ans       = liftST (k r) s
+          STret _ r = ans
+      in
+      case ans of STret s' x -> (# s', x #)
+
+We knew that lazy blackholing could cause the computation to be re-run if the
+result was demanded strictly, but we thought that would be okay in the case of
+ST. However, that is not the case (see Trac #15349). Notably, the first time
+the computation is executed, it may mutate variables that cause it to behave
+*differently* the second time it's run. That may allow it to terminate when it
+should not. More frighteningly, Arseniy Alekseyev produced a somewhat contrived
+example ( https://mail.haskell.org/pipermail/libraries/2018-July/028889.html )
+demonstrating that it can break reasonable assumptions in "trustworthy" code,
+causing a memory safety violation. So now we implement fixST much like we do
+fixIO. See also the implementation notes for fixIO. Simon Marlow wondered
+whether we could get away with an IORef instead of an MVar. I believe we
+cannot. The function passed to fixST may spark a parallel computation that
+demands the final result. Such a computation should block until the final
+result is available.
+-}