Strictness of foldr' from containers

This blog post will present Haskell's evaluation order based on an interesting issue I discovered in foldr'. As a result of this investigation, the original implementation was altered. Let us start with with some basics on lazy evaluation.

Lazy evaluation

There are two important aspects of lazy evaluation: evaluation order and sharing of suspended computations (thunks). Understanding both is important for writing efficient code or to avoid space leaks. In this blog post we focus on the evaluation order. Haskell is using the so called call-by-need evaluation strategy, which is defined as: always reduce the left most, outer most redex (i.e. unevaluated reducible expression) found in an expression. Let's consider a simple example:

head (map (+1) ([1,2,3] :: [Int])
Here we have two unevaluated sub-expressions: the application of head and the application of map. We cannot reduce the application of head since we don't match any of its patterns:
head :: [a] -> a
head (a : _) = a
head []      = error "Prelude.head: empty list"
We can reduce the application of map, which makes it the only redex in this expression (note that if we didn't specify that the literals are all integers they would be redexes too, desugared to fromInteger 1 :: forall a. Num a => a, see Haskell2010 report, in what follows we write 1 for 1 :: Int for brevity). After applying the definition of repeat we get:
head ((+1) 1 : map (+1) [2, 3])
Now we have three redexes that could be evaluated: either the application of head, the (+1) 1 expression or the application map. According to the call-by-need strategy we reduce the application of head first and, as a result the expression evaluates to (+1) 1, which further can be reduced to its normal form 2 :: Int (i.e. an expression without redexes).

Evaluation to weak head normal form (WHNF), means that evaluation stops as soon as at the top of the expression we get a constructor, or a lambda abstraction.

Bang patterns

The BangPatterns extension allows to influence the evaluation order. We'll explain it on a simple example: const and its strict version const':

const :: a -> b -> a
const a _ = a

const' :: a -> b -> a
const' a !_ = a
The difference between them is that application of const' will reduce its second argument b to WHNF when the result is demanded while const won't. They can be distinguished by the end user using the calls below:
> const () (error "BANG!")
()
> const' () (error "BANG!")
CallStack (from HasCallStack):
  error, called at :5:12 in interactive:Ghci2

A bang pattern makes a promise to evaluate the expression under the bang when the right hand side is demanded and it is not already in WHNF. These patterns are also forced when a guard protects it, e.g.

> let f !_ | False = 0; f _ = 1 in f undefined
*** Exception: Prelude.undefined
CallStack (from HasCallStack):
  error, called at libraries/base/GHC/Err.hs:79:14 in base:GHC.Err
  undefined, called at :21:36 in interactive:Ghci4
This also applies to bang patterns which appear when we bind a name using where or let clauses (which otherwise are lazy!). Though in this case the expressions will be reduced to WHNF when the whole expression defined by where or let is demanded. Bang patterns are desugared into applications of the seq :: a -> b -> b. This function makes a promise to evaluate the first expression to WHNF when the second expression is evaluated to WHNF, but it does not make any guarantees which expression will be reduced first. An equivalent definition of const' using seq:
const' a b = b `seq` a
In this case seq guarantees that b is evaluated to WHNF when the right hand side of const' is evaluated to WHNF, which happens to be a. Let's also consider another example:
id :: a -> a
id a = a
  where
    Left _ = undefined
Do you know what is the result of id 'a'? The bindings done by where (or let) are lazy which means that they are only demanded when needed. This version of id will return. But the modified version id' will always error:
id' :: a -> a
id' a = a
  where
    !(Left _) = undefined
id' is desugared to:
id' :: a -> a
id' a = b `seq` a
  where
    b@(Left _) = undefined
When id' is called, seq will force b and that's why it errors.

Strict constructor fields

For the sake of completeness, let us mention that one can use bangs within a constructor's fields (which does not require BangPattern extension). For example a strict Maybe type can be defined as

data SMaybe = SNothing | SJust !a
Whenever an expression of type SMaybe is evaluated to WHNF and it turns out that it is an application of SJust, the argument to the constructor will also be evaluated to WHNF.

With this knowledge at hand we can analyse the strict fold from containers.

Strict right-associative fold

The right-associative strict fold in container-0.6.4.1 is defined in the following way:

-- | O(n). A strict version of foldr. Each application of the operator
-- is evaluated before using the result in the next application.
-- This function is strict in the starting value.
foldr' :: (a -> b -> b) -> b -> Map k a -> b
foldr' f b0 = go b0
  where
    go !b Tip             = b
    go  b (Bin _ _ x l r) = go (f x (go b r)) l

Let us consider what happens when foldr' applied to a function which is lazy in its second argument. It turns out that this use of foldr' will not force the accumulator. Let us consider:

-- | @fromList [(0,0+0), (1,1+0), (2,2+0)]@
a :: Map Int Int
a = Bin 3 1 (1+0) (Bin 1 0 (0+0) Tip Tip) (Bin 1 2 (2+0) Tip Tip)
and analyse how foldr' cons nil a would evaluate for some cons and nil.

Strict right-associative fold reduction

Let us analyse how foldr' cons nil a reduces:
foldr' cons nil a
---------------------------------------------- (1)
go nil (Bin 3 1 x1 l r)
  where
    x1 = 1+0
    l = Bin 1 0 (0+0) Tip Tip
    r = Bin 1 2 (2+0) Tip Tip
---------------------------------------------- (2)
go (cons x1 (go nil r)) l
---------------------------------------------- (3)
go (cons x0 (go (cons x1 (go nil r)) Tip)) Tip
  where
    x0 = 0+0
---------------------------------------------- (4)
cons x0 (go (cons x1 (go nil r)) Tip)
In step (1) the topmost, leftmost redex is the foldr' application. In steps (2), (3), and (4) it is the top most go application. In step (2) and (4) we just use the definition of go, in step (3) we substitute l and use the definition of go. Whenever cons is lazy in its second argument then the go redex will not be evaluated. We can verify this in ghci:
> let !x = foldr' (:) [] a 
> :sprint x
x = _ : _
Furthermore, with the help of ghc-heap-view library, we can explore the heap layout of x:
> HeapView.buildHeapTree 100 (HeapView.asBox x)
    >>= print . HeapView.ppHeapTree
_ind (_bh (_thunk() :
            _thunk _fun (_thunk() :
              _thunk _fun
                ([] 5053632)
                (Bin (I# 2) _thunk() (Tip 4471056) (Tip 4471056) 1)
              )
            (Tip 4471056)))
We can see which parts of the expression were already evaluated to WHNF, which gives more precise information than :sprint does. The first _thunk() is the unevaluated 0+0 expression, then : is the evaluated list constructor. The tail of the list is an unevaluated call to go, which is represented by _thunk _fun. It is applied to cons x1 (go nil r) which itself is evaluated to WHNF in step (2), hence it has the form (:) x1 (go [] r), x1 is an unevaluated thunk represented by _thunk() which stands for the expression 1+0, the empty list constructor is evaluated in the step (1).

Stricter right-associative fold

There is another possible definition of a strict right-associative fold:

foldr'' :: (a -> b -> b) -> b -> Map k a -> b
foldr'' f b0 = go b0
  where
    go b Tip             = b
    go b (Bin _ _ x l r) = go (f x $! go b r) l
Instead of forcing the accumulator b when it is applied to go, we force it when it is applied to f (as mentioned in the original haddock documentation). Let us consider how foldr'' cons nil a would be evaluated by Haskell:
foldr'' cons nil a
------------------------------------------------- (1)
go nil (Bin 3 1 x1 l r)
  where
    x1 = 1+0
    l = Bin 1 0 (0+0) Tip Tip
    r = Bin 1 2 (2+0) Tip Tip
------------------------------------------------- (2)
go (cons x1 $! go nil r) l
------------------------------------------------- (3)
go (cons x0 $! go (cons x1 $! go nil r) Tip) Tip
  where
    x0 = 0+0
------------------------------------------------- (4) force first $!
cons x0 $! go (cons x1 $! go nil r) Tip
------------------------------------------------- (5) force second $!
cons x0 $! cons x1 $! go nil r
------------------------------------------------- (6)
cons x0 $! cons x1 $! go (cons x2 $! go nil Tip) Tip
------------------------------------------------- (7) force last $!
cons x0 $! cons x1 $! cons x2 $! go nil Tip
------------------------------------------------- (8)
cons x0 $! cons x1 $! cons x2 $! nil
The difference is that now, when we have cons y $! go ... we need to force the application of go to WHNF as well, that's why the evaluation goes further. We can verify this in a new session in ghci (or using a fresh copy of a):
> let !x = foldr'' (:) [] a
> :sprint x
[_,_,_]
or using the ghc-heap-view library:
> HeapView.buildHeapTree 100 (HeapView.asBox x)
    >>= print . HeapView.ppHeapTree
_ind (_bh (_thunk() : _thunk() : _thunk() : [] 5056128))

Consequences

My patch was recently accepted and various strict folds in containers-0.6.5.1 became more strict. The version of foldr' in containers-0.6.5.1 is a combination of foldr' and foldr'' from this blog post. If you want to play around this, checkout the gist.