Monad to sample without replacement

Question

I created a monad in Haskell that lets you sample without replacement from user-defined urns, and then at the end gives you a list of all possible outcomes. It looks like it's similar to the list monad (except that one only ever samples with replacement), and to the ST monad. Here's the interface I want to present:

data Draw s a
data Urn s a
instance MonadPlus (Draw s)
newUrn :: [a] -> Draw s (Urn s a)
drawFrom :: Urn s a -> Draw s a
drawList :: [a] -> Draw s a -- so that you can still sample with replacement, like in the list monad
runDraw :: (forall s. Draw s a) -> [a]

Here's an example of how I want to use it:

runDraw $ do
  l <- newUrn [1,2,3,3]
  x <- drawFrom l
  y <- drawFrom l
  return (x, y)

-- produces [(1,2),(1,3),(1,3),(2,1),(2,3),(2,3),(3,1),(3,2),(3,3),(3,1),(3,2),(3,3)]

And here's what I came up with to implement that:

{-# LANGUAGE MagicHash, RankNTypes, RoleAnnotations #-}

module Draw (Draw, Urn, newUrn, drawFrom, drawList, runDraw) where

import Control.Applicative (Alternative(..))
import Control.Monad (MonadPlus, ap, liftM)
import Data.List (genericSplitAt)
import GHC.Exts (Any, unsafeCoerce#)
import Numeric.Natural (Natural)

newtype Draw s a = Draw { unDraw :: (Natural, [Any]) -> [(a, (Natural, [Any]))] }
type role Draw nominal representational

newtype Urn s a = Urn Natural
type role Urn nominal representational

instance Functor (Draw s) where
  fmap = liftM

instance Applicative (Draw s) where
  pure = drawList . pure
  (<*>) = ap

instance Alternative (Draw s) where
  empty = drawList empty
  Draw m1 <|> Draw m2 = Draw $ \s -> m1 s <|> m2 s

instance Monad (Draw s) where
  Draw m >>= f = Draw $ \s -> m s >>= uncurry (unDraw . f)

instance MonadPlus (Draw s)

drawList :: [a] -> Draw s a
drawList xs = Draw $ \s -> flip (,) s <$> xs

runDraw :: (forall s. Draw s a) -> [a]
runDraw (Draw f) = map fst (f (0, []))

newUrn :: [a] -> Draw s (Urn s a)
newUrn xs = Draw $ \(n, us) -> pure (Urn n, (n + 1, us ++ [toAny xs]))

drawFrom :: Urn s a -> Draw s a
drawFrom (Urn i) = Draw go where
  go :: (Natural, [Any]) -> [(a, (Natural, [Any]))]
  go (n, us) = map (\(x, remainingContents) -> (x, (n, before ++ toAny remainingContents : after))) (removeEach (fromAny urnContents)) where
    (before, urnContents:after) = genericSplitAt i us

fromAny :: Any -> [a]
fromAny = unsafeCoerce#

toAny :: [a] -> Any
toAny = unsafeCoerce#

removeEach :: [a] -> [(a, [a])]
removeEach [] = []
removeEach (x:xs) = (x, xs):map (fmap (x:)) (removeEach xs)

This seems to work, at least with the example I posted above.

Here's my concerns:

• I'm doing a lot of unsafeCoerce#, which is obviously not very safe
• (before, urnContents:after) = genericSplitAt i us is an incomplete pattern match, which may be able to fail at runtime
• I'm building the list of urns with xs ++ [x], which is quadratically slow
• I'm not confident that this satisfies all of the typeclass laws, in particular the monad law of associativity I now realize that my type is isomorphic to StateT (Natural, [Any]) [], with equivalent instances, so I'm no longer concerned about this.
• I'm not sure if the way I'm handling the urns is correct, or if it's somehow possible to use an urn where it doesn't belong and thus break type safety

Answer 1

This is an interesting problem, and you've written an interesting solution.

With respect to your inefficient urn "store" -- the mapping of immutable urn references (Urn Natural) to mutable urn contents -- it might be worth considering that because of the nature of your monad, most monadic computations involving urns will scale exponentially in the number of urns anyway, so big-O performance of urn list building and lookups is essentially irrelevant. You can start thinking about it when people want to use your monad for 100000-urn problems; or you could probably port everything over to a Data.Map Int or Data.IntMap in a few minutes).

The bigger problem, as you've noted, is that because this all has to run in a specific monotyped monad, unless you want to pre-declare the set of urns and their element types as used in a particular computation, you need an ugly, unsafe generic type like [Any] to represent your set of urns.

One method of dealing with this would be to represent the mutable contents of an urn by a set of always-integer indices while packaging the actual elements as part of the immutable Urn reference. That is, the Urn references you pass around can be represented as:

data Urn s a = Urn { tag :: Key
                   , labels :: Int -> a }
type role Urn nominal representational

with monotyped mutable state:

data UrnState = UrnState { nextTag :: Key
                         , urns :: IntMap [Int] }

So urns urnState ! tag1 is the set of integer indices still in play for that urn, and the actual elements are available by looking up those indices in the labels urnRef map.

SPOILERS

A complete code example, which seems to work on your test case is:

{-# LANGUAGE DeriveFunctor, RoleAnnotations, RankNTypes #-}
import Data.List
import Control.Monad
import qualified Data.IntMap as IntMap
import Data.IntMap (Key, IntMap, (!))

data Urn s a = Urn { tag :: Key
                   , labels :: Int -> a }
type role Urn nominal representational

data UrnState = UrnState { nextTag :: Key
                         , urns :: IntMap [Int] }

newtype Draw s a = Draw { unDraw :: UrnState -> [(a, UrnState)] } deriving (Functor)
type role Draw nominal representational
instance Applicative (Draw s) where
  pure x = Draw (\s -> [(x,s)])
  (<*>) = ap
instance Monad (Draw s) where
  Draw d >>= f = Draw $ \s -> do  -- list monad
    (a', s') <- d s
    unDraw (f a') s'

evalDraw :: (forall s. Draw s a) -> [a]
evalDraw (Draw d) = map fst $ d $ UrnState 0 IntMap.empty

newUrn :: [a] -> Draw s (Urn s a)
newUrn xs = Draw $ \(UrnState nxttag urs) ->
  let -- list of labels keyed by indexes [0..n-1]
      lbls = IntMap.fromAscList (zip [0..] xs)
      -- new urn has tag "nxttag" and the immutable labelling function
      u = Urn nxttag (lbls !)
      -- add urn to state
      urs' = IntMap.insert nxttag (IntMap.keys lbls) urs
  in  [(u, UrnState (nxttag+1) urs')]

draws :: [a] -> [(a,[a])]
draws xs = zipWith3 go (inits xs) xs (tail (tails xs))
  where go l a r = (a, l++r)

drawFrom :: Urn s a -> Draw s a
drawFrom (Urn tg lbls) = Draw $ \(UrnState nxttag urs) ->
  case urs ! tg of
    [] -> fail "empty urn"
    xs -> do  -- list monad
      (a, xs') <- draws xs
      return $ (lbls a, UrnState nxttag $ IntMap.insert tg xs' urs)

main :: IO ()
main = print $ evalDraw $ do
  l <- newUrn [1,2,3,3]
  x <- drawFrom l
  y <- drawFrom l
  return (x, y)