将手动 realWorld# 状态与任意 Monad 交错传递是否安全

Question

考虑这个为任意 Monad:

生成列表的函数

generateListM :: Monad m => Int -> (Int -> m a) -> m [a]
generateListM sz f = go 0
  where go i | i < sz = do x <- f i
                           xs <- go (i + 1)
                           return (x:xs)
             | otherwise = pure []

实现可能并不完美，但这里只是为了演示所需的效果，这是非常简单的。例如，如果 monad 是一个列表，我们将获取列表列表：

λ> generateListM 3 (\i -> [0 :: Int64 .. fromIntegral i])
[[0,0,0],[0,0,1],[0,0,2],[0,1,0],[0,1,1],[0,1,2]]

我想做的是实现相同的效果，但 ByteArray 而不是列表。事实证明，这比我第一次偶然发现这个问题时想象的要棘手得多。最终目标是使用该生成器在 massiv 中实现 mapM，但这不是重点。

最简单的方法是使用 vector 包中的函数 generateM，同时进行一些手动转换。但事实证明，有一种方法可以通过手动处理状态令牌并将其与 monad 交错的巧妙小技巧，至少实现 x2 的性能增益：

{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples       #-}
import           Data.Primitive.ByteArray
import           Data.Primitive.Types
import qualified Data.Vector.Primitive    as VP
import           GHC.Int
import           GHC.Magic
import           GHC.Prim

-- | Can't `return` unlifted types, so we need a wrapper for the state and MutableByteArray
data MutableByteArrayState s = MutableByteArrayState !(State# s) !(MutableByteArray# s)

generatePrimM :: forall m a . (Prim a, Monad m) => Int -> (Int -> m a) -> m (VP.Vector a)
generatePrimM (I# sz#) f =
  runRW# $ \s0# -> do
    let go i# = do
          case i# <# sz# of
            0# ->
              case newByteArray# (sz# *# sizeOf# (undefined :: a)) (noDuplicate# s0#) of
                (# s1#, mba# #) -> return (MutableByteArrayState s1# mba#)
            _ -> do
              res <- f (I# i#)
              MutableByteArrayState si# mba# <- go (i# +# 1#)
              return (MutableByteArrayState (writeByteArray# mba# i# res si#) mba#)
    MutableByteArrayState s# mba# <- go 0#
    case unsafeFreezeByteArray# mba# s# of
      (# _, ba# #) -> return (VP.Vector 0 (I# sz#) (ByteArray ba#))

我们可以像以前一样使用它，除了现在我们将得到一个原始 Vector，它由 ByteArray 支持，这正是我真正需要的：

λ> generatePrimM 3 (\i -> [0 :: Int64 .. fromIntegral i])
[[0,0,0],[0,0,1],[0,0,2],[0,1,0],[0,1,1],[0,1,2]]

这似乎工作得很好，对于 ghc 8.0 和 8.2 版本表现良好，除了在 8.4 和 8.6 中有回归，但这个问题是正交的。

终于进入正题了。这种方法真的安全吗？是否有一些我不知道的极端情况可能会在以后咬我？关于上述功能，欢迎任何其他建议或意见。

PS。 m 不必局限于 Monad，Applicative 也可以，但使用 do 语法时示例会更清晰一些。

Answer 1

TLDR; 从我目前收集到的信息来看，以我最初提出的方式生成基元 Vector 似乎确实是一种安全的方法。此外，使用 noDuplicate# 并不是真正必要的，因为所有操作都是幂等的，并且操作顺序不会对结果数组产生影响。

披露：从我第一次想到这个问题到现在已经一年多了。直到上个月，我才试图回到它。我之所以这么说是因为查看 primitive package now I noticed a new module Data.Primitive.PrimArray to me. As @chi mentioned in the comments, there isn't really a need to drop down to the low-level primitives in order to get a solution, since it might already exist. Which contains exactly the function generatePrimArrayA，这正是我要找的（源代码的简化副本）：

newtype STA a = STA {_runSTA :: forall s. MutableByteArray# s -> ST s (PrimArray a)}

runSTA :: forall a. Prim a => Int -> STA a -> PrimArray a
runSTA !sz =
  \(STA m) -> runST $ newPrimArray sz >>= \(ar :: MutablePrimArray s a) -> m (unMutablePrimArray ar)

generatePrimArrayA :: (Applicative f, Prim a) => Int -> (Int -> f a) -> f (PrimArray a)
generatePrimArrayA len f =
  let go !i
        | i == len = pure $ STA $ \mary -> unsafeFreezePrimArray (MutablePrimArray mary)
        | otherwise =
          liftA2
            (\b (STA m) -> STA $ \mary -> writePrimArray (MutablePrimArray mary) i b >> m mary)
            (f i)
            (go (i + 1))
   in runSTA len <$> go 0

作为一个有趣的练习，如果我们使用通常的归约规则进行基本简化，我们得到的结果与我最初的结果非常相似：

generatePrimArrayA :: forall f a. (Applicative f, Prim a) => Int -> (Int -> f a) -> f (PrimArray a)
generatePrimArrayA !(I# n#) f =
  let go i# = case i# <# n# of
                0# -> pure $ \mary s# ->
                        case unsafeFreezeByteArray# mary s# of
                          (# s'#, arr'# #) -> (# s'#, PrimArray arr'# #)
                _ -> liftA2
                     (\b m ->
                        \mary s ->
                          case writeByteArray# mary i# b s of
                            s'# -> m mary s'#)
                     (f (I# i#))
                     (go (i# +# 1#))
   in (\m -> runRW# $ \s0# ->
                case newByteArray# (n# *# sizeOf# (undefined :: a)) s0# of
                  (# s'#, arr# #) -> case m arr# s'# of
                                       (# _, a #) -> a)
      <$> go 0#

这是我针对 Applicative 而不是 Monad 调整后的版本：

generatePrimM :: forall m a . (Prim a, Applicative m) => Int -> (Int -> m a) -> m (PrimArray a)
generatePrimM (I# sz#) f =
  let go i# = case i# <# sz# of
                0# -> runRW# $ \s0# ->
                      case newByteArray# (sz# *# sizeOf# (undefined :: a)) s0# of
                        (# s1#, mba# #) -> pure (MutableByteArrayState s1# mba#)
                _  -> liftA2
                      (\b (MutableByteArrayState si# mba#) ->
                         MutableByteArrayState (writeByteArray# mba# i# b si#) mba#)
                      (f (I# i#))
                      (go (i# +# 1#))
   in (\(MutableByteArrayState s# mba#) ->
         case unsafeFreezeByteArray# mba# s# of
           (# _, ba# #) -> PrimArray ba#) <$>
      (go 0#)

在功能和性能方面，它们彼此非常接近，最终它们都会产生完全相同的答案。不同之处在于内部循环 go 最后产生的结果。后者将 return 一个包含可以构造 MutableByteArray#s 的闭包的应用程序，稍后将被冻结。前者有一个循环 returns 一个包含将创建冻结 ByteArray#s 的操作的应用程序，一旦提供了可以创建 MutableByteArray# 的操作。

尽管如此，使这两种方法都安全的原因是循环中每个生成的数组的每个元素都只写入一次，并且每个 MutableByteArray# 创建确实在获取 return 之前被冻结由生成函数编辑，但在完成写入之前不会。

将手动 realWorld# 状态与任意 Monad 交错传递是否安全

Is it safe to interleave manual realWorld# state passing with an arbitrary Monad

arrays

haskell

vector

primitive

massiv