Рассмотрим эту функцию, которая генерирует список для произвольного 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 []
Возможно, реализация не идеальна, но она представлена здесь исключительно для демонстрации желаемого эффекта, что довольно просто.,Например, если монада является списком, получите список списков:
λ> 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
вместо списка.Оказывается, это гораздо сложнее, чем я думал, когда впервые наткнулся на эту проблему.Конечная цель - использовать этот генератор для реализации mapM
в massiv , но это не главное.
Подход, который требует наименьших усилий, заключается в использовании функции generateM
из vector package при выполнении небольшого ручного преобразования.Но, как оказалось, есть способ добиться увеличения производительности по меньшей мере в два раза с помощью этого небольшого хитрого способа обработки маркера состояния вручную и чередования его с монадой:
{-# 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
.