Memory leak in for_, traverse_ and forM_ for Maybe and many other types. Test case for https://www.yesodweb.com/blog/2014/05/foldable-mapm-maybe-recursive

TraverseMaybePerformance.hs

#!/usr/bin/env stack
-- stack --resolver lts-11.22 script --optimize

-- The problem:
-- traverse_ / forM_ / for_ leak memory for `Maybe`.
--
-- There are a couple background posts about it:
--     2014-05 Yesodweb: https://www.yesodweb.com/blog/2014/05/foldable-mapm-maybe-recursive
--     2014-05 School of Haskell: https://www.schoolofhaskell.com/user/snoyberg/general-haskell/basics/foldable-mapm-maybe-and-recursive-functions
--     2017-01 snoyman.com: https://www.snoyman.com/blog/2017/01/foldable-mapm-maybe-and-recursive-functions
-- The ones from 2014 mention to a stack overflow crash, and in
-- the newer one it measured memory since the stack limit was removed
-- (https://ghc.haskell.org/trac/ghc/ticket/8189):
--
-- Reddit discussions:
--     2014: https://www.reddit.com/r/haskell/comments/24t0sj/foldablemapm_maybe_and_recursive_functions_school/
--     2017: https://www.reddit.com/r/haskell/comments/5n63ow/foldablemapm_maybe_and_recursive_functions/

-- Instructions:
-- First compile with
--   stack ./Traverse
-- Then run with:
--   /usr/bin/time ./TraverseMaybePerformance 1
--   /usr/bin/time ./TraverseMaybePerformance 2
--   /usr/bin/time ./TraverseMaybePerformance 3
--   /usr/bin/time ./TraverseMaybePerformance 4
--   /usr/bin/time ./TraverseMaybePerformance 5
--   ...

{-# OPTIONS_GHC -Wall #-}

module Main (main, united_traverse_2, united_for_3, united_traverse_4 ,united_for_4) where -- TODO remove

import qualified Data.Text.Lazy as T
import qualified Data.Foldable as F
import Control.Monad (when)
import System.Environment (getArgs)

printChars_maybe :: Int -> T.Text -> IO ()
printChars_maybe idx t = maybe (return ()) (\(c, t') -> do
    when (idx `mod` 100000 == 0)
        $ putStrLn $ "Character #" ++ show idx ++ ": " ++ show c
    printChars_maybe (idx + 1) t') (T.uncons t)

printChars_forM_ :: Int -> T.Text -> IO ()
printChars_forM_ idx t = F.forM_ (T.uncons t) $ \(c, t') -> do
    when (idx `mod` 100000 == 0)
        $ putStrLn $ "Character #" ++ show idx ++ ": " ++ show c
    printChars_forM_ (idx + 1) t'

printChars_for_ :: Int -> T.Text -> IO ()
printChars_for_ idx t = F.for_ (T.uncons t) $ \(c, t') -> do
    when (idx `mod` 100000 == 0)
        $ putStrLn $ "Character #" ++ show idx ++ ": " ++ show c
    printChars_for_ (idx + 1) t'

forM_Maybe :: Monad m => Maybe a -> (a -> m ()) -> m ()
forM_Maybe Nothing _ = return ()
forM_Maybe (Just x) f = f x

printChars_forM_Maybe :: Int -> T.Text -> IO ()
printChars_forM_Maybe idx t = forM_Maybe (T.uncons t) $ \(c, t') -> do
    when (idx `mod` 100000 == 0)
        $ putStrLn $ "Character #" ++ show idx ++ ": " ++ show c
    printChars_forM_Maybe (idx + 1) t'

united_for_ :: (Foldable t, Applicative f) => t a -> (a -> f ()) -> f ()
united_for_ = F.for_

printChars_united_for_ :: Int -> T.Text -> IO ()
printChars_united_for_ idx t = united_for_ (T.uncons t) $ \(c, t') -> do
    when (idx `mod` 100000 == 0)
        $ putStrLn $ "Character #" ++ show idx ++ ": " ++ show c
    printChars_united_for_ (idx + 1) t'

united_traverse_2 :: (Foldable f, Applicative app) => (a -> app ()) -> f a -> app ()
united_traverse_2 f = go . F.toList
  where
    go l = case l of
      [] -> pure ()
      x : xs -> f x *> go xs

united_for_2 :: (Foldable f, Applicative app) => f a -> (a -> app ()) -> app ()
united_for_2 = flip united_traverse_2

printChars_united_for_2 :: Int -> T.Text -> IO ()
printChars_united_for_2 idx t = united_for_2 (T.uncons t) $ \(c, t') -> do
    when (idx `mod` 100000 == 0)
        $ putStrLn $ "Character #" ++ show idx ++ ": " ++ show c
    printChars_united_for_2 (idx + 1) t'

-- Provided by Michael Snoyman
united_for_3 :: (Foldable f, Applicative app) => f a -> (a -> app ()) -> app ()
united_for_3 t f =
  case F.toList t of
    [] -> pure ()
    x -> go x
  where
    go [x] = f x
    go (x:xs) = f x *> go xs

printChars_united_for_3 :: Int -> T.Text -> IO ()
printChars_united_for_3 idx t = united_for_3 (T.uncons t) $ \(c, t') -> do
    when (idx `mod` 100000 == 0)
        $ putStrLn $ "Character #" ++ show idx ++ ": " ++ show c
    printChars_united_for_3 (idx + 1) t'

united_traverse_4 :: (Foldable f, Applicative app) => (a -> app ()) -> f a -> app ()
united_traverse_4 f t =
  case F.toList t of
    [] -> pure ()
    x -> go x
  where
    go [x] = f x
    go (x:xs) = f x *> go xs

united_for_4 :: (Foldable f, Applicative app) => f a -> (a -> app ()) -> app ()
united_for_4 = flip united_traverse_4

printChars_united_for_4 :: Int -> T.Text -> IO ()
printChars_united_for_4 idx t = united_for_4 (T.uncons t) $ \(c, t') -> do
    when (idx `mod` 100000 == 0)
        $ putStrLn $ "Character #" ++ show idx ++ ": " ++ show c
    printChars_united_for_4 (idx + 1) t'

united_traverse_5 :: (Foldable f, Applicative app) => (a -> app ()) -> f a -> app ()
united_traverse_5 f = F.foldl' (\c a -> c *> f a) (pure ())

united_for_5 :: (Foldable f, Applicative app) => f a -> (a -> app ()) -> app ()
united_for_5 = flip united_traverse_5

printChars_united_for_5 :: Int -> T.Text -> IO ()
printChars_united_for_5 idx t = united_for_5 (T.uncons t) $ \(c, t') -> do
    when (idx `mod` 100000 == 0)
        $ putStrLn $ "Character #" ++ show idx ++ ": " ++ show c
    printChars_united_for_5 (idx + 1) t'

test1 :: IO ()
test1 = printChars_maybe 1 $ T.replicate 5000000 $ T.singleton 'x'

test2 :: IO ()
test2 = printChars_forM_ 1 $ T.replicate 5000000 $ T.singleton 'x'

test3 :: IO ()
test3 = printChars_forM_Maybe 1 $ T.replicate 5000000 $ T.singleton 'x'

test4 :: IO ()
test4 = printChars_for_ 1 $ T.replicate 5000000 $ T.singleton 'x'

test5 :: IO ()
test5 = printChars_united_for_ 1 $ T.replicate 5000000 $ T.singleton 'x'

test6 :: IO ()
test6 = printChars_united_for_2 1 $ T.replicate 5000000 $ T.singleton 'x'

test7 :: IO ()
test7 = printChars_united_for_3 1 $ T.replicate 5000000 $ T.singleton 'x'

test8 :: IO ()
test8 = printChars_united_for_4 1 $ T.replicate 5000000 $ T.singleton 'x'

test9 :: IO ()
test9 = printChars_united_for_5 1 $ T.replicate 5000000 $ T.singleton 'x'

main :: IO ()
main = do
  args <- getArgs
  case args of
    -- Note that for me, the 6 MB variants run 2x as fast with -O2 as with -O
    -- (stack script's `--optimize` does -O2).
    -- But the memory usage is stay equal for the two.
    ["1"] -> test1 --  6 MB maxresident for maybe
    ["2"] -> test2 -- 55 MB maxresident for forM_
    ["3"] -> test3 --  6 MB maxresident for forM_Maybe
    ["4"] -> test4 -- 55 MB maxresident for for_
    ["5"] -> test5 -- 55 MB maxresident for united_for_
    ["6"] -> test6 -- 94 MB maxresident for united_for_2
    ["7"] -> test7 --  6 MB maxresident for united_for_3
    ["8"] -> test8 --  6 MB maxresident for united_for_4 (via flip of united_for_3's traverse_ equivalent)
    ["9"] -> test9 --  6 MB maxresident for united_for_5
    _ -> error "bad arguments"