Make some things more eager

* Perform array indexing eagerly in general to avoid
  useless thunks.

* Make `munzip` stricter for `Array`, to match the
  `SmallArray` instance and avoid loads of thunks.

* Give `Array` and `SmallArray` much less inefficient
  `MonadFix` instances. Leaning on the instance for `[]`
  is bad because indexing into lists is expensive, and that's
  effectively what the `MonadFix` instance does.

Author: treeowl

Data/Primitive/Array.hs

@@ -42,7 +42,7 @@ import Data.Primitive.Internal.Compat ( isTrue#, mkNoRepType )
 import Control.Monad.ST(ST,runST)
 
 import Control.Applicative
-import Control.Monad (MonadPlus(..))
+import Control.Monad (MonadPlus(..), when)
 import Control.Monad.Fix
 #if MIN_VERSION_base(4,4,0)
 import Control.Monad.Zip
@@ -339,7 +339,9 @@ die fun problem = error $ "Data.Primitive.Array." ++ fun ++ ": " ++ problem
 instance Eq a => Eq (Array a) where
   a1 == a2 = sizeofArray a1 == sizeofArray a2 && loop (sizeofArray a1 - 1)
    where loop i | i < 0     = True
-                | otherwise = indexArray a1 i == indexArray a2 i && loop (i-1)
+                | (# x1 #) <- indexArray## a1 i
+                , (# x2 #) <- indexArray## a2 i
+                = x1 == x2 && loop (i-1)
 
 instance Eq (MutableArray s a) where
   ma1 == ma2 = isTrue# (sameMutableArray# (marray# ma1) (marray# ma2))
@@ -349,7 +351,10 @@ instance Ord a => Ord (Array a) where
    where
    mn = sizeofArray a1 `min` sizeofArray a2
    loop i
-     | i < mn    = compare (indexArray a1 i) (indexArray a2 i) `mappend` loop (i+1)
+     | i < mn
+     , (# x1 #) <- indexArray## a1 i
+     , (# x2 #) <- indexArray## a2 i
+     = compare x1 x2 `mappend` loop (i+1)
      | otherwise = compare (sizeofArray a1) (sizeofArray a2)
 
 instance Foldable Array where
@@ -474,9 +479,11 @@ fromList l = fromListN (length l) l
 instance Functor Array where
   fmap f a =
     createArray (sizeofArray a) (die "fmap" "impossible") $ \mb ->
-      let go i | i == sizeofArray a = return ()
-               | otherwise          = writeArray mb i (f $ indexArray a i)
-                                      >> go (i+1)
+      let go i | i == sizeofArray a
+               = return ()
+               | otherwise
+               = do x <- indexArrayM a i
+                    writeArray mb i (f x) >> go (i+1)
        in go 0
 #if MIN_VERSION_base(4,8,0)
   e <$ a = runST $ newArray (sizeofArray a) e >>= unsafeFreezeArray
@@ -486,12 +493,15 @@ instance Applicative Array where
   pure x = runST $ newArray 1 x >>= unsafeFreezeArray
   ab <*> a = runST $ do
     mb <- newArray (szab*sza) $ die "<*>" "impossible"
-    let go1 i
-          | i < szab  = go2 (i*sza) (indexArray ab i) 0 >> go1 (i+1)
-          | otherwise = return ()
-        go2 off f j
-          | j < sza   = writeArray mb (off + j) (f $ indexArray a j)
-          | otherwise = return ()
+    let go1 i = when (i < szab) $
+            do
+              f <- indexArrayM ab i
+              go2 (i*sza) f 0
+              go1 (i+1)
+        go2 off f j = when (j < sza) $
+            do
+              x <- indexArrayM a j
+              writeArray mb (off + j) (f x)
     go1 0
     unsafeFreezeArray mb
    where szab = sizeofArray ab ; sza = sizeofArray a
@@ -503,7 +513,9 @@ instance Applicative Array where
   a <* b = createArray (sza*szb) (die "<*" "impossible") $ \ma ->
     let fill off i e | i < szb   = writeArray ma (off+i) e >> fill off (i+1) e
                      | otherwise = return ()
-        go i | i < sza   = fill (i*szb) 0 (indexArray a i) >> go (i+1)
+        go i | i < sza
+             = do x <- indexArrayM a i
+                  fill (i*szb) 0 x >> go (i+1)
              | otherwise = return ()
      in go 0
    where sza = sizeofArray a ; szb = sizeofArray b
@@ -518,22 +530,35 @@ instance Alternative Array where
   many a | sizeofArray a == 0 = pure []
          | otherwise = die "many" "infinite arrays are not well defined"
 
+data ArrayStack a
+  = PushArray !(Array a) !(ArrayStack a)
+  | EmptyStack
+-- See the note in SmallArray about how we might improve this.
+
 instance Monad Array where
   return = pure
   (>>) = (*>)
 
-  ary >>= f = collect 0 [] (la-1)
+  ary >>= f = collect 0 EmptyStack (la-1)
    where
    la = sizeofArray ary
    collect sz stk i
      | i < 0 = createArray sz (die ">>=" "impossible") $ fill 0 stk
-     | otherwise = let sb = f $ indexArray ary i in
-         collect (sz + sizeofArray sb) (sb:stk) (i-1)
-
-   fill _   [      ] _   = return ()
-   fill off (sb:sbs) smb =
-     copyArray smb off sb 0 (sizeofArray sb)
-       *> fill (off + sizeofArray sb) sbs smb
+     | (# x #) <- indexArray## ary i
+     , let sb = f x
+           lsb = sizeofArray sb
+       -- If we don't perform this check, we could end up allocating
+       -- a stack full of empty arrays if someone is filtering most
+       -- things out. So we refrain from pushing empty arrays.
+     = if lsb == 0
+       then collect sz stk (i - 1)
+       else collect (sz + lsb) (PushArray sb stk) (i-1)
+
+   fill _   EmptyStack         _   = return ()
+   fill off (PushArray sb sbs) smb
+     | let lsb = sizeofArray sb
+     = copyArray smb off sb 0 (lsb)
+         *> fill (off + lsb) sbs smb
 
   fail _ = empty
 
@@ -543,10 +568,13 @@ instance MonadPlus Array where
 
 zipW :: String -> (a -> b -> c) -> Array a -> Array b -> Array c
 zipW s f aa ab = createArray mn (die s "impossible") $ \mc ->
-  let go i
-        | i < mn    = writeArray mc i (f (indexArray aa i) (indexArray ab i))
-                   >> go (i+1)
-        | otherwise = return ()
+  let go i | i < mn
+           = do
+               x <- indexArrayM aa i
+               y <- indexArrayM ab i
+               writeArray mc i (f x y)
+               go (i+1)
+           | otherwise = return ()
    in go 0
  where mn = sizeofArray aa `min` sizeofArray ab
 {-# INLINE zipW #-}
@@ -560,7 +588,7 @@ instance MonadZip Array where
     ma <- newArray sz (die "munzip" "impossible")
     mb <- newArray sz (die "munzip" "impossible")
     let go i | i < sz = do
-          let (a, b) = indexArray aab i
+          (a, b) <- indexArrayM aab i
           writeArray ma i a
           writeArray mb i b
           go (i+1)
@@ -570,7 +598,14 @@ instance MonadZip Array where
 #endif
 
 instance MonadFix Array where
-  mfix f = let l = mfix (toList . f) in fromListN (length l) l
+  mfix f = createArray (sizeofArray (f err))
+                       (die "mfix" "impossible") $ flip fix 0 $
+    \r !i !mary -> when (i < sz) $ do
+                      writeArray mary i (fix (\xi -> f xi `indexArray` i))
+                      r (i + 1) mary
+    where
+      sz = sizeofArray (f err)
+      err = error "mfix for Data.Primitive.Array applied to strict function."
 
 #if MIN_VERSION_base(4,9,0)
 instance Semigroup (Array a) where

Data/Primitive/SmallArray.hs

@@ -75,7 +75,7 @@ import Control.Monad.ST
 import Control.Monad.Zip
 #endif
 import Data.Data
-import Data.Foldable
+import Data.Foldable as Foldable
 import Data.Functor.Identity
 #if !(MIN_VERSION_base(4,11,0))
 import Data.Monoid
@@ -121,7 +121,7 @@ instance IsList (SmallArray a) where
   type Item (SmallArray a) = a
   fromListN n l = SmallArray (fromListN n l)
   fromList l = SmallArray (fromList l)
-  toList (SmallArray a) = toList a
+  toList a = Foldable.toList a
 #endif
 #endif
 
@@ -419,19 +419,27 @@ instance Eq a => Eq (SmallArray a) where
   sa1 == sa2 = length sa1 == length sa2 && loop (length sa1 - 1)
    where
    loop i
-     | i < 0     = True
-     | otherwise = indexSmallArray sa1 i == indexSmallArray sa2 i && loop (i-1)
+     | i < 0
+     = True
+     | (# x #) <- indexSmallArray## sa1 i
+     , (# y #) <- indexSmallArray## sa2 i
+     = x == y && loop (i-1)
 
 instance Eq (SmallMutableArray s a) where
   SmallMutableArray sma1# == SmallMutableArray sma2# =
     isTrue# (sameSmallMutableArray# sma1# sma2#)
 
 instance Ord a => Ord (SmallArray a) where
-  compare sl sr = fix ? 0 $ \go i ->
-    if i < l
-      then compare (indexSmallArray sl i) (indexSmallArray sr i) <> go (i+1)
-      else compare (length sl) (length sr)
-   where l = length sl `min` length sr
+  compare a1 a2 = loop 0
+   where
+   mn = length a1 `min` length a2
+   loop i
+     | i < mn
+     , (# x1 #) <- indexSmallArray## a1 i
+     , (# x2 #) <- indexSmallArray## a2 i
+     = compare x1 x2 `mappend` loop (i+1)
+     | otherwise = compare (length a1) (length a2)
+
 
 instance Foldable SmallArray where
   -- Note: we perform the array lookups eagerly so we won't
@@ -532,8 +540,9 @@ instance Traversable SmallArray where
 instance Functor SmallArray where
   fmap f sa = createSmallArray (length sa) (die "fmap" "impossible") $ \smb ->
     fix ? 0 $ \go i ->
-      when (i < length sa) $
-        writeSmallArray smb i (f $ indexSmallArray sa i) *> go (i+1)
+      when (i < length sa) $ do
+        x <- indexSmallArrayM sa i
+        writeSmallArray smb i (f x) *> go (i+1)
   {-# INLINE fmap #-}
 
   x <$ sa = createSmallArray (length sa) x noOp
@@ -548,22 +557,23 @@ instance Applicative SmallArray where
    where
    la = length sa ; lb = length sb
 
-  sa <* sb = createSmallArray (la*lb) (indexSmallArray sa $ la-1) $ \sma ->
-    fix ? 0 $ \outer i -> when (i < la-1) $ do
-      let a = indexSmallArray sa i
-      fix ? 0 $ \inner j ->
-        when (j < lb) $
-          writeSmallArray sma (la*i + j) a *> inner (j+1)
-      outer $ i+1
-   where
-   la = length sa ; lb = length sb
+  a <* b = createSmallArray (sza*szb) (die "<*" "impossible") $ \ma ->
+    let fill off i e = when (i < szb) $
+                         writeSmallArray ma (off+i) e >> fill off (i+1) e
+        go i = when (i < sza) $ do
+                 x <- indexSmallArrayM a i
+                 fill (i*szb) 0 x
+                 go (i+1)
+     in go 0
+   where sza = sizeofSmallArray a ; szb = sizeofSmallArray b
 
   sf <*> sx = createSmallArray (lf*lx) (die "<*>" "impossible") $ \smb ->
     fix ? 0 $ \outer i -> when (i < lf) $ do
-      let f = indexSmallArray sf i
+      f <- indexSmallArrayM sf i
       fix ? 0 $ \inner j ->
-        when (j < lx) $
-          writeSmallArray smb (lf*i + j) (f $ indexSmallArray sx j)
+        when (j < lx) $ do
+          x <- indexSmallArrayM sx j
+          writeSmallArray smb (lf*i + j) (f x)
             *> inner (j+1)
       outer $ i+1
    where
@@ -583,20 +593,41 @@ instance Alternative SmallArray where
   some sa | null sa   = emptySmallArray
           | otherwise = die "some" "infinite arrays are not well defined"
 
+data ArrayStack a
+  = PushArray !(SmallArray a) !(ArrayStack a)
+  | EmptyStack
+-- TODO: This isn't terribly efficient. It would be better to wrap
+-- ArrayStack with a type like
+--
+-- data NES s a = NES !Int !(SmallMutableArray s a) !(ArrayStack a)
+--
+-- We'd copy incoming arrays into the mutable array until we would
+-- overflow it. Then we'd freeze it, push it on the stack, and continue.
+-- Any sufficiently large incoming arrays would go straight on the stack.
+-- Such a scheme would make the stack much more compact in the case
+-- of many small arrays.
+
 instance Monad SmallArray where
   return = pure
   (>>) = (*>)
 
-  sa >>= f = collect 0 [] (la-1)
+  sa >>= f = collect 0 EmptyStack (la-1)
    where
    la = length sa
    collect sz stk i
      | i < 0 = createSmallArray sz (die ">>=" "impossible") $ fill 0 stk
-     | otherwise = let sb = f $ indexSmallArray sa i in
-         collect (sz + length sb) (sb:stk) (i-1)
-
-   fill _   [      ] _   = return ()
-   fill off (sb:sbs) smb =
+     | (# x #) <- indexSmallArray## sa i
+     , let sb = f x
+           lsb = length sb
+       -- If we don't perform this check, we could end up allocating
+       -- a stack full of empty arrays if someone is filtering most
+       -- things out. So we refrain from pushing empty arrays.
+     = if lsb == 0
+       then collect sz stk (i-1)
+       else collect (sz + lsb) (PushArray sb stk) (i-1)
+
+   fill _ EmptyStack _ = return ()
+   fill off (PushArray sb sbs) smb =
      copySmallArray smb off sb 0 (length sb)
        *> fill (off + length sb) sbs smb
 
@@ -609,9 +640,11 @@ instance MonadPlus SmallArray where
 zipW :: String -> (a -> b -> c) -> SmallArray a -> SmallArray b -> SmallArray c
 zipW nm = \f sa sb -> let mn = length sa `min` length sb in
   createSmallArray mn (die nm "impossible") $ \mc ->
-    fix ? 0 $ \go i -> when (i < mn) $
-      writeSmallArray mc i (f (indexSmallArray sa i) (indexSmallArray sb i))
-        *> go (i+1)
+    fix ? 0 $ \go i -> when (i < mn) $ do
+      x <- indexSmallArrayM sa i
+      y <- indexSmallArrayM sb i
+      writeSmallArray mc i (f x y)
+      go (i+1)
 {-# INLINE zipW #-}
 
 instance MonadZip SmallArray where
@@ -631,7 +664,14 @@ instance MonadZip SmallArray where
         <*> unsafeFreezeSmallArray smb
 
 instance MonadFix SmallArray where
-  mfix f = fromList . mfix $ toList . f
+  mfix f = createSmallArray (sizeofSmallArray (f err))
+                            (die "mfix" "impossible") $ flip fix 0 $
+    \r !i !mary -> when (i < sz) $ do
+                      writeSmallArray mary i (fix (\xi -> f xi `indexSmallArray` i))
+                      r (i + 1) mary
+    where
+      sz = sizeofSmallArray (f err)
+      err = error "mfix for Data.Primitive.SmallArray applied to strict function."
 
 #if MIN_VERSION_base(4,9,0)
 instance Sem.Semigroup (SmallArray a) where
@@ -658,7 +698,7 @@ instance IsList (SmallArray a) where
         [] -> pure ()
         x:xs -> writeSmallArray sma i x *> go (i+1) xs
   fromList l = fromListN (length l) l
-  toList sa = indexSmallArray sa <$> [0 .. length sa - 1]
+  toList = Foldable.toList
 
 instance Show a => Show (SmallArray a) where
   showsPrec p sa = showParen (p > 10) $