{-# LANGUAGE TemplateHaskell #-}
module Control.Monad.Free.Foil.Binary.TH (deriveBinaryPattern) where
import Control.Monad (unless, zipWithM)
import Data.Binary (Binary (..))
import Data.Binary.Get (Get, getWord8)
import Data.Binary.Put (putWord8)
import qualified Data.Map as Map
import Language.Haskell.TH
import Unsafe.Coerce (unsafeCoerce)
deriveBinaryPattern :: Name -> Q [Dec]
deriveBinaryPattern :: Name -> Q [Dec]
deriveBinaryPattern Name
tyName = do
info <- Name -> Q Info
reify Name
tyName
(tvs, cons) <- case info of
TyConI (DataD Cxt
_ Name
_ [TyVarBndr BndrVis]
tvs Maybe Kind
_ [Con]
cons [DerivClause]
_) -> ([TyVarBndr BndrVis], [Con]) -> Q ([TyVarBndr BndrVis], [Con])
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([TyVarBndr BndrVis]
tvs, [Con]
cons)
TyConI (NewtypeD Cxt
_ Name
_ [TyVarBndr BndrVis]
tvs Maybe Kind
_ Con
con [DerivClause]
_) -> ([TyVarBndr BndrVis], [Con]) -> Q ([TyVarBndr BndrVis], [Con])
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([TyVarBndr BndrVis]
tvs, [Con
con])
Info
_ -> String -> Q ([TyVarBndr BndrVis], [Con])
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"deriveBinaryPattern: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Name -> String
forall a. Show a => a -> String
show Name
tyName String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" is not a data type")
unless (length tvs >= 2) $
fail "deriveBinaryPattern: expected a type of kind ... -> S -> S -> Type"
params <- mapM (\Int
i -> String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName (String
"p" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show (Int
i :: Int))) [1 .. length tvs - 2]
nVar <- newName "n"
lVar <- newName "l"
ctors <- concat <$> mapM flatten cons
unless (length ctors <= 256) $
fail "deriveBinaryPattern: more than 256 constructors"
putClauses <- zipWithM (putClause) [0 ..] ctors
getMatches <- zipWithM (getMatch params nVar) [0 ..] ctors
tagName <- newName "tag"
let headTy = (Kind -> Kind -> Kind) -> Kind -> Cxt -> Kind
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Kind -> Kind -> Kind
AppT (Name -> Kind
ConT Name
tyName) ((Name -> Kind) -> [Name] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map Name -> Kind
VarT ([Name]
params [Name] -> [Name] -> [Name]
forall a. Semigroup a => a -> a -> a
<> [Name
nVar, Name
lVar]))
context = [Kind -> Kind -> Kind
AppT (Name -> Kind
ConT ''Binary) (Name -> Kind
VarT Name
p) | Name
p <- [Name]
params]
failMatch = Pat -> Body -> [Dec] -> Match
Match Pat
WildP
(Exp -> Body
NormalB (Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'fail) (Lit -> Exp
LitE (String -> Lit
StringL String
"unknown pattern tag")))) []
getBody = Maybe Exp -> Exp -> Maybe Exp -> Exp
InfixE (Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
VarE 'getWord8)) (Name -> Exp
VarE '(>>=))
(Exp -> Maybe Exp
forall a. a -> Maybe a
Just ([Pat] -> Exp -> Exp
LamE [Name -> Pat
VarP Name
tagName]
(Exp -> [Match] -> Exp
CaseE (Name -> Exp
VarE Name
tagName) ([Match]
getMatches [Match] -> [Match] -> [Match]
forall a. Semigroup a => a -> a -> a
<> [Match
failMatch]))))
pure
[ InstanceD Nothing context (AppT (ConT ''Binary) headTy)
[ FunD 'put putClauses
, ValD (VarP 'get) (NormalB getBody) []
]
]
where
flatten :: Con -> f [(Name, Cxt, Kind)]
flatten (ForallC [TyVarBndr Specificity]
_ Cxt
_ Con
con) = Con -> f [(Name, Cxt, Kind)]
flatten Con
con
flatten (GadtC [Name]
names [BangType]
bts Kind
t) = [(Name, Cxt, Kind)] -> f [(Name, Cxt, Kind)]
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [(Name
c, (BangType -> Kind) -> [BangType] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map BangType -> Kind
forall a b. (a, b) -> b
snd [BangType]
bts, Kind
t) | Name
c <- [Name]
names]
flatten Con
_ =
String -> f [(Name, Cxt, Kind)]
forall a. String -> f a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"deriveBinaryPattern: expected GADT constructors (a pattern's indices vary per constructor)"
putClause :: Integer -> (Name, t a, c) -> m Clause
putClause Integer
tag (Name
cname, t a
fields, c
_) = do
args <- (Int -> m Name) -> [Int] -> m [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\Int
i -> String -> m Name
forall (m :: * -> *). Quote m => String -> m Name
newName (String
"x" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show (Int
i :: Int))) [Int
1 .. t a -> Int
forall a. t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
fields]
let puts = Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'putWord8) (Lit -> Exp
LitE (Integer -> Lit
IntegerL Integer
tag))
Exp -> [Exp] -> [Exp]
forall a. a -> [a] -> [a]
: [Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'put) (Name -> Exp
VarE Name
a) | Name
a <- [Name]
args]
pure (Clause [ConP cname [] (map VarP args)]
(NormalB (AppE (VarE 'mconcat) (ListE puts))) [])
getMatch :: [Name] -> Name -> Integer -> (Name, Cxt, Kind) -> f Match
getMatch [Name]
params Name
nVar Integer
tag (Name
cname, Cxt
fields, Kind
result) = do
let (Kind
_, Cxt
resultArgs) = Kind -> (Kind, Cxt)
unfoldApps Kind
result
paramPairs :: [(Name, Kind)]
paramPairs =
[ (Name
v, Name -> Kind
VarT Name
p)
| (VarT Name
v, Name
p) <- Cxt -> [Name] -> [(Kind, Name)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Int -> Cxt -> Cxt
forall a. Int -> [a] -> [a]
take ([Name] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Name]
params) Cxt
resultArgs) [Name]
params ]
subst :: Map Name Kind
subst = [(Name, Kind)] -> Map Name Kind
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Name, Kind)]
paramPairs
substTy :: Kind -> Kind
substTy Kind
t = case Kind
t of
VarT Name
v -> Kind -> Name -> Map Name Kind -> Kind
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault (Name -> Kind
VarT Name
nVar) Name
v Map Name Kind
subst
AppT Kind
f Kind
x -> Kind -> Kind -> Kind
AppT (Kind -> Kind
substTy Kind
f) (Kind -> Kind
substTy Kind
x)
SigT Kind
x Kind
k -> Kind -> Kind -> Kind
SigT (Kind -> Kind
substTy Kind
x) Kind
k
ParensT Kind
x -> Kind -> Kind
ParensT (Kind -> Kind
substTy Kind
x)
Kind
_ -> Kind
t
getField :: Kind -> Exp
getField Kind
ft = Exp -> Kind -> Exp
SigE (Name -> Exp
VarE 'get) (Kind -> Kind -> Kind
AppT (Name -> Kind
ConT ''Get) (Kind -> Kind
substTy Kind
ft))
chain :: Exp
chain = case Cxt
fields of
[] -> Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'pure) (Name -> Exp
ConE Name
cname)
(Kind
f : Cxt
fs) -> (Exp -> Kind -> Exp) -> Exp -> Cxt -> Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
(\Exp
acc Kind
ft -> Maybe Exp -> Exp -> Maybe Exp -> Exp
InfixE (Exp -> Maybe Exp
forall a. a -> Maybe a
Just Exp
acc) (Name -> Exp
VarE '(<*>)) (Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Kind -> Exp
getField Kind
ft)))
(Maybe Exp -> Exp -> Maybe Exp -> Exp
InfixE (Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Name -> Exp
ConE Name
cname)) (Name -> Exp
VarE '(<$>)) (Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Kind -> Exp
getField Kind
f)))
Cxt
fs
diagTy :: Kind
diagTy = Kind -> Kind -> Kind
AppT (Name -> Kind
ConT ''Get) (Kind -> Kind
substTy Kind
result)
body :: Exp
body = Exp -> Exp -> Exp
AppE (Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'fmap) (Name -> Exp
VarE 'unsafeCoerce)) (Exp -> Kind -> Exp
SigE Exp
chain Kind
diagTy)
Match -> f Match
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Pat -> Body -> [Dec] -> Match
Match (Lit -> Pat
LitP (Integer -> Lit
IntegerL Integer
tag)) (Exp -> Body
NormalB Exp
body) [])
unfoldApps :: Kind -> (Kind, Cxt)
unfoldApps = Cxt -> Kind -> (Kind, Cxt)
go []
where
go :: Cxt -> Kind -> (Kind, Cxt)
go Cxt
args (AppT Kind
f Kind
x) = Cxt -> Kind -> (Kind, Cxt)
go (Kind
x Kind -> Cxt -> Cxt
forall a. a -> [a] -> [a]
: Cxt
args) Kind
f
go Cxt
args Kind
t = (Kind
t, Cxt
args)