How to make ST computation produce a lazy result stream (or work as a joint procedure)?

Question

I am struggling with a general problem on how to make stateful computation in Haskell a lazy result. For instance. The following simple algorithm can be expressed using the Python generator as a calculated, but "lazy" calculation, performing only the steps necessary to achieve the next yield , 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 easily be translated into Haskell's lazy calculation (for example, when m specializes in Logic or [] ), for example

 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 would like to be closer to the original Python implementation using mutable vectors and changing one vmax0 vector in place (since I just need to increase / decrease one element and copying the whole vector just to replace one element - this is quite an overhead than the vector becomes longer); note that this is just a toy example for the class of algorithms that I tried to implement

So my question is - if you have a way to do this - how can I express such a stateful algorithm in the ST monad, still having the ability to return the results back to the caller as soon as they are created during the calculation? I tried to combine ST monad with list monad with transformer monad, but I could not figure out how to make it work ...

Answer 1

Just use the lazy ST. In Haskell, plain old lists are basically identical to Python generators, so we will return a list of results (where the result is [Int] ). Here's the 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 for example. solveLP [1,undefined,3] [[True,True,False],[True,False,True]] to see that it does return results lazily.

Answer 2

It is too early for me to understand your algorithm in time. But if I read the main question correctly, you can use lazy ST. Here's a trivial example:

 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 

This accurately creates a lazy result stream from the ST action, which maintains its state.

Answer 3

Let’s do a more direct translation of the Python code. You use coroutines in Python, so why not just use coroutines in Haskell? Then the question arises of mutable vectors; see below for more details.

First of all, tons of imports:

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

Here are some utility definitions that allow coroutines to be processed as if they behaved as in Python, more or less:

 -- A generator that behaves like a "generator function" in Python type Generator ma = Coroutine (Yield a) m () -- Run a generator, collecting the results into a list generateList :: Monad m => Generator ma -> 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) 

Now we need to use the STVector . You stated that you want to use a lazy ST, and the predefined operations on the STVector defined only for strict ST, so we need to make several STVector functions.I don’t do statements for such things, but you could if you really want to make pythonic code ( Eg $= for writeLazy or something else, you need to somehow handle the index projection, but it's possible that it looks the best anyway).

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

All the tools are here, so let's just translate the algorithm:

 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 

Unfortunately, no one bothered to define MonadPlus for Coroutine s, so guard is not available here. But this is probably not what you wanted, as it causes an error when stopping in some / most monads. Of course, we also need to lift all operations performed in the ST monad from the Coroutine monad; minor nuisance.

That is all the code, so you can just run it:

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

The list variable is clean and lazily generated.

I'm a little tired, so if something does not make sense, please feel free to point it out.