Traverse better

Instead of traversing a list and then converting to an array,
be more direct. GHC seems to be able to optimize this better,
at least in common cases.

Addresses #85.

Author: treeowl

Data/Primitive/Array.hs

@@ -444,10 +444,33 @@ instance Foldable Array where
   {-# INLINE product #-}
 #endif
 
+newtype STA a = STA {_runSTA :: forall s. MutableArray# s a -> ST s (Array a)}
+
+runSTA :: Int -> STA a -> Array a
+runSTA !sz = \ (STA m) -> runST $ newArray_ sz >>= \ ar -> m (marray# ar)
+{-# INLINE runSTA #-}
+
+newArray_ :: Int -> ST s (MutableArray s a)
+newArray_ !n = newArray n badTraverseValue
+
+badTraverseValue :: a
+badTraverseValue = die "traverse" "bad indexing"
+{-# NOINLINE badTraverseValue #-}
+
 instance Traversable Array where
-  traverse f a =
-    fromListN (sizeofArray a)
-      <$> traverse (f . indexArray a) [0 .. sizeofArray a - 1]
+  traverse f = \ !ary ->
+    let
+      !len = sizeofArray ary
+      go !i
+        | i == len = pure $ STA $ \mary -> unsafeFreezeArray (MutableArray mary)
+        | (# x #) <- indexArray## ary i
+        = liftA2 (\b (STA m) -> STA $ \mary ->
+                    writeArray (MutableArray mary) i b >> m mary)
+                 (f x) (go (i + 1))
+    in if len == 0
+       then pure emptyArray
+       else runSTA len <$> go 0
+  {-# INLINE traverse #-}
 
 #if MIN_VERSION_base(4,7,0)
 instance Exts.IsList (Array a) where

Data/Primitive/SmallArray.hs

@@ -525,9 +525,35 @@ instance Foldable SmallArray where
   {-# INLINE product #-}
 #endif
 
+newtype STA a = STA {_runSTA :: forall s. SmallMutableArray# s a -> ST s (SmallArray a)}
+
+runSTA :: Int -> STA a -> SmallArray a
+runSTA !sz = \ (STA m) -> runST $ newSmallArray_ sz >>=
+                        \ (SmallMutableArray ar#) -> m ar#
+{-# INLINE runSTA #-}
+
+newSmallArray_ :: Int -> ST s (SmallMutableArray s a)
+newSmallArray_ !n = newSmallArray n badTraverseValue
+
+badTraverseValue :: a
+badTraverseValue = die "traverse" "bad indexing"
+{-# NOINLINE badTraverseValue #-}
+
 instance Traversable SmallArray where
-  traverse f sa = fromListN l <$> traverse (f . indexSmallArray sa) [0..l-1]
-   where l = length sa
+  traverse f = \ !ary ->
+    let
+      !len = sizeofSmallArray ary
+      go !i
+        | i == len
+        = pure $ STA $ \mary -> unsafeFreezeSmallArray (SmallMutableArray mary)
+        | (# x #) <- indexSmallArray## ary i
+        = liftA2 (\b (STA m) -> STA $ \mary ->
+                    writeSmallArray (SmallMutableArray mary) i b >> m mary)
+                 (f x) (go (i + 1))
+    in if len == 0
+       then pure emptySmallArray
+       else runSTA len <$> go 0
+  {-# INLINE traverse #-}
 
 instance Functor SmallArray where
   fmap f sa = createSmallArray (length sa) (die "fmap" "impossible") $ \smb ->