Specializing bind for monads over special typeclasses in Haskell

Question

In the second last chapter For a Few Monads More of the very nice tutorial "Learn You a Haskell for a Great Good" the author defines the following monad:

import Data.Ratio  
newtype Prob a = Prob { getProb :: [(a,Rational)] } deriving Show
flatten :: Prob (Prob a) -> Prob a  
flatten (Prob xs) = Prob $ concat $ map multAll xs  
  where multAll (Prob innerxs,p) = map (\(x,r) -> (x,p*r)) innerxs
instance Monad Prob where  
  return x = Prob [(x,1%1)]  
  m >>= f = flatten (fmap f m)  
  fail _ = Prob []

I wondered if it is possible in Haskell to specialize the bind operator ">>=" in case the value in the monad belongs to a special typeclass like Eq, as I'd like to add up all probabilities belonging to the same value.

Answer 1

This is called a "restricted monad" and you define it like this:

{-# LANGUAGE ConstraintKinds, TypeFamilies, KindSignatures, FlexibleContexts, UndecidableInstances #-}
module Control.Restricted (RFunctor(..),
                           RApplicative(..),
                           RMonad(..),
                           RMonadPlus(..),) where
import Prelude hiding (Functor(..), Monad(..))
import Data.Foldable (Foldable(foldMap))
import GHC.Exts (Constraint)

class RFunctor f where
    type Restriction f a :: Constraint
    fmap :: (Restriction f a, Restriction f b) => (a -> b) -> f a -> f b

class (RFunctor f) => RApplicative f where
    pure :: (Restriction f a) => a -> f a
    (<*>) :: (Restriction f a, Restriction f b) => f (a -> b) -> f a -> f b

class (RApplicative m) => RMonad m where
    (>>=) :: (Restriction m a, Restriction m b) => m a -> (a -> m b) -> m b
    (>>) :: (Restriction m a, Restriction m b)  => m a -> m b ->  m b
    a >> b = a >>= \_ -> b
    join :: (Restriction m a, Restriction m (m a)) => m (m a) -> m a
    join a = a >>= id
    fail :: (Restriction m a) => String -> m a
    fail = error

return :: (RMonad m, Restriction m a) => a -> m a
return = pure

class (RMonad m) => RMonadPlus m where
    mplus :: (Restriction m a) => m a -> m a -> m a
    mzero :: (Restriction m a) => m a
    msum :: (Restriction m a, Foldable t) => t (m a) -> m a
    msum t = getRMonadPlusMonoid $ foldMap RMonadPlusMonoid t

data RMonadPlusMonoid m a = RMonadPlusMonoid { getRMonadPlusMonoid :: m a }

instance (RMonadPlus m, Restriction m a) => Monoid (RMonadPlusMonoid m a) where
    mappend (RMonadPlusMonoid x) (RMonadPlusMonoid y) = RMonadPlusMonoid $ mplus x y
    mempty = RMonadPlusMonoid mzero
    mconcat t = RMonadPlusMonoid . msum $ map getRMonadPlusMonoid t

guard :: (RMonadPlus m, Restriction m a) => Bool -> m ()
guard p = if p then return () else mzero

To use a restricted monad, you need to begin your file like this:

{-# LANGUAGE ConstraintKinds, TypeFamilies, RebindableSyntax #-}
module {- module line -} where
import Prelude hiding (Functor(..), Monad(..))
import Control.Restricted

Answer 2

Thanks to Ptharien's Flame's answer (please upvote it!) I managed to adapt the example monad from "Learn You a Haskell for a Great Good" running. As I had to google some details (being a Haskell-newbie) here is what I did at the end (the example flipThree in "Learn..." gives now [(True,9 % 40), (False,31 % 40)]):

file Control/Restricted.hs (to shorten it I removed RApplicative, RMonadPlus etc):

{-# LANGUAGE ConstraintKinds, TypeFamilies, KindSignatures, FlexibleContexts, UndecidableInstances #-}
module Control.Restricted (RFunctor(..),
                           RMonad(..)) where
import Prelude hiding (Functor(..), Monad(..))
import Data.Foldable (Foldable(foldMap))
import Data.Monoid
import GHC.Exts (Constraint)

class RFunctor f where
  type Restriction f a :: Constraint
  fmap :: (Restriction f a, Restriction f b) => (a -> b) -> f a -> f b

class (RFunctor m) => RMonad m where
  return :: (Restriction m a) => a -> m a
  (>>=) :: (Restriction m a, Restriction m b) => m a -> (a -> m b) -> m b
  (>>) :: (Restriction m a, Restriction m b)  => m a -> m b -> m b
  a >> b = a >>= \_ -> b
  join :: (Restriction m a, Restriction m (m a)) => m (m a) -> m a
  join a = a >>= id
  fail :: (Restriction m a) => String -> m a
  fail = error

file Prob.hs:

{-# LANGUAGE ConstraintKinds, TypeFamilies, RebindableSyntax, FlexibleContexts #-}
import Data.Ratio
import Control.Restricted
import Prelude hiding (Functor(..), Monad(..))
import Control.Arrow (first, second)
import Data.List (all)

newtype Prob a = Prob { getProb :: [(a, Rational)] } deriving Show

instance RFunctor Prob where
  type Restriction Prob a = (Eq a)
  fmap f (Prob as) = Prob $ map (first f) as

flatten :: Prob (Prob a) -> Prob a
flatten (Prob xs) = Prob $ concat $ map multAll xs
  where multAll (Prob innerxs, p) = map (\(x, r) -> (x, p*r)) innerxs

compress :: Eq a => Prob a -> Prob a
compress (Prob as) = Prob $ foldr f [] as
  where f a [] = [a]
        f (a, p) ((b, q):bs) | a == b    = (a, p+q):bs
                             | otherwise = (b, q):f (a, p) bs

instance Eq a => Eq (Prob a) where
  (==) (Prob as) (Prob bs) = all (`elem` bs) as

instance RMonad Prob where
  return x = Prob [(x, 1%1)]
  m >>= f = compress $ flatten (fmap f m)
  fail _ = Prob []

Answer 3

Here another possibility based on Generalized Algebraic Datatypes using a technique by Ganesh Sittampalam:

{-# LANGUAGE GADTs #-}

import Control.Arrow (first, second)
import Data.Ratio
import Data.List (foldl')

-- monads over typeclass Eq
class EqMonad m where
  eqReturn :: Eq a => a -> m a
  eqBind :: (Eq a, Eq b) => m a -> (a -> m b) -> m b
  eqFail :: Eq a => String -> m a
  eqFail = error

data AsMonad m a where
  Embed :: (EqMonad m, Eq a) => m a -> AsMonad m a
  Return :: EqMonad m => a -> AsMonad m a
  Bind :: EqMonad m => AsMonad m a -> (a -> AsMonad m b) -> AsMonad m b

instance EqMonad m => Monad (AsMonad m) where
  return = Return
  (>>=) = Bind
  fail = error

unEmbed :: Eq a => AsMonad m a -> m a
unEmbed (Embed m) = m
unEmbed (Return v) = eqReturn v
unEmbed (Bind (Embed m) f) = m `eqBind` (unEmbed . f)
unEmbed (Bind (Return v) f) = unEmbed (f v)
unEmbed (Bind (Bind m f) g) = unEmbed (Bind m (\x -> Bind (f x) g))

-- the example monad from "Learn you a Haskell for a Great good"
newtype Prob a = Prob { getProb :: [(a, Rational)] }
  deriving Show

instance Functor Prob where
  fmap f (Prob as) = Prob $ map (first f) as

flatten :: Prob (Prob a) -> Prob a
flatten (Prob xs) = Prob $ concat $ map multAll xs
  where multAll (Prob innerxs, p) = map (\(x, r) -> (x, p*r)) innerxs

compress :: Eq a => Prob a -> Prob a
compress (Prob as) = Prob $ foldl' f [] as
  where f [] a = [a]
        f ((b, q):bs) (a, p) | a == b    = (a, p+q):bs
                             | otherwise = (b, q):f bs (a, p)

instance Eq a => Eq (Prob a) where
  (==) (Prob as) (Prob bs) = all (`elem` bs) as

instance EqMonad Prob where
  eqReturn x = Prob [(x, 1%1)]
  m `eqBind` f = compress $ flatten (fmap f m)
  eqFail _ = Prob []

newtype Probability a = Probability { getProbability :: AsMonad Prob a }

instance Monad Probability where
  return = Probability . Return
  a >>= f = Probability $ Bind (getProbability a) (getProbability . f)
  fail = error

instance (Show a, Eq a) => Show (Probability a) where
  show = show . getProb . unEmbed . getProbability

-- Example flipping four coins (now as 0/1)
prob :: Eq a => [(a, Rational)] -> Probability a
prob = Probability . Embed . Prob

coin :: Probability Int
coin = prob [(0, 1%2), (1, 1%2)]

loadedCoin :: Probability Int
loadedCoin = prob [(0, 1%10), (1, 9%10)]

flipFour :: Probability Int
flipFour = do
  a <- coin
  b <- coin
  c <- coin
  d <- loadedCoin
  return (a+b+c+d)