How to make ST computation produce lazy result stream (or operate like a co-routine)?

Question

I'm struggling with the general problem on how to make a stateful computation in Haskell generate results lazily. E.g. the following simple algorithm can be expressed with the help of Python's generator facility as a stateful but "lazy" computation, performing only the steps necessary to reach the next yield statement and then returning the control-flow to the caller until the next element is requested:

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]

This can be easily translated to a lazy Haskell computation (e.g. when m is specialized to Logic or []), e.g.

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 I'd like to be able to get closer to the original Python implementation, by using mutable vectors, and modifying a single vmax0 vector in-place (as I just need to increment/decrement a single element, and copying the whole vector just to replace a single element is quite an overhead the longer the vector becomes); please note that this is just a toy example for a class of algorithms I've been trying to implement

So my question is -- assuming there is a way to accomplish that -- how can I express such a stateful algorithm in the ST monad while still being able to return the results back to the caller as soon as they are produced during the computation? I've tried combining the ST monad with a list monad via monad-transformers but I couldn't figure out how to make it work...

Answer 1

Just use lazy ST. In Haskell, plain old lists are basically identical to Python generators, so we'll return a list of results (where a result is an [Int]). Here's a transliteration of your Python code:

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)

Try e.g. solveLP [1,undefined,3] [[True,True,False],[True,False,True]] to see that it really does return results lazily.