{-# LANGUAGE Safe #-}

{-# OPTIONS_HADDOCK hide #-}

{-# LANGUAGE CPP #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeFamilies #-}

{- |
Module      : Text.Gigaparsec.Internal.TH.TypeUtils
Description : Common Template Haskell Utils
License     : BSD-3-Clause
Maintainer  : Jamie Willis, Gigaparsec Maintainers
Stability   : experimental

Common utils for manipulating Template Haskell Type AST.

@since 0.2.2.0
-}
module Text.Gigaparsec.Internal.TH.TypeUtils (module Text.Gigaparsec.Internal.TH.TypeUtils) where

import Text.Gigaparsec.Internal.TH.VersionAgnostic

import Data.Bifunctor (Bifunctor (bimap))
import Data.Set (Set)
import Data.Set qualified as Set
import Language.Haskell.TH (pprint)

-- When KindSignatures is off, the default (a :: *) that TH generates is broken!
sanitiseStarT :: TyVarBndr flag -> TyVarBndr flag
sanitiseStarT :: forall flag. TyVarBndr flag -> TyVarBndr flag
sanitiseStarT = (Name -> flag -> TyVarBndr flag)
-> (Name -> flag -> Kind -> TyVarBndr flag)
-> TyVarBndr flag
-> TyVarBndr flag
forall flag a.
(Name -> flag -> a)
-> (Name -> flag -> Kind -> a) -> TyVarBndr flag -> a
recTyVarBndr Name -> flag -> TyVarBndr flag
forall flag. Name -> flag -> TyVarBndr flag
mkPlainTV (\Name
x flag
y -> TyVarBndr flag -> Kind -> TyVarBndr flag
forall a b. a -> b -> a
const (Name -> flag -> TyVarBndr flag
forall flag. Name -> flag -> TyVarBndr flag
mkPlainTV Name
x flag
y))

-- | Remove stars from binder annotations when we don't have `KindSignatures` enabled.
sanitiseBndrStars :: [TyVarBndr flag] -> Q [TyVarBndr flag]
sanitiseBndrStars :: forall flag. [TyVarBndr flag] -> Q [TyVarBndr flag]
sanitiseBndrStars [TyVarBndr flag]
bndrs = do
  kindSigs <- Extension -> Q Bool
isExtEnabled Extension
KindSignatures
  return (if kindSigs then bndrs else map sanitiseStarT bndrs)

sanitiseTypeStars :: Type -> Q Type
sanitiseTypeStars :: Kind -> Q Kind
sanitiseTypeStars = (TypeF (Q Kind) -> Q Kind) -> Kind -> Q Kind
forall a. (TypeF a -> a) -> Kind -> a
cataType TypeF (Q Kind) -> Q Kind
go
 where
  go :: TypeF (Q Type) -> Q Type
  go :: TypeF (Q Kind) -> Q Kind
go (ForallTF [TyVarBndrF Specificity (Q Kind)]
bnds [Q Kind]
ctx Q Kind
tp) =
    [TyVarBndr Specificity] -> Cxt -> Kind -> Kind
ForallT ([TyVarBndr Specificity] -> Cxt -> Kind -> Kind)
-> Q [TyVarBndr Specificity] -> Q (Cxt -> Kind -> Kind)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TyVarBndrF Specificity (Q Kind) -> Q (TyVarBndr Specificity))
-> [TyVarBndrF Specificity (Q Kind)] -> Q [TyVarBndr Specificity]
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 TyVarBndrF Specificity (Q Kind) -> Q (TyVarBndr Specificity)
forall flag. TyVarBndrF flag (Q Kind) -> Q (TyVarBndr flag)
helpBnd [TyVarBndrF Specificity (Q Kind)]
bnds Q (Cxt -> Kind -> Kind) -> Q Cxt -> Q (Kind -> Kind)
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Q Kind] -> Q Cxt
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => [m a] -> m [a]
sequence [Q Kind]
ctx Q (Kind -> Kind) -> Q Kind -> Q Kind
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Q Kind
tp
  go (ForallVisTF [TyVarBndrF () (Q Kind)]
bnds Q Kind
tp) = [TyVarBndr ()] -> Kind -> Kind
mkForallVisT ([TyVarBndr ()] -> Kind -> Kind)
-> Q [TyVarBndr ()] -> Q (Kind -> Kind)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TyVarBndrF () (Q Kind) -> Q (TyVarBndr ()))
-> [TyVarBndrF () (Q Kind)] -> Q [TyVarBndr ()]
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 TyVarBndrF () (Q Kind) -> Q (TyVarBndr ())
forall flag. TyVarBndrF flag (Q Kind) -> Q (TyVarBndr flag)
helpBnd [TyVarBndrF () (Q Kind)]
bnds Q (Kind -> Kind) -> Q Kind -> Q Kind
forall a b. Q (a -> b) -> Q a -> Q b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Q Kind
tp
  go TypeF (Q Kind)
e = TypeF Kind -> Kind
embedType (TypeF Kind -> Kind) -> Q (TypeF Kind) -> Q Kind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeF (Q Kind) -> Q (TypeF Kind)
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => TypeF (m a) -> m (TypeF a)
sequence TypeF (Q Kind)
e

  helpBnd :: TyVarBndrF flag (Q Type) -> Q (TyVarBndr flag)
  helpBnd :: forall flag. TyVarBndrF flag (Q Kind) -> Q (TyVarBndr flag)
helpBnd (PlainTVF Name
n flag
f) = TyVarBndr flag -> Q (TyVarBndr flag)
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (TyVarBndr flag -> Q (TyVarBndr flag))
-> TyVarBndr flag -> Q (TyVarBndr flag)
forall a b. (a -> b) -> a -> b
$ Name -> flag -> TyVarBndr flag
forall flag. Name -> flag -> TyVarBndr flag
mkPlainTV Name
n flag
f
  helpBnd (KindedTVF Name
n flag
f ~Q Kind
_) = TyVarBndr flag -> Q (TyVarBndr flag)
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (TyVarBndr flag -> Q (TyVarBndr flag))
-> TyVarBndr flag -> Q (TyVarBndr flag)
forall a b. (a -> b) -> a -> b
$ Name -> flag -> TyVarBndr flag
forall flag. Name -> flag -> TyVarBndr flag
mkPlainTV Name
n flag
f

unTyVarBndrF :: TyVarBndrF flag k -> (Name, flag, Maybe k)
unTyVarBndrF :: forall flag k. TyVarBndrF flag k -> (Name, flag, Maybe k)
unTyVarBndrF = (Name -> flag -> (Name, flag, Maybe k))
-> (Name -> flag -> k -> (Name, flag, Maybe k))
-> TyVarBndrF flag k
-> (Name, flag, Maybe k)
forall flag a k.
(Name -> flag -> a)
-> (Name -> flag -> k -> a) -> TyVarBndrF flag k -> a
recTyVarBndrF (,,Maybe k
forall a. Maybe a
Nothing) (\Name
x flag
y k
z -> (Name
x, flag
y, k -> Maybe k
forall a. a -> Maybe a
Just k
z))

reTyVarBndr :: Name -> flag -> Maybe Type -> TyVarBndr flag
reTyVarBndr :: forall flag. Name -> flag -> Maybe Kind -> TyVarBndr flag
reTyVarBndr Name
n flag
f Maybe Kind
mt = case Maybe Kind
mt of
  Maybe Kind
Nothing -> Name -> flag -> TyVarBndr flag
forall flag. Name -> flag -> TyVarBndr flag
mkPlainTV Name
n flag
f
  Just Kind
t -> Name -> flag -> Kind -> TyVarBndr flag
forall flag. Name -> flag -> Kind -> TyVarBndr flag
mkKindedTV Name
n flag
f Kind
t

getBndrFName :: TyVarBndrF flag k -> Name
getBndrFName :: forall flag k. TyVarBndrF flag k -> Name
getBndrFName = (\(Name
a, flag
_, Maybe k
_) -> Name
a) ((Name, flag, Maybe k) -> Name)
-> (TyVarBndrF flag k -> (Name, flag, Maybe k))
-> TyVarBndrF flag k
-> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyVarBndrF flag k -> (Name, flag, Maybe k)
forall flag k. TyVarBndrF flag k -> (Name, flag, Maybe k)
unTyVarBndrF

removeUnusedTVars :: Type -> Type
removeUnusedTVars :: Kind -> Kind
removeUnusedTVars = (TypeF (Set Name) -> Set Name)
-> (TypeF (Set Name, Kind) -> Kind) -> Kind -> Kind
forall b a. (TypeF b -> b) -> (TypeF (b, a) -> a) -> Kind -> a
zygoType TypeF (Set Name) -> Set Name
typeFreeVarsAlg TypeF (Set Name, Kind) -> Kind
go
 where
  go :: TypeF (Set Name, Type) -> Type
  go :: TypeF (Set Name, Kind) -> Kind
go (ForallTF [TyVarBndrF Specificity (Set Name, Kind)]
bnds [(Set Name, Kind)]
ctx (Set Name
tpNames, Kind
tp)) =
    let ([Set Name]
ctxNames, Cxt
ctx') = [(Set Name, Kind)] -> ([Set Name], Cxt)
forall a b. [(a, b)] -> ([a], [b])
unzip [(Set Name, Kind)]
ctx
        -- All names that *do* occur in the rest of the type/constraints
        allFreeNames :: Set Name
allFreeNames = [Set Name] -> Set Name
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions (Set Name
tpNames Set Name -> [Set Name] -> [Set Name]
forall a. a -> [a] -> [a]
: [Set Name]
ctxNames)
        ([TyVarBndr Specificity]
bnds', Set Name
_) = (TyVarBndrF Specificity (Set Name, Kind)
 -> ([TyVarBndr Specificity], Set Name)
 -> ([TyVarBndr Specificity], Set Name))
-> ([TyVarBndr Specificity], Set Name)
-> [TyVarBndrF Specificity (Set Name, Kind)]
-> ([TyVarBndr Specificity], Set Name)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TyVarBndrF Specificity (Set Name, Kind)
-> ([TyVarBndr Specificity], Set Name)
-> ([TyVarBndr Specificity], Set Name)
forall s.
TyVarBndrF s (Set Name, Kind)
-> ([TyVarBndr s], Set Name) -> ([TyVarBndr s], Set Name)
discardUnusedTVars ([], Set Name
allFreeNames) [TyVarBndrF Specificity (Set Name, Kind)]
bnds
     in [TyVarBndr Specificity] -> Cxt -> Kind -> Kind
ForallT [TyVarBndr Specificity]
bnds' Cxt
ctx' Kind
tp
  go (ForallVisTF [TyVarBndrF () (Set Name, Kind)]
bnds (Set Name
tpNames, Kind
tp)) =
    let ([TyVarBndr ()]
bnds', Set Name
_) = (TyVarBndrF () (Set Name, Kind)
 -> ([TyVarBndr ()], Set Name) -> ([TyVarBndr ()], Set Name))
-> ([TyVarBndr ()], Set Name)
-> [TyVarBndrF () (Set Name, Kind)]
-> ([TyVarBndr ()], Set Name)
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TyVarBndrF () (Set Name, Kind)
-> ([TyVarBndr ()], Set Name) -> ([TyVarBndr ()], Set Name)
forall s.
TyVarBndrF s (Set Name, Kind)
-> ([TyVarBndr s], Set Name) -> ([TyVarBndr s], Set Name)
discardUnusedTVars ([], Set Name
tpNames) [TyVarBndrF () (Set Name, Kind)]
bnds
     in [TyVarBndr ()] -> Kind -> Kind
ForallVisT [TyVarBndr ()]
bnds' Kind
tp
  go TypeF (Set Name, Kind)
e = TypeF Kind -> Kind
embedType ((Set Name, Kind) -> Kind
forall a b. (a, b) -> b
snd ((Set Name, Kind) -> Kind) -> TypeF (Set Name, Kind) -> TypeF Kind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TypeF (Set Name, Kind)
e)

  discardUnusedTVars ::
    TyVarBndrF s (Set Name, Type) ->
    ([TyVarBndr s], Set Name) ->
    ([TyVarBndr s], Set Name)
  discardUnusedTVars :: forall s.
TyVarBndrF s (Set Name, Kind)
-> ([TyVarBndr s], Set Name) -> ([TyVarBndr s], Set Name)
discardUnusedTVars TyVarBndrF s (Set Name, Kind)
bnd ([TyVarBndr s]
bnds, Set Name
names) =
    let (Name
n, s
f, Maybe (Set Name, Kind)
mk) = TyVarBndrF s (Set Name, Kind) -> (Name, s, Maybe (Set Name, Kind))
forall flag k. TyVarBndrF flag k -> (Name, flag, Maybe k)
unTyVarBndrF TyVarBndrF s (Set Name, Kind)
bnd
     in if Name
n Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set Name
names
          -- We keep n, and add any free vars in its kind (if it has one)
          then
            ( Name -> s -> Maybe Kind -> TyVarBndr s
forall flag. Name -> flag -> Maybe Kind -> TyVarBndr flag
reTyVarBndr Name
n s
f ((Set Name, Kind) -> Kind
forall a b. (a, b) -> b
snd ((Set Name, Kind) -> Kind) -> Maybe (Set Name, Kind) -> Maybe Kind
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Set Name, Kind)
mk) TyVarBndr s -> [TyVarBndr s] -> [TyVarBndr s]
forall a. a -> [a] -> [a]
: [TyVarBndr s]
bnds
            , Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.union (Set Name
-> ((Set Name, Kind) -> Set Name)
-> Maybe (Set Name, Kind)
-> Set Name
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Set Name
forall a. Set a
Set.empty (Set Name, Kind) -> Set Name
forall a b. (a, b) -> a
fst Maybe (Set Name, Kind)
mk) Set Name
names
            )
          -- We discard n as it does not appear in subterms
          else ([TyVarBndr s]
bnds, Set Name
names)

typeFreeVarsAlg :: TypeF (Set Name) -> Set Name
typeFreeVarsAlg :: TypeF (Set Name) -> Set Name
typeFreeVarsAlg = TypeF (Set Name) -> Set Name
go
 where
  go :: TypeF (Set Name) -> Set Name
  go :: TypeF (Set Name) -> Set Name
go (VarTF Name
x) = Name -> Set Name
forall a. a -> Set a
Set.singleton Name
x
  go (ForallTF [TyVarBndrF Specificity (Set Name)]
bnds [Set Name]
ctx Set Name
tp) = [TyVarBndrF Specificity (Set Name)] -> Set Name -> Set Name
forall flag. [TyVarBndrF flag (Set Name)] -> Set Name -> Set Name
handleBnds [TyVarBndrF Specificity (Set Name)]
bnds ([Set Name] -> Set Name
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions ([Set Name] -> Set Name) -> [Set Name] -> Set Name
forall a b. (a -> b) -> a -> b
$ (Set Name
tp Set Name -> [Set Name] -> [Set Name]
forall a. a -> [a] -> [a]
: [Set Name]
ctx))
  go (ForallVisTF [TyVarBndrF () (Set Name)]
bnds Set Name
tp) = [TyVarBndrF () (Set Name)] -> Set Name -> Set Name
forall flag. [TyVarBndrF flag (Set Name)] -> Set Name -> Set Name
handleBnds [TyVarBndrF () (Set Name)]
bnds Set Name
tp
  go TypeF (Set Name)
e = (Set Name -> Set Name -> Set Name)
-> Set Name -> TypeF (Set Name) -> Set Name
forall a b. (a -> b -> b) -> b -> TypeF a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set Name
forall a. Set a
Set.empty TypeF (Set Name)
e

  handleBnds :: [TyVarBndrF flag (Set Name)] -> Set Name -> Set Name
  handleBnds :: forall flag. [TyVarBndrF flag (Set Name)] -> Set Name -> Set Name
handleBnds [TyVarBndrF flag (Set Name)]
bnds Set Name
ns =
    let (Set Name
as, Set Name
ks) =
          ([Name] -> Set Name)
-> ([Set Name] -> Set Name)
-> ([Name], [Set Name])
-> (Set Name, Set Name)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap [Name] -> Set Name
forall a. Ord a => [a] -> Set a
Set.fromList [Set Name] -> Set Name
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions (([Name], [Set Name]) -> (Set Name, Set Name))
-> ([Name], [Set Name]) -> (Set Name, Set Name)
forall a b. (a -> b) -> a -> b
$ [(Name, Set Name)] -> ([Name], [Set Name])
forall a b. [(a, b)] -> ([a], [b])
unzip ((TyVarBndrF flag (Set Name) -> (Name, Set Name))
-> [TyVarBndrF flag (Set Name)] -> [(Name, Set Name)]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndrF flag (Set Name) -> (Name, Set Name)
forall flag. TyVarBndrF flag (Set Name) -> (Name, Set Name)
bndrFreeAndBoundNames [TyVarBndrF flag (Set Name)]
bnds)
     in Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.difference (Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set Name
ks Set Name
ns) Set Name
as

  bndrFreeAndBoundNames :: TyVarBndrF flag (Set Name) -> (Name, Set Name)
  bndrFreeAndBoundNames :: forall flag. TyVarBndrF flag (Set Name) -> (Name, Set Name)
bndrFreeAndBoundNames (PlainTVF Name
x flag
_) = (Name
x, Set Name
forall a. Set a
Set.empty)
  bndrFreeAndBoundNames (KindedTVF Name
x flag
_ Set Name
k) = (Name
x, Set Name
k)

typeFreeVars :: Type -> Set Name
typeFreeVars :: Kind -> Set Name
typeFreeVars = (TypeF (Set Name) -> Set Name) -> Kind -> Set Name
forall a. (TypeF a -> a) -> Kind -> a
cataType TypeF (Set Name) -> Set Name
typeFreeVarsAlg

getRecordFields :: Info -> Q [(Name, Type)]
getRecordFields :: Info -> Q [(Name, Kind)]
getRecordFields Info
i = case Info
i of
  TyConI (DataD Cxt
_ Name
_ [TyVarBndr BndrVis]
_ Maybe Kind
_ [Con]
cstrs [DerivClause]
_) -> [[(Name, Kind)]] -> [(Name, Kind)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Name, Kind)]] -> [(Name, Kind)])
-> Q [[(Name, Kind)]] -> Q [(Name, Kind)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Con -> Q [(Name, Kind)]) -> [Con] -> Q [[(Name, Kind)]]
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 -> Q [(Name, Kind)]
getFieldNames [Con]
cstrs
  TyConI (NewtypeD Cxt
_ Name
_ [TyVarBndr BndrVis]
_ Maybe Kind
_ Con
cstr [DerivClause]
_) -> Con -> Q [(Name, Kind)]
getFieldNames Con
cstr
  DataConI Name
_ Kind
_ Name
tname -> Info -> Q [(Name, Kind)]
getRecordFields (Info -> Q [(Name, Kind)]) -> Q Info -> Q [(Name, Kind)]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Name -> Q Info
reify Name
tname
  Info
info -> String -> Q [(Name, Kind)]
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q [(Name, Kind)]) -> String -> Q [(Name, Kind)]
forall a b. (a -> b) -> a -> b
$ [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [String
"getRecordFields: given info is not for a record: `", Info -> String
forall a. Ppr a => a -> String
pprint Info
info, String
"`"]
  where
    getFieldNames :: Con -> Q [(Name, Type)]
    getFieldNames :: Con -> Q [(Name, Kind)]
getFieldNames (RecC Name
_ [VarBangType]
tps) = [(Name, Kind)] -> Q [(Name, Kind)]
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return ([(Name, Kind)] -> Q [(Name, Kind)])
-> [(Name, Kind)] -> Q [(Name, Kind)]
forall a b. (a -> b) -> a -> b
$ (VarBangType -> (Name, Kind)) -> [VarBangType] -> [(Name, Kind)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Name
nm, Bang
_, Kind
tp) -> (Name
nm, Kind
tp)) [VarBangType]
tps
    getFieldNames Con
c = String -> Q [(Name, Kind)]
forall a. String -> Q a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"getRecordFields: Constructor is not a record: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Con -> String
forall a. Ppr a => a -> String
pprint Con
c)