{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE PolyKinds #-}


module Data.Aeson.GADT.TH
  ( deriveJSONGADT
  , deriveToJSONGADT
  , deriveFromJSONGADT

  , deriveJSONGADTWithOptions
  , deriveToJSONGADTWithOptions
  , deriveFromJSONGADTWithOptions

  , JSONGADTOptions(JSONGADTOptions, gadtConstructorModifier)
  , defaultJSONGADTOptions

  ) where

import Control.Monad (forM, replicateM)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Writer (WriterT, runWriterT, tell)
import Data.Aeson (FromJSON(..), ToJSON(..))
import Data.List (intercalate, partition, sort)
import qualified Data.List.NonEmpty as NonEmpty
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Some (Some(..))
import Language.Haskell.TH hiding (cxt)
import Language.Haskell.TH.Datatype (ConstructorInfo(..), applySubstitution, datatypeCons, reifyDatatype, unifyTypes)
import Language.Haskell.TH.Datatype.TyVarBndr (tvName)

#if !MIN_VERSION_template_haskell(2,21,0)
#if MIN_VERSION_th_abstraction(0,6,0)
import Language.Haskell.TH.Datatype.TyVarBndr (TyVarBndrVis)
#else
import Language.Haskell.TH.Datatype.TyVarBndr (TyVarBndr_)
type TyVarBndrVis = TyVarBndr_ ()
#endif
#endif

#if MIN_VERSION_dependent_sum(0,5,0)
#else
pattern Some :: tag a -> Some tag
pattern Some x = This x
#endif

-- Do not export this type 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 [Type]
cxt Type
t' -> [TyVarBndr Specificity] -> [Type] -> Type -> Type
ForallT [TyVarBndr Specificity]
bndrs [Type]
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
t1 Type
k -> Type -> Type -> Type
SigT (Set Name -> Type -> Type
skolemize Set Name
rigids Type
t1) 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
t1 -> Type -> Type
ParensT (Set Name -> Type -> Type
skolemize Set Name
rigids Type
t1)
  Type
_ -> Type
t

reifyInstancesWithRigids :: Set Name -> Name -> [Type] -> Q [InstanceDec]
reifyInstancesWithRigids :: Set Name -> Name -> [Type] -> Q [InstanceDec]
reifyInstancesWithRigids Set Name
rigids Name
cls [Type]
tys = Name -> [Type] -> Q [InstanceDec]
reifyInstances Name
cls ((Type -> Type) -> [Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map (Set Name -> Type -> Type
skolemize Set Name
rigids) [Type]
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 [Type]
_ 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
t1 Type
_ -> Type -> Set Name
freeTypeVariables Type
t1
  VarT Name
n -> Name -> Set Name
forall a. a -> Set a
Set.singleton Name
n
  Type
_ -> Set Name
forall a. Set a
Set.empty

newtype JSONGADTOptions = JSONGADTOptions
  { JSONGADTOptions -> String -> String
gadtConstructorModifier :: String -> String }

defaultJSONGADTOptions :: JSONGADTOptions
defaultJSONGADTOptions :: JSONGADTOptions
defaultJSONGADTOptions = JSONGADTOptions
  { gadtConstructorModifier :: String -> String
gadtConstructorModifier = String -> String
forall a. a -> a
id }

-- | Derive 'ToJSON' and 'FromJSON' instances for the named GADT
deriveJSONGADT :: Name -> DecsQ
deriveJSONGADT :: Name -> Q [InstanceDec]
deriveJSONGADT = JSONGADTOptions -> Name -> Q [InstanceDec]
deriveJSONGADTWithOptions JSONGADTOptions
defaultJSONGADTOptions

deriveJSONGADTWithOptions :: JSONGADTOptions -> Name -> DecsQ
deriveJSONGADTWithOptions :: JSONGADTOptions -> Name -> Q [InstanceDec]
deriveJSONGADTWithOptions JSONGADTOptions
opts Name
n = do
  tj <- JSONGADTOptions -> Name -> Q [InstanceDec]
deriveToJSONGADTWithOptions JSONGADTOptions
opts Name
n
  fj <- deriveFromJSONGADTWithOptions opts n
  return (tj ++ fj)

deriveToJSONGADT :: Name -> DecsQ
deriveToJSONGADT :: Name -> Q [InstanceDec]
deriveToJSONGADT = JSONGADTOptions -> Name -> Q [InstanceDec]
deriveToJSONGADTWithOptions JSONGADTOptions
defaultJSONGADTOptions

deriveToJSONGADTWithOptions :: JSONGADTOptions -> Name -> DecsQ
deriveToJSONGADTWithOptions :: JSONGADTOptions -> Name -> Q [InstanceDec]
deriveToJSONGADTWithOptions JSONGADTOptions
opts Name
n = do
  info <- Name -> Q DatatypeInfo
reifyDatatype Name
n
  let cons = DatatypeInfo -> [ConstructorInfo]
datatypeCons DatatypeInfo
info
  topVars <- makeTopVars n
  let n' = (Type -> Name -> Type) -> Type -> [Name] -> 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
c Name
v -> Type -> Type -> Type
AppT Type
c (Name -> Type
VarT Name
v)) (Name -> Type
ConT Name
n) [Name]
topVars
  (matches, constraints') <- runWriterT (mapM (fmap pure . conMatchesToJSON opts topVars) cons)
  let constraints = (NonEmpty Type -> Type) -> [NonEmpty Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map NonEmpty Type -> Type
forall a. NonEmpty a -> a
NonEmpty.head ([NonEmpty Type] -> [Type])
-> ([Type] -> [NonEmpty Type]) -> [Type] -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> [NonEmpty Type]
forall (f :: * -> *) a. (Foldable f, Eq a) => f a -> [NonEmpty a]
NonEmpty.group ([Type] -> [NonEmpty Type])
-> ([Type] -> [Type]) -> [Type] -> [NonEmpty Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> [Type]
forall a. Ord a => [a] -> [a]
sort ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ [Type]
constraints'
  impl <- funD 'toJSON
    [ clause [] (normalB $ lamCaseE matches) []
    ]
  return [ InstanceD Nothing constraints (AppT (ConT ''ToJSON) n') [impl] ]

makeTopVars :: Name -> Q [Name]
makeTopVars :: Name -> Q [Name]
makeTopVars Name
tyConName = do
  (tyVarBndrs, kArity) <- Name -> Q ([TyVarBndrVis], Int)
tyConArity' Name
tyConName
  extraVars <- replicateM kArity (newName "topvar")
  return (map tvName tyVarBndrs ++ extraVars)

deriveFromJSONGADT :: Name -> DecsQ
deriveFromJSONGADT :: Name -> Q [InstanceDec]
deriveFromJSONGADT = JSONGADTOptions -> Name -> Q [InstanceDec]
deriveFromJSONGADTWithOptions JSONGADTOptions
defaultJSONGADTOptions

deriveFromJSONGADTWithOptions :: JSONGADTOptions -> Name -> DecsQ
deriveFromJSONGADTWithOptions :: JSONGADTOptions -> Name -> Q [InstanceDec]
deriveFromJSONGADTWithOptions JSONGADTOptions
opts Name
n = do
  info <- Name -> Q DatatypeInfo
reifyDatatype Name
n

  let cons = DatatypeInfo -> [ConstructorInfo]
datatypeCons DatatypeInfo
info
      allConNames =
        String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$
          (ConstructorInfo -> String) -> [ConstructorInfo] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (JSONGADTOptions -> String -> String
gadtConstructorModifier JSONGADTOptions
opts (String -> String)
-> (ConstructorInfo -> String) -> ConstructorInfo -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase (Name -> String)
-> (ConstructorInfo -> Name) -> ConstructorInfo -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ConstructorInfo -> Name
constructorName) [ConstructorInfo]
cons
  wildName <- newName "s"
  let wild = Q Pat -> Q Body -> [Q InstanceDec] -> Q Match
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m InstanceDec] -> m Match
match (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
wildName) (Q Exp -> Q Body
forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [e|
        fail $
          "Expected tag to be one of [" ++ allConNames ++ "] but got: "
          ++ $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
wildName)
        |]) []
  topVars <- makeTopVars n
  let n' = (Type -> Name -> Type) -> Type -> [Name] -> 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
c Name
v -> Type -> Type -> Type
AppT Type
c (Name -> Type
VarT Name
v)) (Name -> Type
ConT Name
n) ([Name] -> Type) -> [Name] -> Type
forall a b. (a -> b) -> a -> b
$ [Name] -> [Name]
forall a. HasCallStack => [a] -> [a]
init [Name]
topVars
  (matches, constraints') <- runWriterT $ mapM (conMatchesParseJSON opts topVars [|_v'|]) cons
  let constraints = (NonEmpty Type -> Type) -> [NonEmpty Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map NonEmpty Type -> Type
forall a. NonEmpty a -> a
NonEmpty.head ([NonEmpty Type] -> [Type])
-> ([Type] -> [NonEmpty Type]) -> [Type] -> [Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> [NonEmpty Type]
forall (f :: * -> *) a. (Foldable f, Eq a) => f a -> [NonEmpty a]
NonEmpty.group ([Type] -> [NonEmpty Type])
-> ([Type] -> [Type]) -> [Type] -> [NonEmpty Type]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> [Type]
forall a. Ord a => [a] -> [a]
sort ([Type] -> [Type]) -> [Type] -> [Type]
forall a b. (a -> b) -> a -> b
$ [Type]
constraints'
  v <- newName "v"
  parser <- funD 'parseJSON
    [ clause [varP v] (normalB [e|
        do (tag', _v') <- parseJSON $(varE v)
           $(caseE [|tag' :: String|] $ map pure matches ++ [wild])
      |]) []
    ]
  return [ InstanceD Nothing constraints (AppT (ConT ''FromJSON) (AppT (ConT ''Some) n')) [parser] ]

splitTopVars :: [Name] -> (Set Name, Name)
splitTopVars :: [Name] -> (Set Name, Name)
splitTopVars [Name]
allTopVars = case [Name] -> [Name]
forall a. [a] -> [a]
reverse [Name]
allTopVars of
  (Name
x:[Name]
xs) -> ([Name] -> Set Name
forall a. Ord a => [a] -> Set a
Set.fromList [Name]
xs, Name
x)
  [Name]
_ -> String -> (Set Name, Name)
forall a. HasCallStack => String -> a
error String
"splitTopVars: Empty set of variables"

-- | Implementation of 'toJSON'
conMatchesToJSON :: JSONGADTOptions -> [Name] -> ConstructorInfo -> WriterT [Type] Q Match
conMatchesToJSON :: JSONGADTOptions
-> [Name] -> ConstructorInfo -> WriterT [Type] Q Match
conMatchesToJSON JSONGADTOptions
opts [Name]
allTopVars ConstructorInfo
c = do
  let (Set Name
topVars, Name
lastVar) = [Name] -> (Set Name, Name)
splitTopVars [Name]
allTopVars
      name :: Name
name = ConstructorInfo -> Name
constructorName ConstructorInfo
c
      base :: String
base = JSONGADTOptions -> String -> String
gadtConstructorModifier JSONGADTOptions
opts (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Name -> String
nameBase Name
name
      toJSONExp :: m Exp -> m Exp
toJSONExp m Exp
e = [| toJSON $(m Exp
e) |]
  vars <- Q [Name] -> WriterT [Type] Q [Name]
forall (m :: * -> *) a. Monad m => m a -> WriterT [Type] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q [Name] -> WriterT [Type] Q [Name])
-> ((Type -> Q Name) -> Q [Name])
-> (Type -> Q Name)
-> WriterT [Type] Q [Name]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Type] -> (Type -> Q Name) -> Q [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (ConstructorInfo -> [Type]
constructorFields ConstructorInfo
c) ((Type -> Q Name) -> WriterT [Type] Q [Name])
-> (Type -> Q Name) -> WriterT [Type] Q [Name]
forall a b. (a -> b) -> a -> b
$ \Type
_ -> String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"x"
  let body = Q Exp -> Q Exp
forall {m :: * -> *}. Quote m => m Exp -> m Exp
toJSONExp (Q Exp -> Q Exp) -> Q Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Q Exp] -> Q Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
tupE
        [ [| base :: String |]
         -- The singleton is special-cased because of
         -- https://downloads.haskell.org/ghc/8.10.1-rc1/docs/html/users_guide/8.10.1-notes.html#template-haskell
        , case [Name]
vars of
            [Name
v] -> Q Exp -> Q Exp
forall {m :: * -> *}. Quote m => m Exp -> m Exp
toJSONExp (Q Exp -> Q Exp) -> Q Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
v
            [Name]
vs -> [Q Exp] -> Q Exp
forall (m :: * -> *). Quote m => [m Exp] -> m Exp
tupE ([Q Exp] -> Q Exp) -> [Q Exp] -> Q Exp
forall a b. (a -> b) -> a -> b
$ (Name -> Q Exp) -> [Name] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Q Exp -> Q Exp
forall {m :: * -> *}. Quote m => m Exp -> m Exp
toJSONExp (Q Exp -> Q Exp) -> (Name -> Q Exp) -> Name -> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE) [Name]
vs
        ]
  _ <- conMatches ''ToJSON topVars lastVar c
  lift $ match (conP name (map varP vars)) (normalB body) []

-- | Implementation of 'parseJSON'
conMatchesParseJSON :: JSONGADTOptions -> [Name] -> ExpQ -> ConstructorInfo -> WriterT [Type] Q Match
conMatchesParseJSON :: JSONGADTOptions
-> [Name] -> Q Exp -> ConstructorInfo -> WriterT [Type] Q Match
conMatchesParseJSON JSONGADTOptions
opts [Name]
allTopVars Q Exp
e ConstructorInfo
c = do
  let (Set Name
topVars, Name
lastVar) = [Name] -> (Set Name, Name)
splitTopVars [Name]
allTopVars
  (pat, conApp) <- Name
-> Set Name
-> Name
-> ConstructorInfo
-> WriterT [Type] Q (Pat, Exp)
conMatches ''FromJSON Set Name
topVars Name
lastVar ConstructorInfo
c
  let match' = Q Pat -> Q Body -> [Q InstanceDec] -> Q Match
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m InstanceDec] -> m Match
match (Lit -> Q Pat
forall (m :: * -> *). Quote m => Lit -> m Pat
litP (String -> Lit
StringL (JSONGADTOptions -> String -> String
gadtConstructorModifier JSONGADTOptions
opts (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Name -> String
nameBase (ConstructorInfo -> Name
constructorName ConstructorInfo
c))))
      body = [Q Stmt] -> Q Exp
forall (m :: * -> *). Quote m => [m Stmt] -> m Exp
doE [ Q Pat -> Q Exp -> Q Stmt
forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Stmt
bindS (Pat -> Q Pat
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Pat
pat) [| parseJSON $Q Exp
e |]
                 , Q Exp -> Q Stmt
forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS [| return (Some $(Exp -> Q Exp
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
conApp)) |]
                 ]
  lift $ match' (normalB body) []

conMatches
  :: Name -- ^ Name of class (''ToJSON or ''FromJSON)
  -> Set Name -- Names of variables used in the instance head in argument order
  -> Name -- Final type variable (the index type)
  -> ConstructorInfo
  -> WriterT [Type] Q (Pat, Exp)
conMatches :: Name
-> Set Name
-> Name
-> ConstructorInfo
-> WriterT [Type] Q (Pat, Exp)
conMatches Name
clsName Set Name
topVars Name
ixVar ConstructorInfo
c = do
  let name :: Name
name = ConstructorInfo -> Name
constructorName ConstructorInfo
c
      types :: [Type]
types = ConstructorInfo -> [Type]
constructorFields ConstructorInfo
c
      ([Type]
constraints, [Type]
equalities') = ((Type -> Bool) -> [Type] -> ([Type], [Type]))
-> [Type] -> (Type -> Bool) -> ([Type], [Type])
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Type -> Bool) -> [Type] -> ([Type], [Type])
forall a. (a -> Bool) -> [a] -> ([a], [a])
partition (ConstructorInfo -> [Type]
constructorContext ConstructorInfo
c) ((Type -> Bool) -> ([Type], [Type]))
-> (Type -> Bool) -> ([Type], [Type])
forall a b. (a -> b) -> a -> b
$ \case
        AppT (AppT Type
EqualityT Type
_) Type
_ -> Bool
False
        Type
_ -> Bool
True
      equalities :: [(Type, Type)]
equalities = [[(Type, Type)]] -> [(Type, Type)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ [(Type
a, Type
b), (Type
b, Type
a)] | AppT (AppT Type
EqualityT Type
a) Type
b <- [Type]
equalities' ]
  unifiedEqualities :: [Map Name Type] <- Q [Map Name Type] -> WriterT [Type] Q [Map Name Type]
forall (m :: * -> *) a. Monad m => m a -> WriterT [Type] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q [Map Name Type] -> WriterT [Type] Q [Map Name Type])
-> Q [Map Name Type] -> WriterT [Type] Q [Map Name Type]
forall a b. (a -> b) -> a -> b
$ [(Type, Type)]
-> ((Type, Type) -> Q (Map Name Type)) -> Q [Map Name Type]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [(Type, Type)]
equalities (((Type, Type) -> Q (Map Name Type)) -> Q [Map Name Type])
-> ((Type, Type) -> Q (Map Name Type)) -> Q [Map Name Type]
forall a b. (a -> b) -> a -> b
$ \(Type
a, Type
b) -> [Type] -> Q (Map Name Type)
unifyTypes [Type
a, Type
b]
  let rigidImplications :: Map Name (Set Name)
      rigidImplications = (Set Name -> Set Name -> Set Name)
-> [Map Name (Set Name)] -> Map Name (Set Name)
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.union ([Map Name (Set Name)] -> Map Name (Set Name))
-> [Map Name (Set Name)] -> Map Name (Set Name)
forall a b. (a -> b) -> a -> b
$ (Type -> Set Name) -> Map Name Type -> Map Name (Set Name)
forall a b. (a -> b) -> Map Name a -> Map Name b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Set Name
freeTypeVariables (Map Name Type -> Map Name (Set Name))
-> [Map Name Type] -> [Map Name (Set Name)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Map Name Type]
unifiedEqualities
  let expandRigids :: Set Name -> Set Name
      expandRigids Set Name
rigids = Set Name -> Set Name -> Set Name
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set Name
rigids (Set Name -> Set Name) -> Set Name -> Set Name
forall a b. (a -> b) -> a -> b
$ [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
$ Map Name (Set Name) -> [Set Name]
forall k a. Map k a -> [a]
Map.elems (Map Name (Set Name) -> [Set Name])
-> Map Name (Set Name) -> [Set Name]
forall a b. (a -> b) -> a -> b
$
        Map Name (Set Name) -> Set Name -> Map Name (Set Name)
forall k v. Ord k => Map k v -> Set k -> Map k v
restrictKeys Map Name (Set Name)
rigidImplications Set Name
rigids
      expandRigidsFully Set Name
rigids =
        let rigids' :: Set Name
rigids' = Set Name -> Set Name
expandRigids Set Name
rigids
        in if Set Name
rigids' Set Name -> Set Name -> Bool
forall a. Eq a => a -> a -> Bool
== Set Name
rigids then Set Name
rigids else Set Name -> Set Name
expandRigidsFully Set Name
rigids'
      rigidVars = Set Name -> Set Name
expandRigidsFully Set Name
topVars
      ixSpecialization :: Map Name Type
      ixSpecialization = Map Name Type -> Set Name -> Map Name Type
forall k v. Ord k => Map k v -> Set k -> Map k v
restrictKeys ([Map Name Type] -> Map Name Type
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
Map.unions [Map Name Type]
unifiedEqualities) (Set Name -> Map Name Type) -> Set Name -> Map Name Type
forall a b. (a -> b) -> a -> b
$ Name -> Set Name
forall a. a -> Set a
Set.singleton Name
ixVar
      -- We filter out constraints which don't mention variables from the instance head mostly to avoid warnings,
      -- but a good deal more of these occur than one might expect due to the normalisation done by reifyDatatype.
      tellCxt a
cs = do
        a -> WriterT a m ()
forall (m :: * -> *) w. Monad m => w -> WriterT w m ()
tell (a -> WriterT a m ()) -> a -> WriterT a m ()
forall a b. (a -> b) -> a -> b
$ Map Name Type -> a -> a
forall a. TypeSubstitution a => Map Name Type -> a -> a
applySubstitution Map Name Type
ixSpecialization a
cs
  tellCxt constraints
  vars <- forM types $ \Type
typ -> do
    x <- Q Name -> WriterT [Type] Q Name
forall (m :: * -> *) a. Monad m => m a -> WriterT [Type] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q Name -> WriterT [Type] Q Name)
-> Q Name -> WriterT [Type] Q Name
forall a b. (a -> b) -> a -> b
$ String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"x"
    let demandInstanceIfNecessary = do
          insts <- Q [InstanceDec] -> WriterT [Type] Q [InstanceDec]
forall (m :: * -> *) a. Monad m => m a -> WriterT [Type] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q [InstanceDec] -> WriterT [Type] Q [InstanceDec])
-> Q [InstanceDec] -> WriterT [Type] Q [InstanceDec]
forall a b. (a -> b) -> a -> b
$ Set Name -> Name -> [Type] -> Q [InstanceDec]
reifyInstancesWithRigids Set Name
rigidVars Name
clsName [Type
typ]
          case insts of
            [] -> [Type] -> WriterT [Type] Q ()
forall {m :: * -> *} {a}.
(Monad m, TypeSubstitution a) =>
a -> WriterT a m ()
tellCxt [Type -> Type -> Type
AppT (Name -> Type
ConT Name
clsName) Type
typ]
            [InstanceD Maybe Overlap
_ [Type]
cxt (AppT Type
_className Type
ityp) [InstanceDec]
_] -> do
              sub <- Q (Map Name Type) -> WriterT [Type] Q (Map Name Type)
forall (m :: * -> *) a. Monad m => m a -> WriterT [Type] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q (Map Name Type) -> WriterT [Type] Q (Map Name Type))
-> Q (Map Name Type) -> WriterT [Type] Q (Map Name Type)
forall a b. (a -> b) -> a -> b
$ [Type] -> Q (Map Name Type)
unifyTypes [Type
ityp, Type
typ]
              tellCxt $ applySubstitution sub cxt

            [InstanceDec]
_ -> String -> WriterT [Type] Q ()
forall a. HasCallStack => String -> a
error (String -> WriterT [Type] Q ()) -> String -> WriterT [Type] Q ()
forall a b. (a -> b) -> a -> b
$ String
"The following instances of " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show Name
clsName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" for " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
typ String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" exist (rigids: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [String] -> String
unwords ((Name -> String) -> [Name] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Name -> String
forall a. Show a => a -> String
show ([Name] -> [String]) -> [Name] -> [String]
forall a b. (a -> b) -> a -> b
$ Set Name -> [Name]
forall a. Set a -> [a]
Set.toList Set Name
rigidVars) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"), and I don't know which to pick:\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ [String] -> String
unlines ((InstanceDec -> String) -> [InstanceDec] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> String
forall a. Show a => a -> String
show (Doc -> String) -> (InstanceDec -> Doc) -> InstanceDec -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InstanceDec -> Doc
forall a. Ppr a => a -> Doc
ppr) [InstanceDec]
insts)
    case typ of
      AppT Type
tn (VarT Name
_) -> do
        -- This may be a nested GADT, so check for special FromJSON instance
        insts <- Q [InstanceDec] -> WriterT [Type] Q [InstanceDec]
forall (m :: * -> *) a. Monad m => m a -> WriterT [Type] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q [InstanceDec] -> WriterT [Type] Q [InstanceDec])
-> Q [InstanceDec] -> WriterT [Type] Q [InstanceDec]
forall a b. (a -> b) -> a -> b
$ Set Name -> Name -> [Type] -> Q [InstanceDec]
reifyInstancesWithRigids Set Name
rigidVars Name
clsName [Type -> Type -> Type
AppT (Name -> Type
ConT ''Some) Type
tn]
        case insts of
          [] -> do
            -- No special instance, try to check the type for a straightforward one, since if we don't have one, we need to demand it.
            WriterT [Type] Q ()
demandInstanceIfNecessary
            (Pat, Exp) -> WriterT [Type] Q (Pat, Exp)
forall a. a -> WriterT [Type] Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> Pat
VarP Name
x, Name -> Exp
VarE Name
x)
          [InstanceD Maybe Overlap
_ [Type]
cxt (AppT Type
_className (AppT (ConT Name
_some) Type
ityp)) [InstanceDec]
_] -> do
            sub <- Q (Map Name Type) -> WriterT [Type] Q (Map Name Type)
forall (m :: * -> *) a. Monad m => m a -> WriterT [Type] m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (Q (Map Name Type) -> WriterT [Type] Q (Map Name Type))
-> Q (Map Name Type) -> WriterT [Type] Q (Map Name Type)
forall a b. (a -> b) -> a -> b
$ [Type] -> Q (Map Name Type)
unifyTypes [Type
ityp, Type
tn]
            tellCxt $ applySubstitution sub cxt
            return ( ConP
                       'Some
#if MIN_VERSION_template_haskell(2,18,0)
                       []
#endif
                       [VarP x]
                   , VarE x
                   )
          [InstanceDec]
_ -> String -> WriterT [Type] Q (Pat, Exp)
forall a. HasCallStack => String -> a
error (String -> WriterT [Type] Q (Pat, Exp))
-> String -> WriterT [Type] Q (Pat, Exp)
forall a b. (a -> b) -> a -> b
$ String
"The following instances of " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show Name
clsName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" for " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Doc -> String
forall a. Show a => a -> String
show ([Type] -> Doc
forall a. Ppr a => a -> Doc
ppr [Type -> Type -> Type
AppT (Name -> Type
ConT ''Some) Type
tn]) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" exist (rigids: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [String] -> String
unwords ((Name -> String) -> [Name] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Name -> String
forall a. Show a => a -> String
show ([Name] -> [String]) -> [Name] -> [String]
forall a b. (a -> b) -> a -> b
$ Set Name -> [Name]
forall a. Set a -> [a]
Set.toList Set Name
rigidVars) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"), and I don't know which to pick:\n" String -> String -> String
forall a. [a] -> [a] -> [a]
++ [String] -> String
unlines ((InstanceDec -> String) -> [InstanceDec] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (Doc -> String
forall a. Show a => a -> String
show (Doc -> String) -> (InstanceDec -> Doc) -> InstanceDec -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. InstanceDec -> Doc
forall a. Ppr a => a -> Doc
ppr) [InstanceDec]
insts)
      Type
_ -> do
        WriterT [Type] Q ()
demandInstanceIfNecessary
        (Pat, Exp) -> WriterT [Type] Q (Pat, Exp)
forall a. a -> WriterT [Type] Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> Pat
VarP Name
x, Name -> Exp
VarE Name
x)
  -- The singleton is special-cased because of
  -- https://downloads.haskell.org/ghc/8.10.1-rc1/docs/html/users_guide/8.10.1-notes.html#template-haskell
  let pat = case [(Pat, Exp)]
vars of
        [(Pat, Exp)
v] -> (Pat, Exp) -> Pat
forall a b. (a, b) -> a
fst (Pat, Exp)
v
        [(Pat, Exp)]
vs -> [Pat] -> Pat
TupP (((Pat, Exp) -> Pat) -> [(Pat, Exp)] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map (Pat, Exp) -> Pat
forall a b. (a, b) -> a
fst [(Pat, Exp)]
vs)
      conApp = (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Exp -> Exp
AppE (Name -> Exp
ConE Name
name) (((Pat, Exp) -> Exp) -> [(Pat, Exp)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Pat, Exp) -> Exp
forall a b. (a, b) -> b
snd [(Pat, Exp)]
vars)
  return (pat, conApp)

-----------------------------------------------------------------------------------------------------

-- | Determine the arity of a kind.
kindArity :: Kind -> Int
kindArity :: Type -> Int
kindArity = \case
  ForallT [TyVarBndr Specificity]
_ [Type]
_ Type
t -> Type -> Int
kindArity Type
t
  AppT (AppT Type
ArrowT Type
_) Type
t -> Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Type -> Int
kindArity Type
t
  SigT Type
t Type
_ -> Type -> Int
kindArity Type
t
  ParensT Type
t -> Type -> Int
kindArity Type
t
  Type
_ -> Int
0

-- | Given the name of a type constructor, determine a list of type variables bound as parameters by
-- its declaration, and the arity of the kind of type being defined (i.e. how many more arguments would
-- need to be supplied in addition to the bound parameters in order to obtain an ordinary type of kind *).
-- If the supplied 'Name' is anything other than a data or newtype, produces an error.
tyConArity' :: Name -> Q ([TyVarBndrVis], Int)
tyConArity' :: Name -> Q ([TyVarBndrVis], Int)
tyConArity' Name
n = Name -> Q Info
reify Name
n Q Info
-> (Info -> Q ([TyVarBndrVis], Int)) -> Q ([TyVarBndrVis], Int)
forall a b. Q a -> (a -> Q b) -> Q b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ([TyVarBndrVis], Int) -> Q ([TyVarBndrVis], Int)
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (([TyVarBndrVis], Int) -> Q ([TyVarBndrVis], Int))
-> (Info -> ([TyVarBndrVis], Int))
-> Info
-> Q ([TyVarBndrVis], Int)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. \case
  TyConI (DataD [Type]
_ Name
_ [TyVarBndrVis]
ts Maybe Type
mk [Con]
_ [DerivClause]
_) -> ([TyVarBndrVis]
ts, Int -> (Type -> Int) -> Maybe Type -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 Type -> Int
kindArity Maybe Type
mk)
  TyConI (NewtypeD [Type]
_ Name
_ [TyVarBndrVis]
ts Maybe Type
mk Con
_ [DerivClause]
_) -> ([TyVarBndrVis]
ts, Int -> (Type -> Int) -> Maybe Type -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 Type -> Int
kindArity Maybe Type
mk)
  Info
_ -> String -> ([TyVarBndrVis], Int)
forall a. HasCallStack => String -> a
error (String -> ([TyVarBndrVis], Int))
-> String -> ([TyVarBndrVis], Int)
forall a b. (a -> b) -> a -> b
$ String
"tyConArity': Supplied name reified to something other than a data declaration: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show Name
n


-----------------------------------------------------------------------------------------------------

restrictKeys :: Ord k => Map k v -> Set k -> Map k v
restrictKeys :: forall k v. Ord k => Map k v -> Set k -> Map k v
restrictKeys Map k v
m Set k
s =
#if MIN_VERSION_containers(0,5,8)
  Map k v -> Set k -> Map k v
forall k v. Ord k => Map k v -> Set k -> Map k v
Map.restrictKeys Map k v
m Set k
s
#else
  Map.intersection m $ Map.fromSet (const ()) s
#endif