从State切换到StateT后，如何恢复对单子构造的列表的惰性计算？

Question

使用以下代码：

（[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

我可以将nMax设置为5，即50,000,000，并且我得到的运行时间大致相同：

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

正如我所期望的，鉴于我对惰性评估机制的理解。

但是，如果我从State切换到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

然后我看到了各个运行时间之间的极大差异：

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

并且我认为那是因为当我切换到基于StateT的实现时，我对单子构造的列表失去了懒惰的评估。

1. 正确吗？

2. 我能否恢复单子构造列表的惰性求值，同时保持基于StateT的实现？

Answer 1

在您的示例中，每个foo仅运行一个runState操作，因此对State和/或StateT的使用本质上是无关紧要的。您可以将foo的使用替换为等效项：

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

并且其行为方式相同。

问题是IO monad的严格性。如果您改为在Identity monad中运行此计算，则：

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

然后它会懒惰地运行。

如果要在IO monad中延迟运行，则需要使用unsafeInterleaveIO明确地执行此操作，因此可以使用以下命令：

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 []

Answer 2

[KA Buhr的另一个答案解释了为什么State vs StateT不是相关因素（IO是，并且还指出了您的示例是如何构造的（因为State(T)部分不是实际用作每个数字使用新状态[]）。但是除了这些要点之外，我不确定我会说“失去对单子构造的列表的惰性求值”，因为如果我们理解诸如“惰性求值=仅在需要时才求值”之类的内容，则确实需要运行foo在输入列表上的每个元素上执行所有效果，因此懒惰的评估不会“丢失”。您正在得到想要的东西。 （碰巧的是foo不执行任何IO，如果编译器/ GHC有可能在此基础上对其进行优化，也许其他人可以发表评论，但是您可以很容易地理解为什么GHC会天真这里的东西。）

这是Haskell中常见的，众所周知的问题。有各种各样的库（其中最知名的是streaming，pipes，conduit）通过为您提供streams（基本上是列表），它们也效果懒惰解决了该问题。如果我以流式方式重新创建您的示例，

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

然后，这两个版本实际上都可以立即运行。为了使区别更加明显，请设想foo也称为print "hi"。然后，对效果不满意的流媒体版本将只打印两次，而您的原始版本都将打印nMax次。由于它们在效果上比较懒惰，因此不需要遍历整个列表即可短路并尽早完成。