Make unsafeInterleaveST less unsafe

* Make `unsafeInterleaveST` use `noDuplicate#` like
`unsafeInterleaveIO` does to prevent the suspended action from
being run in two threads.

* In order to accomplish this without `unsafeCoerce#`, generalize
the type of `noDuplicate#`.

* Add `unsafeDupableInterleaveST` to get the old behavior.

* Document unsafe `ST` functions and clean up some related
documentation.

Fixes #13457

Reviewers: austin, hvr, bgamari, ekmett

Reviewed By: bgamari

Subscribers: rwbarton, thomie

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

5 files changed, 56 insertions, 10 deletions

diff --git a/libraries/base/Control/Monad/ST/Imp.hs b/libraries/base/Control/Monad/ST/Imp.hs
index 984970fc72..c053dcc64d 100644
--- a/libraries/base/Control/Monad/ST/Imp.hs
+++ b/libraries/base/Control/Monad/ST/Imp.hs
@@ -29,10 +29,12 @@ module Control.Monad.ST.Imp (
 
         -- * Unsafe operations
         unsafeInterleaveST,
+        unsafeDupableInterleaveST,
         unsafeIOToST,
         unsafeSTToIO
     ) where
 
-import GHC.ST           ( ST, runST, fixST, unsafeInterleaveST )
+import GHC.ST           ( ST, runST, fixST, unsafeInterleaveST
+                        , unsafeDupableInterleaveST )
 import GHC.Base         ( RealWorld )
 import GHC.IO           ( stToIO, unsafeIOToST, unsafeSTToIO )
diff --git a/libraries/base/Control/Monad/ST/Unsafe.hs b/libraries/base/Control/Monad/ST/Unsafe.hs
index 9fa4b739b1..b8560b1cfd 100644
--- a/libraries/base/Control/Monad/ST/Unsafe.hs
+++ b/libraries/base/Control/Monad/ST/Unsafe.hs
@@ -21,6 +21,7 @@
 module Control.Monad.ST.Unsafe (
         -- * Unsafe operations
         unsafeInterleaveST,
+        unsafeDupableInterleaveST,
         unsafeIOToST,
         unsafeSTToIO
     ) where
diff --git a/libraries/base/GHC/IO.hs b/libraries/base/GHC/IO.hs
index 8459db6b75..63b47ff738 100644
--- a/libraries/base/GHC/IO.hs
+++ b/libraries/base/GHC/IO.hs
@@ -84,22 +84,31 @@ failIO s = IO (raiseIO# (toException (userError s)))
 -- ---------------------------------------------------------------------------
 -- Coercions between IO and ST
 
--- | A monad transformer embedding strict state transformers in the 'IO'
--- monad.  The 'RealWorld' parameter indicates that the internal state
+-- | Embed a strict state transformer in an 'IO'
+-- action.  The 'RealWorld' parameter indicates that the internal state
 -- used by the 'ST' computation is a special one supplied by the 'IO'
 -- monad, and thus distinct from those used by invocations of 'runST'.
 stToIO        :: ST RealWorld a -> IO a
 stToIO (ST m) = IO m
 
+-- | Convert an 'IO' action into an 'ST' action. The type of the result
+-- is constrained to use a 'RealWorld' state, and therefore the result cannot
+-- be passed to 'runST'.
 ioToST        :: IO a -> ST RealWorld a
 ioToST (IO m) = (ST m)
 
--- This relies on IO and ST having the same representation modulo the
--- constraint on the type of the state
---
+-- | Convert an 'IO' action to an 'ST' action.
+-- This relies on 'IO' and 'ST' having the same representation modulo the
+-- constraint on the type of the state.
 unsafeIOToST        :: IO a -> ST s a
 unsafeIOToST (IO io) = ST $ \ s -> (unsafeCoerce# io) s
 
+-- | Convert an 'ST' action to an 'IO' action.
+-- This relies on 'IO' and 'ST' having the same representation modulo the
+-- constraint on the type of the state.
+--
+-- For an example demonstrating why this is unsafe, see
+-- https://mail.haskell.org/pipermail/haskell-cafe/2009-April/060719.html
 unsafeSTToIO :: ST s a -> IO a
 unsafeSTToIO (ST m) = IO (unsafeCoerce# m)
 
diff --git a/libraries/base/GHC/IO/Unsafe.hs b/libraries/base/GHC/IO/Unsafe.hs
index 752353515e..c1c07ae2df 100644
--- a/libraries/base/GHC/IO/Unsafe.hs
+++ b/libraries/base/GHC/IO/Unsafe.hs
@@ -104,7 +104,7 @@ unsafeDupablePerformIO  :: IO a -> a
 unsafeDupablePerformIO (IO m) = case runRW# m of (# _, a #) -> a
 
 {-|
-'unsafeInterleaveIO' allows 'IO' computation to be deferred lazily.
+'unsafeInterleaveIO' allows an 'IO' computation to be deferred lazily.
 When passed a value of type @IO a@, the 'IO' will only be performed
 when the value of the @a@ is demanded.  This is used to implement lazy
 file reading, see 'System.IO.hGetContents'.
@@ -113,6 +113,9 @@ file reading, see 'System.IO.hGetContents'.
 unsafeInterleaveIO :: IO a -> IO a
 unsafeInterleaveIO m = unsafeDupableInterleaveIO (noDuplicate >> m)
 
+-- Note [unsafeDupableInterleaveIO should not be inlined]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+--
 -- We used to believe that INLINE on unsafeInterleaveIO was safe,
 -- because the state from this IO thread is passed explicitly to the
 -- interleaved IO, so it cannot be floated out and shared.
@@ -131,7 +134,18 @@ unsafeInterleaveIO m = unsafeDupableInterleaveIO (noDuplicate >> m)
 -- share and sometimes not (plus it probably breaks the noDuplicate).
 -- So now, we do not inline unsafeDupableInterleaveIO.
 
+{-|
+'unsafeDupableInterleaveIO' allows an 'IO' computation to be deferred lazily.
+When passed a value of type @IO a@, the 'IO' will only be performed
+when the value of the @a@ is demanded.
+
+The computation may be performed multiple times by different threads,
+possibly at the same time. To ensure that the computation is performed
+only once, use 'unsafeInterleaveIO' instead.
+-}
+
 {-# NOINLINE unsafeDupableInterleaveIO #-}
+-- See Note [unsafeDupableInterleaveIO should not be inlined]
 unsafeDupableInterleaveIO :: IO a -> IO a
 unsafeDupableInterleaveIO (IO m)
   = IO ( \ s -> let
diff --git a/libraries/base/GHC/ST.hs b/libraries/base/GHC/ST.hs
index 7982d598af..4e00c0e85f 100644
--- a/libraries/base/GHC/ST.hs
+++ b/libraries/base/GHC/ST.hs
@@ -21,7 +21,7 @@ module GHC.ST (
         fixST, runST,
 
         -- * Unsafe functions
-        liftST, unsafeInterleaveST
+        liftST, unsafeInterleaveST, unsafeDupableInterleaveST
     ) where
 
 import GHC.Base
@@ -84,9 +84,29 @@ data STret s a = STret (State# s) a
 liftST :: ST s a -> State# s -> STret s a
 liftST (ST m) = \s -> case m s of (# s', r #) -> STret s' r
 
-{-# NOINLINE unsafeInterleaveST #-}
+noDuplicateST :: ST s ()
+noDuplicateST = ST $ \s -> (# noDuplicate# s, () #)
+
+-- | 'unsafeInterleaveST' allows an 'ST' computation to be deferred
+-- lazily.  When passed a value of type @ST a@, the 'ST' computation will
+-- only be performed when the value of the @a@ is demanded.
+{-# INLINE unsafeInterleaveST #-}
 unsafeInterleaveST :: ST s a -> ST s a
-unsafeInterleaveST (ST m) = ST ( \ s ->
+unsafeInterleaveST m = unsafeDupableInterleaveST (noDuplicateST >> m)
+
+-- | 'unsafeDupableInterleaveST' allows an 'ST' computation to be deferred
+-- lazily.  When passed a value of type @ST a@, the 'ST' computation will
+-- only be performed when the value of the @a@ is demanded.
+--
+-- The computation may be performed multiple times by different threads,
+-- possibly at the same time. To prevent this, use 'unsafeInterleaveST' instead.
+--
+-- @since 4.11
+{-# NOINLINE unsafeDupableInterleaveST #-}
+-- See Note [unsafeDupableInterleaveIO should not be inlined]
+-- in GHC.IO.Unsafe
+unsafeDupableInterleaveST :: ST s a -> ST s a
+unsafeDupableInterleaveST (ST m) = ST ( \ s ->
     let
         r = case m s of (# _, res #) -> res
     in