How do I recover lazy evaluation of a monadically constructed list, after switching from State to StateT?

user_1549736 0
Viewed 101

With the following code:

(lazy_test.hs)

-- Testing lazy evaluation of monadically constructed lists, using State.
import Control.Monad.State

nMax = 5

foo :: Int -> State [Int] Bool
foo n = do
  modify $ \st -> n : st
  return (n `mod` 2 == 1)

main :: IO ()
main = do
  let ress = for [0..nMax] $ \n -> runState (foo n) []
      sts  = map snd $ dropWhile (not . fst) ress
  print $ head sts

for = flip map

I can set nMax to 5, or 50,000,000, and I get approximately the same run time:

nMax = 5:

$ stack ghc lazy_test.hs
[1 of 1] Compiling Main             ( lazy_test.hs, lazy_test.o )
Linking lazy_test ...

$ time ./lazy_test
[1]

real    0m0.019s
user    0m0.002s
sys     0m0.006s

nMax = 50,000,000:

$ stack ghc lazy_test.hs
[1 of 1] Compiling Main             ( lazy_test.hs, lazy_test.o )
Linking lazy_test ...

$ time ./lazy_test
[1]

real    0m0.020s
user    0m0.002s
sys     0m0.005s

which is as I expect, given my understanding of lazy evaluation mechanics.

However, if I switch from State to StateT:

(lazy_test2.hs)

-- Testing lazy evaluation of monadically constructed lists, using StateT.
import Control.Monad.State

nMax = 5

foo :: Int -> StateT [Int] IO Bool
foo n = do
  modify $ \st -> n : st
  return (n `mod` 2 == 1)

main :: IO ()
main = do
  ress <- forM [0..nMax] $ \n -> runStateT (foo n) []
  let sts  = map snd $ dropWhile (not . fst) ress
  print $ head sts

for = flip map

then I see an extreme difference between the respective run times:

nMax = 5:

$ stack ghc lazy_test2.hs
[1 of 1] Compiling Main             ( lazy_test2.hs, lazy_test2.o )
Linking lazy_test2 ...

$ time ./lazy_test2 
[1]

real    0m0.019s
user    0m0.002s
sys     0m0.004s

nMax = 50,000,000:

$ stack ghc lazy_test2.hs
[1 of 1] Compiling Main             ( lazy_test2.hs, lazy_test2.o )
Linking lazy_test2 ...

$ time ./lazy_test2 
[1]

real    0m29.758s
user    0m25.488s
sys     0m4.231s

And I'm assuming that's because I'm losing lazy evaluation of the monadically constructed list, when I switch to the StateT-based implementation.

  1. Is that correct?

  2. Can I recover lazy evaluation of a monadically constructed list, while keeping with the StateT-based implementation?

io-monad monad-transformers state-monad lazy-evaluation haskell
2 Answers

In your example, you're only running one foo action per runState, so your use of State and/or StateT is essentially irrelevant. You can replace the use of foo with the equivalent:

import Control.Monad

nMax = 50000000

main :: IO ()
main = do
  ress <- forM [0..nMax] $ \n -> return (n `mod` 2 == 1, [n])
  let sts  = map snd $ dropWhile (not . fst) ress
  print $ head sts

and it behaves the same way.

The issue is the strictness of the IO monad. If you ran this computation in the Identity monad instead:

import Control.Monad
import Data.Functor.Identity

nMax = 50000000

main :: IO ()
main = do
  let ress = runIdentity $ forM [0..nMax] $ \n -> return (n `mod` 2 == 1, [n])
  let sts  = map snd $ dropWhile (not . fst) ress
  print $ head sts

then it would run lazily.

If you want to run lazily in the IO monad, you need to do it explicitly with unsafeInterleaveIO, so the following would work:

import System.IO.Unsafe
import Control.Monad

nMax = 50000000

main :: IO ()
main = do
  ress <- lazyForM [0..nMax] $ \n -> return (n `mod` 2 == 1, [n])
  let sts  = map snd $ dropWhile (not . fst) ress
  print $ head sts

lazyForM :: [a] -> (a -> IO b) -> IO [b]
lazyForM (x:xs) f = do
  y <- f x
  ys <- unsafeInterleaveIO (lazyForM xs f)
  return (y:ys)
lazyForM [] _ = return []
user_72030160

The other answer by K A Buhr explains why State vs StateT is not the pertinent factor (IO is), and also points out how your example is strangely constructed (in that the State(T) part isn't actually used as each number uses a new state []). But aside from those points, I'm not sure I would say "losing lazy evaluation of the monadically constructed list", because if we understand something like "lazy evaluation = evaluated only when needed", then foo does indeed need to run on every element on the input list in order to perform all the effects, so lazy evaluation is not being "lost". You are getting what you asked for. (It just so happens that foo doesn't perform any IO, and perhaps someone else can comment with if it's ever possible for a compiler/GHC to optimize it away on this basis, but you can easily see why GHC does the naive thing here.)

This is a common, well-known problem in Haskell. There are various libraries (best known of which are streaming, pipes, conduit) which solve the problem by giving you streams (basically lists) which are lazy in the effects too. If I recreate your example in a streaming style,

import Data.Function ((&))
import Control.Monad.State
import Streaming
import qualified Streaming.Prelude as S

foo :: Int -> StateT [Int] IO Bool
foo n = 
  (n `mod` 2 == 1) <$ modify (n:)

nMax :: Int
nMax = 5000000

main :: IO ()
main = do
  mHead <- S.head_ $ S.each [0..nMax]
                   & S.mapM (flip runStateT [] . foo)
                   & S.dropWhile (not . fst)
  print $ snd <$> mHead

then both versions run practically instantaneously. To make the difference more apparent, imagine that foo also called print "hi". Then the streaming version, being lazy in the effects, would print only twice, whereas your original versions would both print nMax times. As they're lazy in the effects, then the whole list doesn't need to be traversed in order to short-circuit and finish early.

user_112577150
Related