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

6 files changed, 75 insertions, 17 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.
+-}
diff --git a/libraries/base/GHC/ST.hs b/libraries/base/GHC/ST.hs
index e9d79d9220..9a174383f8 100644
--- a/libraries/base/GHC/ST.hs
+++ b/libraries/base/GHC/ST.hs
@@ -18,7 +18,7 @@
 
 module GHC.ST (
         ST(..), STret(..), STRep,
-        fixST, runST,
+        runST,
 
         -- * Unsafe functions
         liftST, unsafeInterleaveST, unsafeDupableInterleaveST
@@ -92,8 +92,7 @@ instance Monoid a => Monoid (ST s a) where
 
 data STret s a = STret (State# s) a
 
--- liftST is useful when we want a lifted result from an ST computation.  See
--- fixST below.
+-- liftST is useful when we want a lifted result from an ST computation.
 liftST :: ST s a -> State# s -> STret s a
 liftST (ST m) = \s -> case m s of (# s', r #) -> STret s' r
 
@@ -126,16 +125,6 @@ unsafeDupableInterleaveST (ST m) = ST ( \ s ->
     (# s, r #)
   )
 
--- | 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
-fixST k = ST $ \ s ->
-    let ans       = liftST (k r) s
-        STret _ r = ans
-    in
-    case ans of STret s' x -> (# s', x #)
-
 -- | @since 2.01
 instance  Show (ST s a)  where
     showsPrec _ _  = showString "<<ST action>>"
diff --git a/libraries/base/tests/T15349.hs b/libraries/base/tests/T15349.hs
new file mode 100644
index 0000000000..6674330924
--- /dev/null
+++ b/libraries/base/tests/T15349.hs
@@ -0,0 +1,17 @@
+import Control.Monad.ST.Strict
+import Control.Monad.Fix
+import Data.STRef
+
+foo :: ST s Int
+foo = do
+  ref <- newSTRef True
+  mfix $ \res -> do
+    x <- readSTRef ref
+    if x
+      then do
+        writeSTRef ref False
+        return $! (res + 5)
+      else return 10
+
+main :: IO ()
+main = print $ runST foo
diff --git a/libraries/base/tests/T15349.stderr b/libraries/base/tests/T15349.stderr
new file mode 100644
index 0000000000..9cb080d93e
--- /dev/null
+++ b/libraries/base/tests/T15349.stderr
@@ -0,0 +1 @@
+T15349: <<loop>>
diff --git a/libraries/base/tests/all.T b/libraries/base/tests/all.T
index 3d3ebbcd0d..715d4c3f53 100644
--- a/libraries/base/tests/all.T
+++ b/libraries/base/tests/all.T
@@ -241,3 +241,4 @@ test('T14425', normal, compile_and_run, [''])
 test('T10412', normal, compile_and_run, [''])
 test('T13896', normal, compile_and_run, [''])
 test('T13167', normal, compile_and_run, [''])
+test('T15349', [exit_code(1)], compile_and_run, [''])