{-# LANGUAGE TemplateHaskell #-}
-- | Derive the 'Binary' instance a client's pattern (binder) type needs,
-- alongside the hand-written instances of "Control.Monad.Free.Foil.Binary".
--
-- A pattern type is a GADT over two scope indices, so its instance cannot
-- come from "GHC.Generics". What the deriver writes is the shape one would
-- write by hand: one tag byte per constructor in declaration order, then
-- the fields in order. Decoding happens /at the diagonal/, with every
-- scope index of a constructor instantiated to the same variable, which any
-- chain of binder indices admits. A single coercion then moves the result to
-- the requested indices. That coercion mints scope evidence, so a derived
-- instance is part of the same trust boundary as the library's own. See the
-- module documentation of "Control.Monad.Free.Foil.Binary".
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' ''Pattern@ writes
-- @instance (Binary p1, …) => Binary (Pattern p1 … n l)@ for a pattern
-- type of kind @… -> S -> S -> Type@ whose fields are parameters, name
-- binders, or nested patterns.
--
-- @since 0.4.0
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))) [])

    -- Decode at the diagonal: the constructor's own scope variables (the
    -- result indices and any intermediates) all become @n@, and its
    -- parameter variables become the instance's parameters. The chain of a
    -- pattern's indices always admits the diagonal. Each field's 'get' is
    -- annotated with the substituted type, pinning the intermediates.
    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)