{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

-- | Shared functions for dependent-sum-template
module Data.GADT.TH.Internal where

import Control.Monad
import Control.Monad.Writer
import Data.List (foldl', drop)
import Data.Maybe
import Data.Map (Map)
import qualified Data.Map as Map
import qualified Data.Map.Merge.Lazy as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Language.Haskell.TH
import Language.Haskell.TH.Datatype
import Language.Haskell.TH.Datatype.TyVarBndr

classHeadToParams :: Type -> (Name, [Type])
classHeadToParams :: Type -> (Name, Cxt)
classHeadToParams Type
t = (Name
h, Cxt -> Cxt
forall a. [a] -> [a]
reverse Cxt
reversedParams)
  where
    (Name
h, Cxt
reversedParams) = Type -> (Name, Cxt)
go Type
t
    go :: Type -> (Name, [Type])
    go :: Type -> (Name, Cxt)
go Type
t = case Type
t of
      AppT Type
f Type
x ->
        let (Name
h, Cxt
reversedParams) = Type -> (Name, Cxt)
classHeadToParams Type
f
        in (Name
h, Type
x Type -> Cxt -> Cxt
forall a. a -> [a] -> [a]
: Cxt
reversedParams)
      Type
_ -> (Type -> Name
headOfType Type
t, [])

-- Do not export this data family, it must remain empty. It's used as a way to trick GHC into not unifying certain type variables.
data family Skolem :: k -> k

skolemize :: Set Name -> Type -> Type
skolemize :: Set Name -> Type -> Type
skolemize Set Name
rigids Type
t = case Type
t of
  ForallT [TyVarBndr Specificity]
bndrs Cxt
cxt Type
t' -> [TyVarBndr Specificity] -> Cxt -> Type -> Type
ForallT [TyVarBndr Specificity]
bndrs Cxt
cxt (Set Name -> Type -> Type
skolemize (Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.difference Set Name
rigids ([Name] -> Set Name
forall a. Ord a => [a] -> Set a
Set.fromList ((TyVarBndr Specificity -> Name)
-> [TyVarBndr Specificity] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr Specificity -> Name
forall flag. TyVarBndr_ flag -> Name
tvName [TyVarBndr Specificity]
bndrs))) Type
t')
  AppT Type
t1 Type
t2 -> Type -> Type -> Type
AppT (Set Name -> Type -> Type
skolemize Set Name
rigids Type
t1) (Set Name -> Type -> Type
skolemize Set Name
rigids Type
t2)
  SigT Type
t Type
k -> Type -> Type -> Type
SigT (Set Name -> Type -> Type
skolemize Set Name
rigids Type
t) Type
k
  VarT Name
v -> if Name -> Set Name -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Name
v Set Name
rigids
    then Type -> Type -> Type
AppT (Name -> Type
ConT ''Skolem) (Name -> Type
VarT Name
v)
    else Type
t
  InfixT Type
t1 Name
n Type
t2 -> Type -> Name -> Type -> Type
InfixT (Set Name -> Type -> Type
skolemize Set Name
rigids Type
t1) Name
n (Set Name -> Type -> Type
skolemize Set Name
rigids Type
t2)
  UInfixT Type
t1 Name
n Type
t2 -> Type -> Name -> Type -> Type
UInfixT (Set Name -> Type -> Type
skolemize Set Name
rigids Type
t1) Name
n (Set Name -> Type -> Type
skolemize Set Name
rigids Type
t2)
  ParensT Type
t -> Type -> Type
ParensT (Set Name -> Type -> Type
skolemize Set Name
rigids Type
t)
  Type
_ -> Type
t

reifyInstancesWithRigids :: Set Name -> Name -> [Type] -> Q [InstanceDec]
reifyInstancesWithRigids :: Set Name -> Name -> Cxt -> Q [InstanceDec]
reifyInstancesWithRigids Set Name
rigids Name
cls Cxt
tys = Name -> Cxt -> Q [InstanceDec]
reifyInstances Name
cls ((Type -> Type) -> Cxt -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (Set Name -> Type -> Type
skolemize Set Name
rigids) Cxt
tys)

-- | Determine the type variables which occur freely in a type.
freeTypeVariables :: Type -> Set Name
freeTypeVariables :: Type -> Set Name
freeTypeVariables Type
t = case Type
t of
  ForallT [TyVarBndr Specificity]
bndrs Cxt
_ Type
t' -> Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.difference (Type -> Set Name
freeTypeVariables Type
t') ([Name] -> Set Name
forall a. Ord a => [a] -> Set a
Set.fromList ((TyVarBndr Specificity -> Name)
-> [TyVarBndr Specificity] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr Specificity -> Name
forall flag. TyVarBndr_ flag -> Name
tvName [TyVarBndr Specificity]
bndrs))
  AppT Type
t1 Type
t2 -> Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.union (Type -> Set Name
freeTypeVariables Type
t1) (Type -> Set Name
freeTypeVariables Type
t2)
  SigT Type
t Type
_ -> Type -> Set Name
freeTypeVariables Type
t
  VarT Name
n -> Name -> Set Name
forall a. a -> Set a
Set.singleton Name
n
  Type
_ -> Set Name
forall a. Set a
Set.empty

subst :: Map Name Type -> Type -> Type
subst :: Map Name Type -> Type -> Type
subst Map Name Type
s = Type -> Type
f
  where
    f :: Type -> Type
f = \case
      ForallT [TyVarBndr Specificity]
bndrs Cxt
cxt Type
t ->
        let s' :: Map Name Type
s' = Map Name Type -> Map Name () -> Map Name Type
forall k a b. Ord k => Map k a -> Map k b -> Map k a
Map.difference Map Name Type
s ([(Name, ())] -> Map Name ()
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Name
k,()) | Name
k <- (TyVarBndr Specificity -> Name)
-> [TyVarBndr Specificity] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map TyVarBndr Specificity -> Name
forall flag. TyVarBndr_ flag -> Name
tvName [TyVarBndr Specificity]
bndrs])
        in [TyVarBndr Specificity] -> Cxt -> Type -> Type
ForallT [TyVarBndr Specificity]
bndrs Cxt
cxt (Map Name Type -> Type -> Type
subst Map Name Type
s' Type
t)
      AppT Type
t Type
t' -> Type -> Type -> Type
AppT (Type -> Type
f Type
t) (Type -> Type
f Type
t')
      SigT Type
t Type
k -> Type -> Type -> Type
SigT (Type -> Type
f Type
t) Type
k
      VarT Name
n -> case Name -> Map Name Type -> Maybe Type
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Name
n Map Name Type
s of
        Just Type
t -> Type
t
        Maybe Type
Nothing -> Name -> Type
VarT Name
n
      InfixT Type
t Name
x Type
t' -> Type -> Name -> Type -> Type
InfixT (Type -> Type
f Type
t) Name
x (Type -> Type
f Type
t')
      UInfixT Type
t Name
x Type
t' -> Type -> Name -> Type -> Type
UInfixT (Type -> Type
f Type
t) Name
x (Type -> Type
f Type
t')
      Type
x -> Type
x

-- Invoke the deriver for the given class instance.  We assume that the type
-- we're deriving for is always the first typeclass parameter, if there are
-- multiple.
deriveForDec
  :: Name
  -> (DatatypeInfo -> WriterT [Type] Q Dec)
  -> Dec
  -> Q [Dec]
deriveForDec :: Name
-> (DatatypeInfo -> WriterT Cxt Q InstanceDec)
-> InstanceDec
-> Q [InstanceDec]
deriveForDec Name
className DatatypeInfo -> WriterT Cxt Q InstanceDec
f (InstanceD Maybe Overlap
overlaps Cxt
cxt Type
instanceHead [InstanceDec]
decs) = do
  let (Name
givenClassName, Type
firstParam : Cxt
_) = Type -> (Name, Cxt)
classHeadToParams Type
instanceHead
  Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Name
givenClassName Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
/= Name
className) (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 -> Q ()) -> String -> Q ()
forall a b. (a -> b) -> a -> b
$ String
"while deriving " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show Name
className String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": wrong class name in prototype declaration: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show Name
givenClassName
  let dataTypeName :: Name
dataTypeName = Type -> Name
headOfType Type
firstParam
  dataTypeInfo <- Name -> Q DatatypeInfo
reifyDatatype Name
dataTypeName
  let instTypes = DatatypeInfo -> Cxt
datatypeInstTypes DatatypeInfo
dataTypeInfo
      paramVars = [Set Name] -> Set Name
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions [Type -> Set Name
freeTypeVariables Type
t | Type
t <- Cxt
instTypes]
      instTypes' = case Cxt -> Cxt
forall a. [a] -> [a]
reverse Cxt
instTypes of
        [] -> String -> Cxt
forall a. String -> [a]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"deriveGEq: Not enough type parameters"
        (Type
_:Cxt
xs) -> Cxt -> Cxt
forall a. [a] -> [a]
reverse Cxt
xs
      generatedInstanceHead = Type -> Type -> Type
AppT (Name -> Type
ConT Name
className) ((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 (Name -> Type
ConT (Name -> Type) -> Name -> Type
forall a b. (a -> b) -> a -> b
$ DatatypeInfo -> Name
datatypeName DatatypeInfo
dataTypeInfo) Cxt
instTypes')
  unifiedTypes <- unifyTypes [generatedInstanceHead, instanceHead]
  let
    newInstanceHead = Map Name Type -> Type -> Type
forall a. TypeSubstitution a => Map Name Type -> a -> a
applySubstitution Map Name Type
unifiedTypes Type
instanceHead
    newContext = Map Name Type -> Cxt -> Cxt
forall a. TypeSubstitution a => Map Name Type -> a -> a
applySubstitution Map Name Type
unifiedTypes Cxt
cxt
  -- We are not using the generated context that we collect from f, instead
  -- relying on a correct instance head from the user
  (dec, _) <- runWriterT $ f dataTypeInfo
  return [InstanceD overlaps newContext newInstanceHead [dec]]
deriveForDec Name
className DatatypeInfo -> WriterT Cxt Q InstanceDec
f InstanceDec
dataDec = do
  dataTypeInfo <- InstanceDec -> Q DatatypeInfo
normalizeDec InstanceDec
dataDec
  let instTypes = DatatypeInfo -> Cxt
datatypeInstTypes DatatypeInfo
dataTypeInfo
      paramVars = [Set Name] -> Set Name
forall (f :: * -> *) a. (Foldable f, Ord a) => f (Set a) -> Set a
Set.unions [Type -> Set Name
freeTypeVariables Type
t | Type
t <- Cxt
instTypes]
      instTypes' = case Cxt -> Cxt
forall a. [a] -> [a]
reverse Cxt
instTypes of
        [] -> String -> Cxt
forall a. String -> [a]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"deriveGEq: Not enough type parameters"
        (Type
_:Cxt
xs) -> Cxt -> Cxt
forall a. [a] -> [a]
reverse Cxt
xs
      instanceHead = Type -> Type -> Type
AppT (Name -> Type
ConT Name
className) ((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 (Name -> Type
ConT (Name -> Type) -> Name -> Type
forall a b. (a -> b) -> a -> b
$ DatatypeInfo -> Name
datatypeName DatatypeInfo
dataTypeInfo) Cxt
instTypes')
  (dec, cxt') <- runWriterT (f dataTypeInfo)
  return [InstanceD Nothing (datatypeContext dataTypeInfo ++ cxt') instanceHead [dec]]

headOfType :: Type -> Name
headOfType :: Type -> Name
headOfType = \case
  ForallT [TyVarBndr Specificity]
_ Cxt
_ Type
ty -> Type -> Name
headOfType Type
ty
  VarT Name
name -> Name
name
  ConT Name
name -> Name
name
  TupleT Int
n -> Int -> Name
tupleTypeName Int
n
  Type
ArrowT -> ''(->)
  Type
ListT -> ''[]
  AppT Type
t Type
_ -> Type -> Name
headOfType Type
t