Question

Consider the following abbreviated code from this excellent blog post:

import System.Random (Random, randomRIO)

newtype Stream m a = Stream { runStream :: m (Maybe (NonEmptyStream m a)) }
type NonEmptyStream m a = (a, Stream m a)

empty :: (Monad m) => Stream m a
empty = Stream $ return Nothing

cons :: (Monad m) => a -> Stream m a -> Stream m a
cons a s = Stream $ return (Just (a, s))

fromList :: (Monad m) => [a] -> NonEmptyStream m a
fromList (x:xs) = (x, foldr cons empty xs)

Not too bad thus far – a monadic, recursive data structure and a way to build one from a list.

Now consider this function that chooses a (uniformly) random element from a stream, using constant memory:

select :: NonEmptyStream IO a -> IO a
select (a, s) = select' (return a) 1 s where
  select' :: IO a -> Int -> Stream IO a -> IO a
  select' a n s = do
    next <- runStream s
    case next of 
      Nothing -> a
      Just (a', s') -> select' someA (n + 1) s' where
        someA = do i <- randomRIO (0, n) 
                   case i of 0 -> return a'
                             _ -> a

I’m not grasping the mysterious cyclic well of infinity that’s going on in the last four lines; the result a' depends on a recursion on someA, which itself could depend on a', but not necessarily.

I get the vibe that the recursive worker is somehow ‘accumulating’ potential values in the IO a accumulator, but I obviously can’t reason about it well enough.

Could anyone provide an explanation as to how this function produces the behaviour that it does?

Answer 1

That code doesn’t actually run in constant space, as it composes a bigger and bigger IO a action which delays all the random choices until it’s reached the end of the stream. Only when we reach the Nothing -> a case does the action in a actually get run.

For example, try running it on an infinite, constant space stream made by this function:

repeat' :: a -> NonEmptyStream IO a
repeat' x = let xs = (x, Stream $ return (Just xs)) in xs

Obviously, running select on this stream won’t terminate, but you should see the memory usage going up as it allocates a lot of thunks for the delayed actions.

Here’s a slightly re-written version of the code which does the choices as it goes along, so it runs in constant space and should hopefully be more clear as well. Note that I’ve replaced the IO a argument with a plain a which makes it clear that there are no delayed actions being built up here.

select :: NonEmptyStream IO a -> IO a
select (x, xs) = select' x 1 xs where
  select' :: a -> Int -> Stream IO a -> IO a
  select' current n xs = do
    next <- runStream xs
    case next of 
      Nothing       -> return current
      Just (x, xs') -> do
        i <- randomRIO (0, n)                         -- (1)
        case i of                                     
          0 -> select' x       (n+1) xs'              -- (2)
          _ -> select' current (n+1) xs'              -- (3)

As the name implies, current stores the currently selected value at each step. Once we’ve extracted the next item from the stream, we (1) pick a random number and use this to decide whether to (2) replace our selection with the new item or (3) keep our current selection before recursing on the rest of the stream.