[Haskell-cafe] Question on optimization

[Viktor Dukhovni]

On Wed, Dec 16, 2020 at 05:04:53PM -0800, Todd Wilson wrote:

> Thanks, Viktor, for your response, but I think you misunderstood my question:

Indeed I read much too little into it.  Sorry about that.

> Here is a similar, but more sophisticated example involving flattening
> of nested lists, coded three ways:
> 
> > data Nest a = Nil | Cons a (Nest a) | Nest (Nest a) (Nest a)
> >
> > flatten :: Nest a -> [a]
> > flatten Nil = []
> > flatten (Cons x n) = x : flatten n
> > flatten (Nest n1 n2) = flatten n1 ++ flatten n2

I would have gone with a right fold:

    flatten = foldr k []
      where
        k Nil a = a
        k (Cons x n) a = x : foldr k a n
        k (Nest x y) a = foldr k (foldr k a y) x

In fact, I'd be rather inclined to do this indirectly by instead
defining an instance of Foldable for this structure, and then just
using that to do the flattening:

    instance Foldable Nest where
        foldr _ z Nil = z
        foldr f z (Cons x n) = f x (foldr f z n)
        foldr f z (Nest x y) = foldr f (foldr f z y) x

    flatten = foldr (:) []

> Again, I wish I had some general principles about what kind of
> compiler optimizations are applied, so that I could answer such
> questions in the abstract and write code in an easier-to-understand
> but seemingly more expensive way knowing that the extra expense was
> going to be compiled away.

I don't have an answer for your original question, but I think
you'll find that flattening foldable/traversable data structures
is often optimally done with some sort of right fold, which avoids
the costs of (++), by producing the first element directly, and
then lazily the rest of the structure.

The Traversable instance in this case would be:

    {-# LANGUAGE DerivingFunctor, StandaloneDeriving #-}
    data Nest a = Nil
                | Cons a (Nest a)
                | Nest (Nest a) (Nest a)
        deriving (Show, Functor)

    instance Traversable Nest where
        traverse f Nil = pure Nil
        traverse f (Cons x n) = Cons <$> f x <*> traverse f n
        traverse f (Nest x y) = Nest <$> traverse f x <*> traverse f y

A complete demo program is below my signature, which produces the
expected output:

    [1,2,3,10,4,20]
    Just (Cons 2 (Cons 3 (Nest (Cons 4 (Cons 11 Nil)) (Cons 5 (Cons 21 Nil)))))

I don't know whether indirection via "foldr = go" with `go`
recursive and `foldr` "INLINABLE" and ditto for `traverse` would
create further opportunities for optimisation when the instances
are in a separate module.  Someone else might clear that up, or
you could benchmark and see.  IIRC, recursive functions don't
directly inline without such tweaks, but the tweaks might not
help much.

-- 
    Viktor.

{-# LANGUAGE DeriveFunctor #-}

module Main (main) where

data Nest a = Nil
            | Cons a (Nest a)
            | Nest (Nest a) (Nest a)
    deriving (Show, Functor)

instance Foldable Nest where
    foldr _ z Nil = z
    foldr f z (Cons x n) = f x (foldr f z n)
    foldr f z (Nest x y) = foldr f (foldr f z y) x

instance Traversable Nest where
    traverse f Nil = pure Nil
    traverse f (Cons x n) = Cons <$> f x <*> traverse f n
    traverse f (Nest x y) = Nest <$> traverse f x <*> traverse f y

flatten :: Nest a -> [a]
flatten = foldr (:) []

nested :: Nest Int
nested =
    Cons 1
    $ Cons 2
    $ Nest (Cons 3 $ Cons 10 Nil)
           (Cons 4 $ Cons 20 Nil)

main :: IO ()
main = do
    print $ flatten nested
    print $ traverse (\i -> Just (i + 1)) nested