Fold that's both constant-space and short-circuiting

Question

I'm trying to build a Haskell function that does basically the same thing as Prelude's product. Unlike that function, however, it should have these two properties:

1. It should operate in constant space (ignoring the fact that some numeric types like Integer aren't). For example, I want myProduct (replicate 100000000 1) to eventually return 1, unlike Prelude's product which uses up all of my RAM and then gives *** Exception: stack overflow.
2. It should short-circuit when it encounters a 0. For example, I want myProduct (0:undefined) to return 0, unlike Prelude's product which gives *** Exception: Prelude.undefined.

Here's what I've come up with so far:

myProduct :: (Eq n, Num n) => [n] -> n
myProduct = go 1
  where go acc (x:xs) = if x == 0 then 0 else acc `seq` go (acc * x) xs
        go acc []     = acc

That works exactly how I want it to for lists, but I'd like to generalize it to have type (Foldable t, Eq n, Num n) => t n -> n. Is it possible to do this with any of the folds? If I just use foldr, then it will short-circuit but won't be constant-space, and if I just use foldl', then it will be constant-space but won't short-circuit.

Gorse

Just an aside: product doesn't stack overflow in a compiled program. – Threnody

Answer 1

You might be looking for foldM. Instantiate it with m = Either b and you get short circuiting behavior (or Maybe, depends if you have many possible early exit values, or one known in advance).

foldM :: (Foldable t, Monad m) => (b -> a -> m b) -> b -> t a -> m b

I recall discussions whether there should be foldM', but IIRC GHC does the right thing most of the time.

import Control.Monad
import Data.Maybe

myProduct :: (Foldable t, Eq n, Num n) => t n -> n
myProduct = fromMaybe 0 . foldM go 1
  where go acc x = if x == 0 then Nothing else Just $! acc * x

Jipijapa

Something cool I just noticed: if you take this answer, inline the definition of foldM, and simplify, the result is the same as the final result in Daniel Wagner's answer. – Gorse

Answer 2

If you spell your function slightly differently, it's more obvious how to turn it into a foldr. Namely:

myProduct :: (Eq n, Num n) => [n] -> n
myProduct = flip go 1 where
    go (x:xs) = if x == 0 then \acc -> 0 else \acc -> acc `seq` go xs (acc * x)
    go [] = \acc -> acc

Now go has got that foldr flavor, and we can just fill in the holes.

myProduct :: (Foldable t, Eq n, Num n) => t n -> n
myProduct = flip go 1 where
    go = foldr
        (\x f -> if x == 0 then \acc -> 0 else \acc -> acc `seq` f (acc * x))
        (\acc -> acc)

Hopefully you can see where each of those pieces came from in the previous explicit-recursion style and how mechanical the transformation is. Then I'd make a few aesthetic tweaks:

myProduct :: (Foldable t, Eq n, Num n) => t n -> n
myProduct xs = foldr step id xs 1 where
    step 0 f acc = 0
    step x f acc = f $! acc * x

And we're all done! A bit of quick testing in ghci reveals that it still short-circuits on 0 as required and uses constant space when specialized to lists.

Threnody

I've seen the pattern of "foldr to build a function that you then immediately call" before, but this is the first time it's been clear exactly how/why it works. Thanks! – Gorse