{-# LANGUAGE TemplateHaskell, DataKinds #-}
module Control.Isomorphism.Partial.TH
  ( constructorIso
  , defineIsomorphisms
  ) where

import           Control.Monad
import           Data.Char                          (toLower)
import           Data.List                          (find)
import           Language.Haskell.TH

import           Control.Isomorphism.Partial.Unsafe (Iso (Iso))

gadtError :: a
gadtError :: forall a. a
gadtError = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Control.Isomorphism.Partial.TH: GADTs currently not supported."
{-# NOINLINE gadtError #-}

-- | Extract the name of a constructor, e.g. ":" or "Just".
conName :: Con -> Name
conName :: Con -> Name
conName (NormalC Name
name [BangType]
_)   =   Name
name
conName (RecC Name
name [VarBangType]
_)      =   Name
name
conName (InfixC BangType
_ Name
name BangType
_)  =   Name
name
conName (ForallC [TyVarBndr Specificity]
_ Cxt
_ Con
con)  =   Con -> Name
conName Con
con
conName (GadtC [Name]
_ [BangType]
_ Type
_)      =   Name
forall a. a
gadtError
conName (RecGadtC [Name]
_ [VarBangType]
_ Type
_)   =   Name
forall a. a
gadtError

-- | Extract the types of the constructor's fields.
conFields :: Con -> [Type]
conFields :: Con -> Cxt
conFields (NormalC Name
_ [BangType]
fields)  =   (BangType -> Type) -> [BangType] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (\(Bang
_, Type
t) -> Type
t) [BangType]
fields
conFields (RecC Name
_ [VarBangType]
fields)     =   (VarBangType -> Type) -> [VarBangType] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (\(Name
_, Bang
_, Type
t) -> Type
t) [VarBangType]
fields
conFields (InfixC BangType
lhs Name
_ BangType
rhs)  =   (BangType -> Type) -> [BangType] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (\(Bang
_, Type
t) -> Type
t) [BangType
lhs, BangType
rhs]
conFields (ForallC [TyVarBndr Specificity]
_ Cxt
_ Con
con)   =   Con -> Cxt
conFields Con
con
conFields (GadtC [Name]
_ [BangType]
_ Type
_)       =   Cxt
forall a. a
gadtError
conFields (RecGadtC [Name]
_ [VarBangType]
_ Type
_)    =   Cxt
forall a. a
gadtError

-- Data dec information
data DecInfo flag = DecInfo Type [TyVarBndr flag] [Con]

-- | Extract data or newtype declaration information
decInfo :: Dec -> Q (DecInfo ())
decInfo :: Dec -> Q (DecInfo ())
decInfo (DataD    Cxt
_ Name
name [TyVarBndr ()]
tyVars Maybe Type
_ [Con]
cs [DerivClause]
_) =  DecInfo () -> Q (DecInfo ())
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (DecInfo () -> Q (DecInfo ())) -> DecInfo () -> Q (DecInfo ())
forall a b. (a -> b) -> a -> b
$ Type -> [TyVarBndr ()] -> [Con] -> DecInfo ()
forall flag. Type -> [TyVarBndr flag] -> [Con] -> DecInfo flag
DecInfo (Name -> Type
ConT Name
name) [TyVarBndr ()]
tyVars [Con]
cs
decInfo (NewtypeD Cxt
_ Name
name [TyVarBndr ()]
tyVars Maybe Type
_ Con
c [DerivClause]
_) =  DecInfo () -> Q (DecInfo ())
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (DecInfo () -> Q (DecInfo ())) -> DecInfo () -> Q (DecInfo ())
forall a b. (a -> b) -> a -> b
$ Type -> [TyVarBndr ()] -> [Con] -> DecInfo ()
forall flag. Type -> [TyVarBndr flag] -> [Con] -> DecInfo flag
DecInfo (Name -> Type
ConT Name
name) [TyVarBndr ()]
tyVars [Con
c]
decInfo Dec
_ = [Char] -> Q (DecInfo ())
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
"partial isomorphisms can only be derived for constructors of data type or newtype declarations."

-- | Convert tyVarBndr to type
tyVarBndrToType :: TyVarBndr () -> Type
tyVarBndrToType :: TyVarBndr () -> Type
tyVarBndrToType (PlainTV Name
n ()
_)   = Name -> Type
VarT Name
n
tyVarBndrToType (KindedTV Name
n ()
_ Type
k) = Type -> Type -> Type
SigT (Name -> Type
VarT Name
n) Type
k

-- | Create Iso type for specified type and conctructor fields (Iso (a, b) (CustomType a b c))
isoType :: Type -> [TyVarBndr ()] -> [Type] -> Q Type
isoType :: Type -> [TyVarBndr ()] -> Cxt -> Q Type
isoType Type
typ [TyVarBndr ()]
tyVarBndrs Cxt
fields = do
    Type
isoCon <- [t| Iso |]
    Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Q Type) -> Type -> Q Type
forall a b. (a -> b) -> a -> b
$ [TyVarBndr Specificity] -> Cxt -> Type -> Type
ForallT ((TyVarBndr () -> TyVarBndr Specificity)
-> [TyVarBndr ()] -> [TyVarBndr Specificity]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr () -> TyVarBndr Specificity
forall {flag}. TyVarBndr flag -> TyVarBndr Specificity
specified [TyVarBndr ()]
tyVarBndrs) [] (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$ Type
isoCon Type -> Type -> Type
`AppT` (Cxt -> Type
isoArgs Cxt
fields) Type -> Type -> Type
`AppT` (Type -> Cxt -> Type
applyAll Type
typ (Cxt -> Type) -> Cxt -> Type
forall a b. (a -> b) -> a -> b
$ (TyVarBndr () -> Type) -> [TyVarBndr ()] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr () -> Type
tyVarBndrToType [TyVarBndr ()]
tyVarBndrs)
    where
      specified :: TyVarBndr flag -> TyVarBndr Specificity
specified (PlainTV Name
name flag
_) = Name -> Specificity -> TyVarBndr Specificity
forall flag. Name -> flag -> TyVarBndr flag
PlainTV Name
name Specificity
SpecifiedSpec
      specified (KindedTV Name
name flag
_ Type
kind) = Name -> Specificity -> Type -> TyVarBndr Specificity
forall flag. Name -> flag -> Type -> TyVarBndr flag
KindedTV Name
name Specificity
SpecifiedSpec Type
kind

isoArgs :: [Type] -> Type
isoArgs :: Cxt -> Type
isoArgs []     = Int -> Type
TupleT Int
0
isoArgs [Type
x]    = Type
x
isoArgs (Type
x:Cxt
xs) = Type -> Type -> Type
AppT (Type -> Type -> Type
AppT (Int -> Type
TupleT Int
2) Type
x) (Cxt -> Type
isoArgs Cxt
xs)

-- | Apply all types to supplied type
applyAll :: Type -> [Type] -> Type
applyAll :: Type -> Cxt -> Type
applyAll = (Type -> Type -> Type) -> Type -> Cxt -> Type
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
AppT

-- | Construct a partial isomorphism expression for a constructor,
-- given the constructor's name.
constructorIso :: Name -> ExpQ
constructorIso :: Name -> ExpQ
constructorIso Name
name = do
  DataConI Name
n Type
_ Name
d    <-  Name -> Q Info
reify Name
name
  TyConI Dec
dec        <-  Name -> Q Info
reify Name
d
  DecInfo Type
_ [TyVarBndr ()]
_ [Con]
cs    <-  Dec -> Q (DecInfo ())
decInfo Dec
dec
  let Just Con
con      =   (Con -> Bool) -> [Con] -> Maybe Con
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find (\Con
c -> Name
n Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Con -> Name
conName Con
c) [Con]
cs
  [MatchQ] -> Con -> ExpQ
isoFromCon ([Con] -> [MatchQ]
wildcard [Con]
cs) Con
con

wildcard :: [Con] -> [MatchQ]
wildcard :: [Con] -> [MatchQ]
wildcard [Con]
cs
  =  if [Con] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Con]
cs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
     then  [Q Pat -> Q Body -> [Q Dec] -> MatchQ
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (Q Pat
forall (m :: * -> *). Quote m => m Pat
wildP) (ExpQ -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [| Nothing |]) []]
     else  []

-- | Converts a constructor name (starting with an upper-case
--   letter) into a function name (starting with a lower-case
--   letter).
rename :: Name -> Name
rename :: Name -> Name
rename Name
n
  = [Char] -> Name
mkName (Char -> Char
toLower Char
c Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: [Char]
cs) where Char
c : [Char]
cs = Name -> [Char]
nameBase Name
n

defineIsomorphisms :: Name -> Q [Dec]
defineIsomorphisms :: Name -> Q [Dec]
defineIsomorphisms Name
d = do
  TyConI Dec
dec  <-  Name -> Q Info
reify Name
d
  DecInfo Type
typ [TyVarBndr ()]
tyVarBndrs [Con]
cs          <-  Dec -> Q (DecInfo ())
decInfo Dec
dec
  [[Dec]] -> [Dec]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ([[Dec]] -> [Dec]) -> Q [[Dec]] -> Q [Dec]
forall a b. (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` (Con -> Q [Dec]) -> [Con] -> Q [[Dec]]
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 (\Con
a -> [MatchQ] -> Type -> [TyVarBndr ()] -> Con -> Q [Dec]
defFromCon ([Con] -> [MatchQ]
wildcard [Con]
cs) Type
typ [TyVarBndr ()]
tyVarBndrs Con
a) [Con]
cs

-- | Constructs a partial isomorphism definition for a
--   constructor, given information about the constructor.
--   The name of the partial isomorphisms is constructed by
--   spelling the constructor name with an initial lower-case
--   letter.
defFromCon :: [MatchQ] -> Type -> [TyVarBndr ()] -> Con -> DecsQ
defFromCon :: [MatchQ] -> Type -> [TyVarBndr ()] -> Con -> Q [Dec]
defFromCon [MatchQ]
matches Type
t [TyVarBndr ()]
tyVarBndrs Con
con = do
    let funName :: Name
funName = Name -> Name
rename (Name -> Name) -> Name -> Name
forall a b. (a -> b) -> a -> b
$ Con -> Name
conName Con
con
    Dec
sig <- Name -> Type -> Dec
SigD Name
funName (Type -> Dec) -> Q Type -> Q Dec
forall a b. (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` Type -> [TyVarBndr ()] -> Cxt -> Q Type
isoType Type
t [TyVarBndr ()]
tyVarBndrs (Con -> Cxt
conFields Con
con)
    Dec
fun <- Name -> [Q Clause] -> Q Dec
forall (m :: * -> *). Quote m => Name -> [m Clause] -> m Dec
funD Name
funName [ [Q Pat] -> Q Body -> [Q Dec] -> Q Clause
forall (m :: * -> *).
Quote m =>
[m Pat] -> m Body -> [m Dec] -> m Clause
clause [] (ExpQ -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB ([MatchQ] -> Con -> ExpQ
isoFromCon [MatchQ]
matches Con
con)) [] ]
    [Dec] -> Q [Dec]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return [Dec
sig, Dec
fun]

-- | Constructs a partial isomorphism expression for a
--   constructor, given information about the constructor.
isoFromCon :: [MatchQ] -> Con -> ExpQ
isoFromCon :: [MatchQ] -> Con -> ExpQ
isoFromCon [MatchQ]
matches Con
con = do
  let c :: Name
c     =   Con -> Name
conName Con
con
  let fs :: Cxt
fs    =   Con -> Cxt
conFields Con
con
  let n :: Int
n     =   Cxt -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Cxt
fs
  ([Q Pat]
ps, [ExpQ]
vs)  <-  Int -> Q ([Q Pat], [ExpQ])
genPE Int
n
  Name
v         <-  [Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"x"
  let f :: ExpQ
f     =   [Q Pat] -> ExpQ -> ExpQ
forall (m :: * -> *). Quote m => [m Pat] -> m Exp -> m Exp
lamE [([Q Pat] -> Q Pat) -> [Q Pat] -> Q Pat
forall t. ([t] -> t) -> [t] -> t
nested [Q Pat] -> Q Pat
forall (m :: * -> *). Quote m => [m Pat] -> m Pat
tupP [Q Pat]
ps]
                  [| Just $((ExpQ -> ExpQ -> ExpQ) -> ExpQ -> [ExpQ] -> ExpQ
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ExpQ -> ExpQ -> ExpQ
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
c) [ExpQ]
vs) |]
  let g :: ExpQ
g     =   [Q Pat] -> ExpQ -> ExpQ
forall (m :: * -> *). Quote m => [m Pat] -> m Exp -> m Exp
lamE [Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
v]
                  (ExpQ -> [MatchQ] -> ExpQ
forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
v) ([MatchQ] -> ExpQ) -> [MatchQ] -> ExpQ
forall a b. (a -> b) -> a -> b
$
                    [ Q Pat -> Q Body -> [Q Dec] -> MatchQ
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (Name -> [Q Pat] -> Q Pat
forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
c [Q Pat]
ps)
                        (ExpQ -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [| Just $(([ExpQ] -> ExpQ) -> [ExpQ] -> ExpQ
forall t. ([t] -> t) -> [t] -> t
nested [ExpQ] -> ExpQ
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
tupE [ExpQ]
vs) |]) []
                    ] [MatchQ] -> [MatchQ] -> [MatchQ]
forall a. [a] -> [a] -> [a]
++ [MatchQ]
matches)
  [| Iso $ExpQ
f $ExpQ
g |]

genPE :: Int -> Q ([PatQ], [ExpQ])
genPE :: Int -> Q ([Q Pat], [ExpQ])
genPE Int
n = do
  [Name]
ids <- Int -> Q Name -> Q [Name]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n ([Char] -> Q Name
forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"x")
  ([Q Pat], [ExpQ]) -> Q ([Q Pat], [ExpQ])
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Name -> Q Pat) -> [Name] -> [Q Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP [Name]
ids, (Name -> ExpQ) -> [Name] -> [ExpQ]
forall a b. (a -> b) -> [a] -> [b]
map Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
ids)

nested :: ([t] -> t) -> [t] -> t
nested :: forall t. ([t] -> t) -> [t] -> t
nested [t] -> t
tup []      =  [t] -> t
tup []
nested [t] -> t
_   [t
x]     =  t
x
nested [t] -> t
tup (t
x:[t]
xs)  =  [t] -> t
tup [t
x, ([t] -> t) -> [t] -> t
forall t. ([t] -> t) -> [t] -> t
nested [t] -> t
tup [t]
xs]