Remove the Ix requirement for Traversable (Array i)
Motivation
The Functor
and Foldable
instances for Array
don't have an Ix
requirement on the index i
used in them, but the Traversable
instance does, because it's defined in a different module and uses only the publicly-accessible functions to define it. For the sake of consistency, we should remove the Ix
requirement for Traversable
.
Proposal
We use the same sort of technology that Data.Primitive.Array
does that avoids creating an intermediate list to define an atraverse
function, which will be exported from GHC.Arr
:
newtype STA a = STA { unSTA :: forall s. MutableArray# s a -> State# s -> State# s }
atraverse :: Applicative f => (a -> f b) -> Array i a -> f (Array i b)
atraverse f = \arr@(Array l u n@(I# n#) _) -> let
runSTA (STA sta) = runST $ ST $ \s -> case newArray# n# undefined s of
(# sn, ma# #) -> case sta ma# sn of
sw -> case unsafeFreezeArray# ma# sw of
(# sf, a# #) -> (# sf, Array l u n a# #)
{-# INLINE runSTA #-}
go i | i == n = pure $ STA $ \_ s -> s
go i@(I# i#) = case unsafeAt# arr i of
(# a #) -> let
push b (STA sta) = STA $ \ma# s -> case writeArray# ma# i# b s of
sw -> sta ma# sw
{-# INLINE push #-}
in liftA2 push (f a) (go (i + 1))
{-# INLINE go #-}
in fmap runSTA $ go 0
{-# INLINE [1] atraverse #-}
{-# RULES "atraverse/Identity" atraverse = (coerce :: ((a -> b) -> Array i a -> Array i b) -> (a -> Identity b) -> Array i a -> Identity (Array i b)) amap #-}
{-# RULES "atraverse/(->)" atraverse = \f a -> \e -> amap (flip f e) a #-}
{-# RULES "atraverse/amap" forall f g a. atraverse f (amap g a) = atraverse (f . g) a #-}
We could add some specialty instances for different Applicatives
, like ST
and Maybe
and Either
:
atraverseST :: (a -> ST s b) -> Array i a -> ST s (Array i b)
atraverseST f = \arr@(Array l u n@(I# n#) _) -> ST $ \s0 -> case newArray# n# undefined s0 of
(# sn, ma# #) -> let
go i s | i == n = case unsafeFreezeArray# ma# s of
(# sf, a# #) -> (# sf, Array l u n a# #)
go i@(I# i#) s = case unsafeAt# arr i of
(# a #) -> case f a of
ST mb -> case mb s of
(# sb, b #) -> case writeArray# ma# i# b sb of
sw -> go (i + 1) sw
{-# INLINE go #-}
in go 0 sn
{-# INLINE [1] atraverseST #-}
atraverseIO :: (a -> IO b) -> Array i a -> IO (Array i b)
atraverseIO f = \arr@(Array l u n@(I# n#) _) -> IO $ \s0 -> case newArray# n# undefined s0 of
(# sn, ma# #) -> let
go i s | i == n = case unsafeFreezeArray# ma# s of
(# sf, a# #) -> (# sf, Array l u n a# #)
go i@(I# i#) s = case unsafeAt# arr i of
(# a #) -> case f a of
IO mb -> case mb s of
(# sb, b #) -> case writeArray# ma# i# b sb of
sw -> go (i + 1) sw
{-# INLINE go #-}
in go 0 sn
{-# INLINE [1] atraverseIO #-}
atraverseMaybe :: (a -> Maybe b) -> Array i a -> Maybe (Array i b)
atraverseMaybe f = \arr@(Array l u n@(I# n#) _) -> runST $ ST $ \s0 -> case newArray# n# undefined s0 of
(# sn, ma# #) -> let
go i s | i == n = case unsafeFreezeArray# ma# s of
(# sf, a# #) -> (# sf, Just (Array l u n a#) #)
go i@(I# i#) s = case unsafeAt# arr i of
(# a #) -> case f a of
Nothing -> (# s, Nothing #)
Just b -> case writeArray# ma# i# b s of
sw -> go (i + 1) sw
{-# INLINE go #-}
in go 0 sn
{-# INLINE [1] atraverseMaybe #-}
atraverseEither :: (a -> Either e b) -> Array i a -> Either e (Array i b)
atraverseEither f = \arr@(Array l u n@(I# n#) _) -> runST $ ST $ \s0 -> case newArray# n# undefined s0 of
(# sn, ma# #) -> let
go i s | i == n = case unsafeFreezeArray# ma# s of
(# sf, a# #) -> (# sf, Right (Array l u n a#) #)
go i@(I# i#) s = case unsafeAt# arr i of
(# a #) -> case f a of
Left e -> (# s, Left e #)
Right b -> case writeArray# ma# i# b s of
sw -> go (i + 1) sw
{-# INLINE go #-}
in go 0 sn
{-# INLINE [1] atraverseEither #-}
{-# RULES "atraverse/ST" atraverse = atraverseST #-}
{-# RULES "atraverse/IO" atraverse = atraverseIO #-}
{-# RULES "atraverse/Maybe" atraverse = atraverseMaybe #-}
{-# RULES "atraverse/Either" atraverse = atraverseEither #-}
Even though this request touches the public interface for Data.Array
, removing a constraint shouldn't break any code and shouldn't require the Libraries committee to weigh in.