A monad for building test data

Question

OK, so I'm trying to write a monad for building test data, but I can't quite get it to work how I want it. It looks something like this:

runBuildM :: [i] -> BuildM i o x -> [o]
-- Given a list of i, build a list of o.

source :: BuildM i o i
-- Fetch unique i.

yield :: o -> BuildM i o ()
-- Return a new o to the caller.

gather :: BuildM i o x -> BuildM i o o
-- Fetch every possible o from sub-computation.

local :: BuildM i o x -> BuildM i o x
-- Isolate any source invocations from the rest of the code.

In other words, it's a supply monad, writer monad and list monad. The idea is that I can write something like this:

build_tests depth = do
  local $ do
    v <- source
    yield v
    yield (map toLower v)
  yield "[]"
  yield "()"
  when (depth > 2) $ do
    t1 <- gather $ build_tests (depth-1)
    yield $ "(" ++ t1 ++ ")"
    yield $ "[" ++ t1 ++ "]"
    t2 <- gather $ build_tests (depth-1)
    yield $ "(" ++ t1 ++ "," ++ t2 ++ ")"

The idea is to generate all possible combinations of data. You can do that just using list comprehensions, but the result ends up syntactically awful. This is much more readable. Unfortunately, it doesn't actually work...

The problem seems to boil down to the local function not behaving correctly. The intention is for any source calls in the sub-computation to have no effect outside of it. (I.e., subsequent calls to source from outside the local block get the first token again.) However, what my implementation of local actually does is reset the next token for everything (i.e., including the contents of the sub-computation). This is clearly incorrect, but I cannot for the life of me bend my mind around how to make it work correctly.

The fact that I'm having this much trouble getting the code to work as required probably means the actual internal representation of my monad is just wrong. Can anybody take a stab at implementing this correctly?

---

EDIT: I should perhaps have realised this, but I didn't actually specify the expected result I'm trying to get. The above code is supposed to produce this:

["A", "a", "[]", "()", "(A)", "(a)", "[A]", "[a]", "(A, B)", "(A, b)", "(a, B)", "(a, b)"]

It's not super-critical that the results appear in exactly this order. I'd like the single cases to appear before the compound ones, but I'm not too fussed exactly what order the compounds appear. The rule is that the same variable never appears twice in any individual expression.

If we allow the depth to be a bit deeper, we additionally get terms such as

"((A))", "([A])", "[(A)]", "((A, B), C)", "(A, (B, C))"

and so on.

---

It's clearly broken, but here's what I have so far:

newtype BuildM i o x = BuildM ([i] -> SEQ.Seq ([i], SEQ.Seq o, x))

instance Functor (BuildM i o) where
  fmap uf (BuildM sf) =
    BuildM $ \ is0 -> do
      (is1, os, x) <- sf is0
      return (is1, os, uf x)

instance Applicative (BuildM i o) where
  pure x = BuildM $ \ is0 -> return (is0, SEQ.empty, x)

  BuildM sf1 <*> BuildM sf2 =
    BuildM $ \ is1 -> do
      (is2, os2, f) <- sf1 is1
      (is3, os3, x) <- sf2 is2
      return (is3, os2 >< os3, f x)

instance Monad (BuildM i o) where
  return = pure

  BuildM sf1 >>= uf =
    BuildM $ \ is1 -> do
      (is2, os2, x) <- sf1 is1
      let BuildM sf2 = uf x
      (is3, os3, y) <- sf2 is2
      return (is3, os2 >< os3, y)

runBuildM :: [i] -> BuildM i o x -> [o]
runBuildM is0 (BuildM sf) =
  toList $ do
    (is, os, x) <- sf is0
    os

source :: BuildM i o i
source =
  BuildM $ \ is ->
    if null is
      then error "AHC.Tests.TestBuilder.source: end of input"
      else return (tail is, SEQ.empty, head is)

yield :: o -> BuildM i o ()
yield o = BuildM $ \ is -> return (is, SEQ.singleton o, () )

gather :: BuildM i o x -> BuildM i o' o
gather (BuildM sf1) =
  BuildM $ \ is1 -> do
    (is2, os2, _) <- sf1 is1
    o <- os2
    return (is2, SEQ.empty, o)

local :: BuildM i o x -> BuildM i o ()
local (BuildM sf1) =
  BuildM $ \ is1 ->
    let os = do (is2, os2, x) <- sf1 is1; os2
    in  return (is1, os, () )

Answer 1

If my other answer is overkill, the continuation monad transformer provides a convenient way to construct any MonadPlus values.

The continuation monad lets us easily capture the idea of doing something mplus the as yet unknown remainder.

import Control.Monad
import Control.Monad.Trans.Cont

once :: MonadPlus m => m a -> ContT a m ()
once m = ContT $ \k -> m `mplus` k ()

Yielding a result is just returning it once.

yield :: MonadPlus m => a -> ContT a m ()
yield = once . return

We can gather up all the results by sticking mzero at the end.

gather :: MonadPlus m => ContT a m r -> m a
gather m = runContT m (const mzero)

Your example is written in terms of yield, gather, once, and lift.

import Data.Char

import Control.Monad.Trans.Class

build_tests :: MonadPlus m => m String -> Int -> ContT String m ()
build_tests source = go
  where
    go depth = do
      once . gather $ do
        v <- lift source
        yield v
        yield (map toLower v)
      yield "[]"
      yield "()"
      when (depth > 2) $ do
        t1 <- lift . gather $ go (depth-1)
        yield $ "(" ++ t1 ++ ")"
        yield $ "[" ++ t1 ++ "]"
        t2 <- lift . gather $ go (depth-1)
        yield $ "(" ++ t1 ++ "," ++ t2 ++ ")"

main = print . gather $ build_tests ["A", "B"] 3

This outputs the following:

Depth 2
["A","a","B","b","[]","()"]
Depth 3
["A","a","B","b","[]","()","(A)","[A]","(A,A)","(A,a)","(A,B)","(A,b)","(A,[])","(A,())","(a)","[a]","(a,A)","(a,a)","(a,B)","(a,b)","(a,[])","(a,())","(B)","[B]","(B,A)","(B,a)","(B,B)","(B,b)","(B,[])","(B,())","(b)","[b]","(b,A)","(b,a)","(b,B)","(b,b)","(b,[])","(b,())","([])","[[]]","([],A)","([],a)","([],B)","([],b)","([],[])","([],())","(())","[()]","((),A)","((),a)","((),B)","((),b)","((),[])","((),())"]

I've taken the liberty of getting rid of the requirement to read the original source from the environment for simplicity. You can add a ReaderT to the transformer stack to get it back. I also haven't chosen a list transfomer for you, the example is running using the ordinary list monad. Since it's written in terms of MonadPlus it will work for any (MonadTrans t, MonadPlus (t m)) => t m as well.