Search code examples
haskellfoldspace-complexityshort-circuiting

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


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.


Solution

  • 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