Use a more stack efficient version of mapM

[josefs]

The function mapM is notorious for eating up a lot of stack space.
This diff provides a version of mapM for the Maybe monad and the
list traversable which doesn't eat up the stack.

Having a version of mapM which doesn't eat a lot of stack is
particularly helpful when trying to find spaceleaks with the
method described here:
http://neilmitchell.blogspot.com/2015/09/detecting-space-leaks.html

[ygale]

Thanks for the great suggestion.

Let's keep in mind that the main goal of this library, like the time library, is to guarantee correctness of time semantics. So we strongly prefer simple, expressive, and clearly correct types and implementations. As a subgoal - that means avoid CPS at almost any cost.

But still, if we can improve stack usage, it would be nice.

First of all, are you sure that the current implementation eats stack? Have you tested it? I'm not convinced that is the case here, even though in some cases mapM can be problematic.

If you can show it is true that we build up stack - can we think of a simpler (non-CPS) way to improve that?

What does mapM mean in the Maybe monad? We traverse the entire result list to verify there aren't any Nothings. At the first Nothing we hit, we abort the whole computation and just return Nothing. If we can get all the way to the end without hitting a Nothing, we go back to the beginning and return the entire result list.

There is no way to avoid traversing the entire spine of the list and saving the thunks before returning a result. We just have to make sure that the thunks are on the heap and not the stack.

[ygale]

@josefs See also Ed Kmett's comment to Neil Mitchell's blog post that you linked. This implies that the behavior of mapM in current base is improved, though not as good as Neil's implementation.

[ygale]

@josefs Here are two more points:

1. The list we are mapping over is small - its length is the total number of historical wall clock transitions for a given time zone. These tend to be no more than about 200 entries, and typically grow at the rate of about 2 entries per year. When I originally wrote my code, I wasn't careful about optimizing for long lists; it is likely even quadratic in a few places. If you could fix that, it would probably be a better optimization than messing around with that mapM.
2. It should be possible to eliminate the mapM completely. We are only in the Maybe monad there to guard against a corrupted Olson file where the number of transition times is not equal to the the number of transition infos. If it's really true that the mapM is a performance bottleneck, we could rewrite the code to do that some other more efficient way.

Looking forward to hearing your thoughts.

[josefs]

Thanks for the feedback!

First of all let me explained the way I tested this. I wrote the following small test program and ran it with a restricted stack size. ("Israel" is the largest time zone file I have easy access to)

import Data.Time.LocalTime.TimeZone.Olson
import Paths_timezone_olson

main = do
  filepath <- getDataFileName "Israel"
  timeSeries <- getTimeZoneSeriesFromOlsonFile filepath
  print timeSeries

Before my change I needed a stack size of 32K. After my change I only needed 1K to run it.

I think it makes a lot of sense to check for corrupted timezone files and use the Maybe monad to capture any problems.

It is certainly possible to simplify my implementation. The fail argument can be removed and replaced with Nothing everywhere. That'll make the code a bit simpler. It would also be possible to defunctionalize the succ continuation to a separate datatype. That would look something like this:

mapMMaybe :: (a -> Maybe b) -> [a] -> Maybe [b]
mapMMaybe f ls = mapMCont f ls JustF

mapMCont :: (a -> Maybe b) -> [a] -> JustFun b -> Maybe [b]
mapMCont f [] = \succ -> apply succ []
mapMCont f (x:xs) = \succ ->
  case f x of
    Nothing -> Nothing
    Just x  -> mapMCont f xs (ConsF succ x)

data JustFun b
  = JustF
  | ConsF (JustFun b) b

apply :: JustFun b -> [b] -> Maybe [b]
apply JustF bs = Just bs
apply (ConsF succ x) bs = (apply succ) (x:bs)

Would you consider this implementation to be preferable? I'm afraid I don't know of a simple way to eliminate the stack leak.

[josefs]

The datatype JustFun is essentially just a list and apply is essentially reverse, so we can simplify things further a little bit:

mapMMaybe :: (a -> Maybe b) -> [a] -> Maybe [b]
mapMMaybe f ls = mapMCont f ls []

mapMCont :: (a -> Maybe b) -> [a] -> [b] -> Maybe [b]
mapMCont f [] acc = Just (reverse acc)
mapMCont f (x:xs) acc =
  case f x of
    Nothing -> Nothing
    Just x  -> mapMCont f xs (x:acc)

Personally I think it's clearer that this code doesn't use any stack space compared to the continuation passing implementation in my patch. All calls are tail calls. So thanks for asking me to simplify things!

[josefs]

Ping

[ygale]

Hi @josefs I am really sorry for leaving this great idea languish for so long.

Are you still with me? If so, I would like to merge this. Here are a last few small suggested tweaks, please let me know what you think:

1. In the last equation for mapCont, are you sure we don't need a bang on !acc? Otherwise it looks like we'll still get a build-up of thunks.
2. Also there, I would like to avoid shadowing the variable x.
3. Also there, the explicit case looks to me like a re-implementation of bind in the Maybe monad. Would we lose anything to write it using bind? Something like this:
   mapCont f (x:xs) !acc = f x >>= mapCont f xs . (: acc)
   It's a matter of style whether that is simpler. I'll leave it up to you.
4. Please merge with master - we switched from Travis CI to Github actions and we need the CI to run.
5. Once you decide on the final code version, please run your test one more time to ensure we really are still plugging the space leak.

Thanks!

Another random observation: When I saw your code I immediately thought of dlist. But dlist is solving a different problem - it intentionally builds up thunks to avoid repeated list traversals. Your implementation is more like TQueue from stm. Anyway, I'm not looking to add new library dependencies.