Fast coproducts for Haskell & Eta

Main.hs

#!/usr/bin/env stack
{- stack
     --resolver lts-6.27
     --install-ghc
     runghc
     --package containers
-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE InstanceSigs #-}
{-
This example is compatible with GHC >= 7.10.3 and any version of Eta.

For large scale free monads, the linked list used to store the handlers
for each component of a coproduct can be quite slow. So we instead use
a map to store the handlers. The performance of this approach is O(log n)
vs O(n) for the standard coproduct approach. Hence, the constant
factor is much better for this approach as n becomes large, but
less than 100 in practice.

This example was generalized from:
https://gist.github.com/puffnfresh/9924680

And inspiration for this idea was taken from the README of:
https://github.com/frees-io/iota

You can this this file with:

stack Main.hs

-}
module Main where

import GHC.Base
import Unsafe.Coerce
import Control.Monad
import Data.Typeable
import Data.Map (Map)
import qualified Data.Map as M

-- Free Monad
data Free f a = Free (f (Free f a)) | Pure a

instance Functor f => Functor (Free f) where
  fmap f x = pure f <*> x

instance Functor f => Applicative (Free f) where
  (<*>) = ap
  pure = return

instance Functor f => Monad (Free f) where
    Pure a >>= f = f a
    Free r >>= f = Free (fmap (>>= f) r)
    return = Pure

liftF :: Functor f => f a -> Free f a
liftF = Free . fmap return

-- Coproducts
type family IsElem' (x :: k) (xs :: [k]) where
  IsElem' x '[] = 'False
  IsElem' x (x ': xs) = 'True
  IsElem' x (y ': xs) = IsElem' x xs

type IsElem a as = IsElem' a as ~ 'True

data Coproduct (fs :: [* -> *]) a where
  Inject :: (Typeable f, Functor f, IsElem f fs) => f a -> Coproduct fs a

inject :: (Typeable f, Functor f, IsElem f fs) => f a -> Coproduct fs a
inject = Inject

instance Functor (Coproduct fs) where
  fmap :: (a -> b) -> Coproduct fs a -> Coproduct fs b
  fmap f (Inject fa) = Inject $ fmap f fa

{- Indexed Coproduct
   A map that stores functions that can handle each case of the coproduct.
   This should *NOT* be exposed outside the module. -}
newtype IndexedCoproduct = IndexedCoproduct (Map TyCon Any)

emptyCop = IndexedCoproduct $ M.empty

selectCop tyrep (IndexedCoproduct m) =
  case M.lookup tyrep m of
    Just x -> unsafeCoerce x
    _      -> error "select: Bad lookup!"

insertCop k v (IndexedCoproduct m) =
  IndexedCoproduct $ M.insert k (unsafeCoerce v) m

{- CoproductMap is a typeclass that takes a list of functions that transform
   the components of the coproduct and build a map that efficiently
   indexes each component.

   CoproductH is an associated type family that allows coproduct to be
   variable-argument and with types corresponding to the order of the
   components of the coproduct. -}
class CoproductMap (fs' :: [* -> *]) (fs :: [* -> *]) (a :: *) (b :: *) where
  type CoproductH fs' fs a b

  coproduct' :: Proxy fs -> Proxy b -> IndexedCoproduct -> Coproduct fs' a -> CoproductH fs' fs a b

instance CoproductMap fs' '[] a b where
  type CoproductH fs' '[] a b = b

  coproduct' :: Proxy '[] -> Proxy b -> IndexedCoproduct -> Coproduct fs' a -> b
  coproduct' p1 p2 m (Inject (fa :: f x)) =
    (selectCop (typeOfF (Proxy :: Proxy f)) m) (unsafeCoerce fa)

instance (CoproductMap fs' fs a b, Typeable f) => CoproductMap fs' (f ': fs) a b where
  type CoproductH fs' (f ': fs) a b = (f a -> b) -> CoproductH fs' fs a b

  coproduct' :: Proxy (f ': fs) -> Proxy b -> IndexedCoproduct -> Coproduct fs' a ->
                 (f a -> b) -> CoproductH fs' fs a b
  coproduct' p1 p2 m cop f =
    coproduct' (Proxy :: Proxy fs) p2
                (insertCop (typeOfF (Proxy :: Proxy f)) f m) cop

typeOfF :: forall f. (Typeable f) => Proxy f -> TyCon
typeOfF = typeRepTyCon . typeRep

coproduct :: forall fs a b. (CoproductMap fs fs a b) => Proxy b -> Coproduct fs a -> CoproductH fs fs a b
coproduct p cop = coproduct' (Proxy :: Proxy fs) p emptyCop cop

-- Lifting coproducts

liftCop :: (Typeable f, Functor f, IsElem f fs) => f a -> Free (Coproduct fs) a
liftCop = liftF . inject

-- Actions

data FPrint a = FPrint String a
  deriving Typeable

instance Functor FPrint where
    fmap f (FPrint s a) = FPrint s $ f a

data FRead a = FRead (String -> a)
  deriving Typeable

instance Functor FRead where
    fmap f (FRead g) = FRead $ f . g

fprint s = liftCop $ FPrint s ()
fread    = liftCop $ FRead id

-- Example program

readPrint :: Free (Coproduct '[FPrint, FRead]) ()
readPrint = do
  fprint "Hello, name?"
  name <- fread
  fprint $ "Hi " ++ name ++ "!"

-- Interpreter
runIO :: Free (Coproduct '[FPrint, FRead]) a -> IO a
runIO (Free c) = coproduct
                   (Proxy :: Proxy (IO a))
                   c
                   (\(FPrint s a) -> putStrLn s >> runIO a)
                   (\(FRead f)    -> getLine >>= runIO . f)
runIO (Pure a) = return a

main :: IO ()
main = runIO readPrint

Would it be possible to have this as a library? (I could do it myself) 😄

@NickSeagull Just noticed this comment! Haven't had time to release this as a library - you are free to do so if you wish.