{-# LANGUAGE Safe #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
module Data.Universe.Some.TH (
  DeriveUniverseSome (..),
  universeSomeQ,
  ) where

import Control.Monad (forM, mapM, unless)
import Data.Some (Some, mkSome)
import Data.Universe.Class (Universe (..))
import Data.Universe.Some (UniverseSome (..))
import Data.Universe.Helpers (interleave, (<+*+>))
import Language.Haskell.TH
import Language.Haskell.TH.Datatype

-- $setup
-- >>> :m + Data.Some Data.Universe.Class Data.Universe.Some
-- >>> import Language.Haskell.TH (DecsQ)

-- | Derive the @'UniverseSome' n@ instance.
--
-- >>> :set -XGADTs -XTemplateHaskell -XStandaloneDeriving
-- >>> import Data.Universe.Class (universe)
-- >>> import Data.GADT.Show
--
-- >>> data Tag b a where IntTag :: Tag b Int; BoolTag :: b -> Tag b Bool
-- >>> deriving instance Show b => Show (Tag b a)
-- >>> instance Show b => GShow (Tag b) where gshowsPrec = showsPrec
--
-- (@data Unused@ is to workaround bug in older GHCi)
-- >>> data Unused; $(deriveUniverseSome ''Tag)
--
-- >>> universe :: [Some (Tag (Maybe Bool))]
-- [Some IntTag,Some (BoolTag Nothing),Some (BoolTag (Just False)),Some (BoolTag (Just True))]
--
-- 'deriveUniverseSome' variant taking a 'Name' guesses simple class constraints.
-- If you need more specific, you can specify them:
-- (Note: on older GHCs this will warn, as the instance definition doesn't have all methods defined).
--
-- >>> data Tag b a where IntTag :: Tag b Int; BoolTag :: b -> Tag b Bool
-- >>> deriving instance Show b => Show (Tag b a)
-- >>> instance Show b => GShow (Tag b) where gshowsPrec = showsPrec
-- >>> data Unused; $(deriveUniverseSome ([d| instance Universe b => UniverseSome (Tag b) |] :: DecsQ))
-- ...
-- >>> universe :: [Some (Tag (Maybe Bool))]
-- [Some IntTag,Some (BoolTag Nothing),Some (BoolTag (Just False)),Some (BoolTag (Just True))]
--
class DeriveUniverseSome a where
  deriveUniverseSome :: a -> DecsQ

instance DeriveUniverseSome a => DeriveUniverseSome [a] where
  deriveUniverseSome :: [a] -> DecsQ
deriveUniverseSome [a]
a = ([[Dec]] -> [Dec]) -> Q [[Dec]] -> DecsQ
forall a b. (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ((a -> DecsQ) -> [a] -> 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 a -> DecsQ
forall a. DeriveUniverseSome a => a -> DecsQ
deriveUniverseSome [a]
a)

instance DeriveUniverseSome a => DeriveUniverseSome (Q a) where
  deriveUniverseSome :: Q a -> DecsQ
deriveUniverseSome Q a
a = a -> DecsQ
forall a. DeriveUniverseSome a => a -> DecsQ
deriveUniverseSome (a -> DecsQ) -> Q a -> DecsQ
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Q a
a

instance DeriveUniverseSome Name where
  deriveUniverseSome :: Name -> DecsQ
deriveUniverseSome Name
name = do
    DatatypeInfo
di <- Name -> Q DatatypeInfo
reifyDatatype Name
name
    let DatatypeInfo { datatypeContext :: DatatypeInfo -> Cxt
datatypeContext = Cxt
ctxt
                     , datatypeName :: DatatypeInfo -> Name
datatypeName    = Name
parentName
                     , datatypeInstTypes :: DatatypeInfo -> Cxt
datatypeInstTypes = Cxt
vars0
                     , datatypeCons :: DatatypeInfo -> [ConstructorInfo]
datatypeCons    = [ConstructorInfo]
cons
                     } = DatatypeInfo
di

    case Cxt -> Maybe (Cxt, Type)
forall a. [a] -> Maybe ([a], a)
safeUnsnoc Cxt
vars0 of
      Maybe (Cxt, Type)
Nothing -> String -> DecsQ
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Datatype should have at least one type variable"
      Just (Cxt
vars, Type
var) -> do
        [Name]
varNames <- Cxt -> (Type -> Q Name) -> Q [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM Cxt
vars ((Type -> Q Name) -> Q [Name]) -> (Type -> Q Name) -> Q [Name]
forall a b. (a -> b) -> a -> b
$ \Type
v -> case Type
v of
          SigT (VarT Name
n) Type
StarT -> String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"x"
          Type
_                   -> String -> Q Name
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Only arguments of kind Type are supported"

        let constrs :: [TypeQ]
            constrs :: [TypeQ]
constrs = (Name -> TypeQ) -> [Name] -> [TypeQ]
forall a b. (a -> b) -> [a] -> [b]
map (\Name
n -> Name -> TypeQ
forall (m :: * -> *). Quote m => Name -> m Type
conT ''Universe TypeQ -> TypeQ -> TypeQ
forall (m :: * -> *). Quote m => m Type -> m Type -> m Type
`appT` Name -> TypeQ
forall (m :: * -> *). Quote m => Name -> m Type
varT Name
n) [Name]
varNames
        let typ :: TypeQ
typ     = (TypeQ -> Name -> TypeQ) -> TypeQ -> [Name] -> TypeQ
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\TypeQ
c Name
n -> TypeQ
c TypeQ -> TypeQ -> TypeQ
forall (m :: * -> *). Quote m => m Type -> m Type -> m Type
`appT` Name -> TypeQ
forall (m :: * -> *). Quote m => Name -> m Type
varT Name
n) (Name -> TypeQ
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
parentName) [Name]
varNames

        Dec
i <- Q Cxt -> TypeQ -> [Q Dec] -> Q Dec
forall (m :: * -> *).
Quote m =>
m Cxt -> m Type -> [m Dec] -> m Dec
instanceD ([TypeQ] -> Q Cxt
forall (m :: * -> *). Quote m => [m Type] -> m Cxt
cxt [TypeQ]
constrs) (Name -> TypeQ
forall (m :: * -> *). Quote m => Name -> m Type
conT ''UniverseSome TypeQ -> TypeQ -> TypeQ
forall (m :: * -> *). Quote m => m Type -> m Type -> m Type
`appT` TypeQ
typ)
            [ DatatypeInfo -> Q Dec
instanceDecFor DatatypeInfo
di
            ]

        [Dec] -> DecsQ
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return [Dec
i]

instanceDecFor :: DatatypeInfo -> Q Dec
instanceDecFor :: DatatypeInfo -> Q Dec
instanceDecFor DatatypeInfo
di = Q Pat -> Q Body -> [Q Dec] -> Q Dec
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP 'universeSome) (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (Q Exp -> Q Body) -> Q Exp -> Q Body
forall a b. (a -> b) -> a -> b
$ DatatypeInfo -> Q Exp
universeSomeQ' DatatypeInfo
di) []

instance DeriveUniverseSome Dec where
  deriveUniverseSome :: Dec -> DecsQ
deriveUniverseSome (InstanceD Maybe Overlap
overlaps Cxt
c Type
classHead []) = do
    let instanceFor :: [Dec] -> Dec
instanceFor = Maybe Overlap -> Cxt -> Type -> [Dec] -> Dec
InstanceD Maybe Overlap
overlaps Cxt
c Type
classHead
    case Type
classHead of
      ConT Name
u `AppT` Type
t | Name
u Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== ''UniverseSome -> do
        Name
name <- Type -> Q Name
headOfType Type
t
        DatatypeInfo
di <- Name -> Q DatatypeInfo
reifyDatatype Name
name
        Dec
i <- ([Dec] -> Dec) -> DecsQ -> Q Dec
forall a b. (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Dec] -> Dec
instanceFor (DecsQ -> Q Dec) -> DecsQ -> Q Dec
forall a b. (a -> b) -> a -> b
$ (Q Dec -> Q Dec) -> [Q Dec] -> DecsQ
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 Q Dec -> Q Dec
forall a. a -> a
id
            [ DatatypeInfo -> Q Dec
instanceDecFor DatatypeInfo
di
            ]
        [Dec] -> DecsQ
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return [Dec
i]
      Type
_ -> String -> DecsQ
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> DecsQ) -> String -> DecsQ
forall a b. (a -> b) -> a -> b
$ String
"deriveUniverseSome: expected an instance head like `UniverseSome (C a b ...)`, got " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
classHead
  deriveUniverseSome Dec
_ = String -> DecsQ
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"deriveUniverseSome: expected an empty instance declaration"

-- | Derive the method for @:: ['Some' tag]@
--
-- >>> :set -XGADTs -XTemplateHaskell -XStandaloneDeriving
-- >>> import Data.GADT.Show
--
-- >>> data Tag b a where IntTag :: Tag b Int; BoolTag :: b -> Tag b Bool
-- >>> deriving instance Show b => Show (Tag b a)
-- >>> instance Show b => GShow (Tag b) where gshowsPrec = showsPrec
--
-- >>> $(universeSomeQ ''Tag) :: [Some (Tag Bool)]
-- [Some IntTag,Some (BoolTag False),Some (BoolTag True)]
--
universeSomeQ :: Name -> ExpQ
universeSomeQ :: Name -> Q Exp
universeSomeQ Name
name = Name -> Q DatatypeInfo
reifyDatatype Name
name Q DatatypeInfo -> (DatatypeInfo -> Q Exp) -> Q Exp
forall a b. Q a -> (a -> Q b) -> Q b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= DatatypeInfo -> Q Exp
universeSomeQ'

universeSomeQ' :: DatatypeInfo -> Q Exp
universeSomeQ' :: DatatypeInfo -> Q Exp
universeSomeQ' DatatypeInfo
di = do
  let DatatypeInfo { datatypeContext :: DatatypeInfo -> Cxt
datatypeContext = Cxt
ctxt
                   , datatypeName :: DatatypeInfo -> Name
datatypeName    = Name
parentName
                   , datatypeInstTypes :: DatatypeInfo -> Cxt
datatypeInstTypes = Cxt
vars0
                   , datatypeCons :: DatatypeInfo -> [ConstructorInfo]
datatypeCons    = [ConstructorInfo]
cons
                   } = DatatypeInfo
di

  -- check
  Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Cxt -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null Cxt
ctxt) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ String -> Q ()
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Datatype context is not empty"

  case Cxt -> Maybe (Cxt, Type)
forall a. [a] -> Maybe ([a], a)
safeUnsnoc Cxt
vars0 of
    Maybe (Cxt, Type)
Nothing -> String -> Q Exp
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Datatype should have at least one type variable"
    Just (Cxt
vars, Type
var) -> do
      let universe' :: Q Exp
universe'   = [| universe |]
      let uap :: Q Exp
uap         = [| (<+*+>) |]
      let interleave' :: Q Exp
interleave' = [| interleave |]
      let mapSome' :: Q Exp
mapSome'    = [| map mkSome |]

      let sums :: [Q Exp]
sums = (ConstructorInfo -> Q Exp) -> [ConstructorInfo] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Q Exp -> Q Exp -> Q Exp -> ConstructorInfo -> Q Exp
forall {m :: * -> *}.
Quote m =>
m Exp -> m Exp -> m Exp -> ConstructorInfo -> m Exp
universeForCon Q Exp
mapSome' Q Exp
universe' Q Exp
uap) [ConstructorInfo]
cons
      Q Exp
interleave' Q Exp -> Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` [Q Exp] -> Q Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
listE [Q Exp]
sums
  where
    universeForCon :: m Exp -> m Exp -> m Exp -> ConstructorInfo -> m Exp
universeForCon m Exp
mapSome' m Exp
universe' m Exp
uap ConstructorInfo
ci =
      let con :: m Exp
con     = [m Exp] -> m Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
listE [ Name -> m Exp
forall (m :: * -> *). Quote m => Name -> m Exp
conE (ConstructorInfo -> Name
constructorName ConstructorInfo
ci) ]
          nargs :: Int
nargs   = Cxt -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (ConstructorInfo -> Cxt
constructorFields ConstructorInfo
ci)
          conArgs :: m Exp
conArgs = (m Exp -> m Exp -> m Exp) -> m Exp -> [m Exp] -> m Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\m Exp
f m Exp
x -> Maybe (m Exp) -> m Exp -> Maybe (m Exp) -> m Exp
forall (m :: * -> *).
Quote m =>
Maybe (m Exp) -> m Exp -> Maybe (m Exp) -> m Exp
infixE (m Exp -> Maybe (m Exp)
forall a. a -> Maybe a
Just m Exp
f) m Exp
uap (m Exp -> Maybe (m Exp)
forall a. a -> Maybe a
Just m Exp
universe')) m Exp
con (Int -> m Exp -> [m Exp]
forall a. Int -> a -> [a]
replicate Int
nargs m Exp
universe')

      in m Exp
mapSome' m Exp -> m Exp -> m Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` m Exp
conArgs

-------------------------------------------------------------------------------
-- helpers
-------------------------------------------------------------------------------

headOfType :: Type -> Q Name
headOfType :: Type -> Q Name
headOfType (AppT Type
t Type
_) = Type -> Q Name
headOfType Type
t
headOfType (VarT Name
n)   = Name -> Q Name
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Name
n
headOfType (ConT Name
n)   = Name -> Q Name
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Name
n
headOfType Type
t          = String -> Q Name
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Name) -> String -> Q Name
forall a b. (a -> b) -> a -> b
$ String
"headOfType: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
t

safeUnsnoc :: [a] -> Maybe ([a], a)
safeUnsnoc :: forall a. [a] -> Maybe ([a], a)
safeUnsnoc [a]
xs = case [a] -> [a]
forall a. [a] -> [a]
reverse [a]
xs of
  []     -> Maybe ([a], a)
forall a. Maybe a
Nothing
  (a
y:[a]
ys) -> ([a], a) -> Maybe ([a], a)
forall a. a -> Maybe a
Just ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
ys, a
y)