{-# LANGUAGE CPP             #-}
{-# LANGUAGE LambdaCase      #-}
{-# LANGUAGE TemplateHaskell #-}

module Data.InvertibleGrammar.TH where

#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ < 710
import Control.Applicative
#endif
import Data.Foldable (toList)
import Data.InvertibleGrammar.Base
import Data.Maybe
import Data.Text (pack)
import Language.Haskell.TH as TH
import Data.Set (Set)
import qualified Data.Set as S
#if !MIN_VERSION_base(4,11,0)
import Data.Semigroup ((<>))
#endif


{- | Build a prism and the corresponding grammar that will match on the
     given constructor and convert it to reverse sequence of :- stacks.

     E.g. consider a data type:

     > data FooBar a b c = Foo a b c | Bar

     For constructor Foo

     > fooGrammar = $(grammarFor 'Foo)

     will expand into

     > fooGrammar = PartialIso
     >   (\(c :- b :- a :- t) -> Foo a b c :- t)
     >   (\case { Foo a b c :- t -> Just $ c :- b :- a :- t; _ -> Nothing })

     Note the order of elements on the stack:

     > ghci> :t fooGrammar
     > fooGrammar :: Grammar p (c :- (b :- (a :- t))) (FooBar a b c :- t)
-}

grammarFor :: Name -> ExpQ
grammarFor :: Name -> ExpQ
grammarFor Name
constructorName = do
#if defined(__GLASGOW_HASKELL__)
# if __GLASGOW_HASKELL__ <= 710
  DataConI realConstructorName _typ parentName _fixity <- reify constructorName
# else
  DataConI Name
realConstructorName Type
_typ Name
parentName <- Name -> Q Info
reify Name
constructorName
# endif
#endif
  TyConI Dec
dataDef <- Name -> Q Info
reify Name
parentName

  let Just (Bool
single, Con
constructorInfo) = do
        (Bool
single, [Con]
allConstr) <- Dec -> Maybe (Bool, [Con])
constructors Dec
dataDef
        Con
constr <- Name -> [Con] -> Maybe Con
findConstructor Name
realConstructorName [Con]
allConstr
        forall (m :: * -> *) a. Monad m => a -> m a
return (Bool
single, Con
constr)

  let ts :: [Type]
ts = Con -> [Type]
fieldTypes Con
constructorInfo
  [Name]
vs <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall a b. a -> b -> a
const forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => String -> m Name
newName String
"x") [Type]
ts
  Name
t <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"t"

  let matchStack :: [Name] -> m Pat
matchStack []      = forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
t
      matchStack (Name
_v:[Name]
vs) = [p| $(varP _v) :- $_vs' |]
        where
          _vs' :: m Pat
_vs' = [Name] -> m Pat
matchStack [Name]
vs
      fPat :: Q Pat
fPat  = forall {m :: * -> *}. Quote m => [Name] -> m Pat
matchStack [Name]
vs
      buildConstructor :: ExpQ
buildConstructor = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Name
v ExpQ
acc -> forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE ExpQ
acc (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
v)) (forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
realConstructorName) [Name]
vs
      fBody :: ExpQ
fBody = [e| $buildConstructor :- $(varE t) |]
      fFunc :: ExpQ
fFunc = forall (m :: * -> *). Quote m => [m Pat] -> m Exp -> m Exp
lamE [Q Pat
fPat] ExpQ
fBody

  let gPat :: Q Pat
gPat  = [p| $_matchConsructor :- $(varP t) |]
        where
          _matchConsructor :: Q Pat
_matchConsructor = forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
realConstructorName (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Pat
varP (forall a. [a] -> [a]
reverse [Name]
vs))
      gBody :: ExpQ
gBody = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Name
v ExpQ
acc -> [e| $(varE v) :- $acc |]) (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
t) [Name]
vs
      gFunc :: ExpQ
gFunc = forall (m :: * -> *). Quote m => [m Match] -> m Exp
lamCaseE forall a b. (a -> b) -> a -> b
$ forall a. [Maybe a] -> [a]
catMaybes
        [ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
TH.match Q Pat
gPat (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [e| Right ($gBody) |]) []
        , if Bool
single
          then forall a. Maybe a
Nothing
          else forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
TH.match forall (m :: * -> *). Quote m => m Pat
wildP (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [e| Left (expected $ "constructor " <> pack ( $(stringE (show constructorName))) ) |]) []
        ]

  [e| PartialIso $fFunc $gFunc |]


{- | Build prisms and corresponding grammars for all data constructors of given
     type. Expects grammars to zip built ones with.

     > $(match ''Maybe)

     Will expand into a lambda:

     > (\nothingG justG -> ($(grammarFor 'Nothing) . nothingG) <>
     >                     ($(grammarFor 'Just)    . justG))
-}
match :: Name -> ExpQ
match :: Name -> ExpQ
match Name
tyName = do
  [Name]
names  <- forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall (t :: * -> *) a. Foldable t => t a -> [a]
toList forall b c a. (b -> c) -> (a -> b) -> a -> c
. Con -> Set Name
constructorNames) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Info -> Q [Con]
extractConstructors forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Name -> Q Info
reify Name
tyName)
  [Name]
argTys <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\Name
_ -> forall (m :: * -> *). Quote m => String -> m Name
newName String
"a") [Name]
names
  let grammars :: [ExpQ]
grammars = forall a b. (a -> b) -> [a] -> [b]
map (\(Name
con, Name
arg) -> [e| $(varE arg) $(grammarFor con) |]) (forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
names [Name]
argTys)
  forall (m :: * -> *). Quote m => [m Pat] -> m Exp -> m Exp
lamE (forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Pat
varP [Name]
argTys) (forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 (\ExpQ
e1 ExpQ
e2 -> [e| $e1 <> $e2 |]) [ExpQ]
grammars)
  where
    extractConstructors :: Info -> Q [Con]
    extractConstructors :: Info -> Q [Con]
extractConstructors (TyConI Dec
dataDef) =
      case Dec -> Maybe (Bool, [Con])
constructors Dec
dataDef of
        Just (Bool
_, [Con]
cs) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [Con]
cs
        Maybe (Bool, [Con])
Nothing      -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"Data type " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Name
tyName forall a. [a] -> [a] -> [a]
++ String
" defines no constructors"
    extractConstructors Info
_ =
      forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"Data definition expected for name " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Name
tyName

----------------------------------------------------------------------
-- Utils

constructors :: Dec -> Maybe (Bool, [Con])
#if defined(__GLASGOW_HASKELL__)
# if __GLASGOW_HASKELL__ <= 710
constructors (DataD _ _ _ cs _)     = Just (length cs == 1, cs)
constructors (NewtypeD _ _ _ c _)   = Just (True, [c])
# else
constructors :: Dec -> Maybe (Bool, [Con])
constructors (DataD [Type]
_ Name
_ [TyVarBndr ()]
_ Maybe Type
_ [Con]
cs [DerivClause]
_)   = forall a. a -> Maybe a
Just (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Con]
cs forall a. Eq a => a -> a -> Bool
== Int
1, [Con]
cs)
constructors (NewtypeD [Type]
_ Name
_ [TyVarBndr ()]
_ Maybe Type
_ Con
c [DerivClause]
_) = forall a. a -> Maybe a
Just (Bool
True, [Con
c])
# endif
#endif
constructors Dec
_                      = forall a. Maybe a
Nothing

findConstructor :: Name -> [Con] -> Maybe Con
findConstructor :: Name -> [Con] -> Maybe Con
findConstructor Name
_    [] = forall a. Maybe a
Nothing
findConstructor Name
name (Con
c:[Con]
cs)
  | Name
name forall a. Ord a => a -> Set a -> Bool
`S.member` Con -> Set Name
constructorNames Con
c = forall a. a -> Maybe a
Just Con
c
  | Bool
otherwise                          = Name -> [Con] -> Maybe Con
findConstructor Name
name [Con]
cs

constructorNames :: Con -> Set Name
constructorNames :: Con -> Set Name
constructorNames = \case
  NormalC Name
name [BangType]
_   -> forall a. a -> Set a
S.singleton Name
name
  RecC Name
name [VarBangType]
_      -> forall a. a -> Set a
S.singleton Name
name
  InfixC BangType
_ Name
name BangType
_  -> forall a. a -> Set a
S.singleton Name
name
  ForallC [TyVarBndr Specificity]
_ [Type]
_ Con
con' -> Con -> Set Name
constructorNames Con
con'
#if MIN_VERSION_template_haskell(2, 11, 0)
  GadtC [Name]
cs [BangType]
_ Type
_     -> forall a. Ord a => [a] -> Set a
S.fromList [Name]
cs
  RecGadtC [Name]
cs [VarBangType]
_ Type
_  -> forall a. Ord a => [a] -> Set a
S.fromList [Name]
cs
#endif

fieldTypes :: Con -> [Type]
fieldTypes :: Con -> [Type]
fieldTypes = \case
  NormalC Name
_ [BangType]
fieldTypes  -> forall a b. (a -> b) -> [a] -> [b]
map forall {a} {b}. (a, b) -> b
extractType [BangType]
fieldTypes
  RecC Name
_ [VarBangType]
fieldTypes     -> forall a b. (a -> b) -> [a] -> [b]
map forall {a} {b} {c}. (a, b, c) -> c
extractType' [VarBangType]
fieldTypes
  InfixC (Bang
_,Type
a) Name
_b (Bang
_,Type
b) -> [Type
a, Type
b]
  ForallC [TyVarBndr Specificity]
_ [Type]
_ Con
con'      -> Con -> [Type]
fieldTypes Con
con'
#if defined(__GLASGOW_HASKELL__) && __GLASGOW_HASKELL__ >= 800
  GadtC [Name]
_ [BangType]
fs Type
_          -> forall a b. (a -> b) -> [a] -> [b]
map forall {a} {b}. (a, b) -> b
extractType [BangType]
fs
  RecGadtC [Name]
_ [VarBangType]
fs Type
_       -> forall a b. (a -> b) -> [a] -> [b]
map forall {a} {b} {c}. (a, b, c) -> c
extractType' [VarBangType]
fs
#endif
  where
    extractType :: (a, b) -> b
extractType (a
_, b
t) = b
t
    extractType' :: (a, b, c) -> c
extractType' (a
_, b
_, c
t) = c
t