如何让ST计算产生懒惰的结果流(或者像协同例程一样操作)？

Question

我正在努力解决如何在Haskell中进行有状态计算懒惰地生成结果的一般问题。例如，在Python生成器工具的帮助下，下面的简单算法可以表示为有状态但“惰性”的计算，只执行到达下一个yield语句所需的步骤，然后将控制流返回给调用者，直到请求下一个元素：

def solveLP(vmax0, elems):
    elem_true_ixs = [ [ ei for ei, b in enumerate(row) if b ] for row in elems ]

    return go(vmax0, elem_true_ixs)

def go(vmax, mms):
    if not mms:
        yield []

    else:
        for ei in mms[0]:
            maxcnt = vmax[ei]

            if not maxcnt > 0:
                continue

            vmax[ei] = maxcnt-1 # modify vmax vector in-place
            for es in go(vmax, mms[1:]):
                # note: inefficient vector-concat operation
                # but not relevant for this question
                yield [ei]+es
            vmax[ei] = maxcnt   # restore original vmax state


for sol in solveLP([1,2,3],[[True,True,False],[True,False,True]]):
    print sol

# prints [0,2], [1,0], and [1,2]

这可以很容易地转换为惰性哈斯克尔计算(例如，当m专门用于Logic或[]时)，例如

import           Control.Monad
import qualified Data.Vector.Unboxed as VU

solveLP :: MonadPlus m => VU.Vector Int -> [[Bool]] -> m [Int]
solveLP vmax0 elems = go vmax0 elemTrueIxs
  where
    -- could be fed to 'sequence'
    elemTrueIxs = [ [ ei | (ei,True) <- zip [0::Int ..] row ] | row <- elems ]

    go vmax []     = return []
    go vmax (m:ms) = do
        ei <- mlist m

        let vmax'  = vmax VU.// [(ei, maxcnt-1)] -- this operation is expensive
            maxcnt = vmax VU.! ei

        guard $ maxcnt > 0

        es <- go vmax' ms

        return $ (ei:es)

    mlist = msum . map return

...but我希望能够更接近原始的vmax0实现，通过使用可变向量，并就地修改单个Python向量(因为我只需要递增/递减单个元素，并且复制整个向量来替换单个元素的开销就越大)；请注意，这只是我一直在尝试实现的一类算法的一个玩具示例

所以我的问题是--假设有一种方法可以做到这一点--我如何在ST monad中表达这样的有状态算法，同时仍然能够在计算过程中生成结果后立即将结果返回给调用者？我尝试过通过monad-transformers将ST monad与list monad结合起来，但我不知道如何让它工作……

Answer 1

只要使用lazy ST即可。在Haskell中，普通的旧列表基本上与Python生成器相同，所以我们将返回一个结果列表(其中的结果是一个[Int])。以下是您的Python代码的音译：

import Control.Monad.ST.Lazy
import Data.Array.ST
import Control.Monad
import Data.List

solveLP :: [Int] -> [[Bool]] -> [[Int]]
solveLP vmax_ elems_ = runST $ do
    vmax <- newListArray (0, length vmax_) vmax_
    let elems = map (findIndices id) elems_
    go vmax elems

go :: STArray s Int Int -> [[Int]] -> ST s [[Int]]
go vmax [] = return [[]]
go vmax (mm:mms) = liftM concat . forM mm $ \ei -> do
    maxcnt <- readArray vmax ei
    if not (maxcnt > 0) then return [] else do
        writeArray vmax ei (maxcnt - 1)
        rest <- go vmax mms
        writeArray vmax ei maxcnt
        return (map (ei:) rest)

例如，试试solveLP [1,undefined,3] [[True,True,False],[True,False,True]]，看看它是否真的懒惰地返回结果。

Answer 2

对我来说，现在花时间理解你的算法还为时过早。但是如果我没看错下面的问题，你可以使用lazy ST。下面是一个简单的例子：

import Control.Monad.ST.Lazy
import Data.STRef.Lazy

generator :: ST s [Integer]
generator = do
    r <- newSTRef 0
    let loop = do
            x <- readSTRef r
            writeSTRef r $ x + 1
            xs <- loop
            return $ x : xs
    loop

main :: IO ()
main = print . take 25 $ runST generator

它实际上是从一个维护其状态的ST操作创建一个惰性结果流。

Answer 3

让我们对Python代码进行更直接的转换。您在Python中使用协程，那么为什么不在Haskell中使用协程呢？然后是可变向量的问题；请参阅下面的更多详细信息。

首先，成吨的进口：

-- Import some coroutines
import Control.Monad.Coroutine -- from package monad-coroutine

-- We want to support "yield" functionality like in Python, so import it:
import Control.Monad.Coroutine.SuspensionFunctors (Yield(..), yield)

-- Use the lazy version of ST for statefulness
import Control.Monad.ST.Lazy

-- Monad utilities
import Control.Monad
import Control.Monad.Trans.Class (lift)

-- Immutable and mutable vectors
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import Data.Vector.Mutable (STVector)
import qualified Data.Vector.Mutable as Vector

以下是一些实用程序定义，让我们可以像对待Python中的行为一样对待协程，或多或少：

-- A generator that behaves like a "generator function" in Python
type Generator m a = Coroutine (Yield a) m ()

-- Run a generator, collecting the results into a list
generateList :: Monad m => Generator m a -> m [a]
generateList generator = do
  s <- resume generator -- Continue where we left off
  case s of
    -- The function exited and returned a value; we don't care about the value
    Right _ -> return []
    -- The function has `yield`ed a value, namely `x`
    Left (Yield x cont) -> do
      -- Run the rest of the function
      xs <- generateList cont
      return (x : xs)

现在我们需要能够以某种方式使用STVector。您声明要使用lazy ST，而STVector上的预定义操作仅为strict ST定义，因此我们需要创建一些包装函数。我不喜欢为这类东西编写运算符，但如果您真的想让代码具有pythonic风格(例如，writeLazy的$=或其他任何东西；您需要以某种方式处理索引投影，但无论如何都可以让它看起来更好)，您可以这样做。

writeLazy :: STVector s a -> Int -> a -> ST s ()
writeLazy vec idx val = strictToLazyST $ Vector.write vec idx val

readLazy :: STVector s a -> Int -> ST s a
readLazy vec idx = strictToLazyST $ Vector.read vec idx

thawLazy :: Vector a -> ST s (STVector s a)
thawLazy = strictToLazyST . Vector.thaw

所有的工具都在这里，所以让我们来翻译一下算法：

solveLP :: STVector s Int -> [[Bool]] -> Generator (ST s) [Int]
solveLP vmax0 elems =
  go vmax0 elemTrueIxs
  where
    elemTrueIxs = [[ei | (ei, True) <- zip [0 :: Int ..] row] | row <- elems]

go :: STVector s Int -> [[Int]] -> Generator (ST s) [Int]
go _ [] = yield []
go vmax (m : ms) = do
  forM_ m $ \ ei -> do
    maxcnt <- lift $ readLazy vmax ei
    when (maxcnt > 0) $ do
      lift $ writeLazy vmax ei $ maxcnt - 1
      sublist <- lift . generateList $ go vmax ms
      forM_ sublist $ \ es -> yield $ ei : es
      lift $ writeLazy vmax ei maxcnt

遗憾的是，没有人费心为Coroutine定义MonadPlus，所以guard在这里不可用。但这可能不是您想要的，因为它在某些/大多数monads中停止时会引发错误。当然，我们还需要将在ST monad中完成的所有操作从Coroutine monad中剔除；这是一个小麻烦。

这就是所有的代码，所以可以简单地运行它：

main :: IO ()
main =
  forM_ list print
  where
    list =  runST $ do
      vmax <- thawLazy . Vector.fromList $ [1, 2, 3]
      generateList (solveLP vmax [[True, True, False], [True, False, True]])

list变量是纯延迟生成的。

我有点累了，所以如果有什么不合理的地方，请毫不犹豫地指出。