{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskell #-}

module Data.Constraint.Extras.TH (deriveArgDict, deriveArgDictV, gadtIndices) where

import Data.Constraint
import Data.Constraint.Extras
import Data.Maybe
import Control.Monad
import Language.Haskell.TH

deriveArgDict :: Name -> Q [Dec]
deriveArgDict :: Name -> Q [Dec]
deriveArgDict n :: Name
n = do
  Name
c <- String -> Q Name
newName "c"
  [Either Type Type]
ts <- Name -> Name -> Q [Either Type Type]
gadtIndices Name
c Name
n
  let xs :: [Type]
xs = ((Either Type Type -> Type) -> [Either Type Type] -> [Type])
-> [Either Type Type] -> (Either Type Type -> Type) -> [Type]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Either Type Type -> Type) -> [Either Type Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map [Either Type Type]
ts ((Either Type Type -> Type) -> [Type])
-> (Either Type Type -> Type) -> [Type]
forall a b. (a -> b) -> a -> b
$ \case
        Left t :: Type
t -> Type -> Type -> Type
AppT (Type -> Type -> Type
AppT (Name -> Type
ConT ''ConstraintsFor) Type
t) (Name -> Type
VarT Name
c)
        Right t :: Type
t -> (Type -> Type -> Type
AppT (Name -> Type
VarT Name
c) Type
t)
      l :: Int
l = [Type] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Type]
xs
      constraints :: Type
constraints = (Type -> Type -> Type) -> Type -> [Type] -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
AppT (Int -> Type
TupleT Int
l) [Type]
xs
  Int
arity <- Name -> Q Int
tyConArity Name
n
  [Name]
tyVars <- Int -> Q Name -> Q [Name]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Int
arity Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1) (String -> Q Name
newName "a")
  let n' :: Type
n' = (Name -> Type -> Type) -> Type -> [Name] -> Type
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\v :: Name
v x :: Type
x -> Type -> Type -> Type
AppT Type
x (Name -> Type
VarT Name
v)) (Name -> Type
ConT Name
n) [Name]
tyVars
  [d| instance ArgDict $(varT c) $(pure n') where
        type ConstraintsFor  $(pure n') $(varT c) = $(pure constraints)
        argDict = $(LamCaseE <$> matches c n 'argDict)
    |]

{-# DEPRECATED deriveArgDictV "Just use 'deriveArgDict'" #-}
deriveArgDictV :: Name -> Q [Dec]
deriveArgDictV :: Name -> Q [Dec]
deriveArgDictV = Name -> Q [Dec]
deriveArgDict

matches :: Name -> Name -> Name -> Q [Match]
matches :: Name -> Name -> Name -> Q [Match]
matches c :: Name
c n :: Name
n argDictName :: Name
argDictName = do
  Name
x <- String -> Q Name
newName "x"
  Name -> Q Info
reify Name
n Q Info -> (Info -> Q [Match]) -> Q [Match]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    TyConI (DataD _ _ _ _ constrs :: [Con]
constrs _) -> ([[Match]] -> [Match]) -> Q [[Match]] -> Q [Match]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Match]] -> [Match]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Q [[Match]] -> Q [Match]) -> Q [[Match]] -> Q [Match]
forall a b. (a -> b) -> a -> b
$ [Con] -> (Con -> Q [Match]) -> Q [[Match]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Con]
constrs ((Con -> Q [Match]) -> Q [[Match]])
-> (Con -> Q [Match]) -> Q [[Match]]
forall a b. (a -> b) -> a -> b
$ \case
      GadtC [name :: Name
name] _ _ -> [Match] -> Q [Match]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Match] -> Q [Match]) -> [Match] -> Q [Match]
forall a b. (a -> b) -> a -> b
$
        [Pat -> Body -> [Dec] -> Match
Match (Name -> [FieldPat] -> Pat
RecP Name
name []) (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE 'Dict) []]
      ForallC _ _ (GadtC [name :: Name
name] bts :: [BangType]
bts (AppT _ (VarT b :: Name
b))) -> do
        [Maybe Name]
ps <- [BangType] -> (BangType -> Q (Maybe Name)) -> Q [Maybe Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [BangType]
bts ((BangType -> Q (Maybe Name)) -> Q [Maybe Name])
-> (BangType -> Q (Maybe Name)) -> Q [Maybe Name]
forall a b. (a -> b) -> a -> b
$ \case
          (_, AppT t :: Type
t (VarT b' :: Name
b')) | Name
b Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
b' -> do
            Bool
hasArgDictInstance <- Bool -> Bool
not (Bool -> Bool) -> ([Dec] -> Bool) -> [Dec] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Dec] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Dec] -> Bool) -> Q [Dec] -> Q Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> [Type] -> Q [Dec]
reifyInstances ''ArgDict [Name -> Type
VarT Name
c, Type
t]
            Maybe Name -> Q (Maybe Name)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Name -> Q (Maybe Name)) -> Maybe Name -> Q (Maybe Name)
forall a b. (a -> b) -> a -> b
$ if Bool
hasArgDictInstance
              then Name -> Maybe Name
forall a. a -> Maybe a
Just Name
x
              else Maybe Name
forall a. Maybe a
Nothing
          _ -> Maybe Name -> Q (Maybe Name)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Name
forall a. Maybe a
Nothing
        [Match] -> Q [Match]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Match] -> Q [Match]) -> [Match] -> Q [Match]
forall a b. (a -> b) -> a -> b
$ case [Maybe Name] -> [Name]
forall a. [Maybe a] -> [a]
catMaybes [Maybe Name]
ps of
          [] -> [Pat -> Body -> [Dec] -> Match
Match (Name -> [FieldPat] -> Pat
RecP Name
name []) (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE 'Dict) []]
          (v :: Name
v:_) ->
            let patf :: Maybe a -> (Bool -> [Pat]) -> Bool -> [Pat]
patf = \v' :: Maybe a
v' rest :: Bool -> [Pat]
rest done :: Bool
done -> if Bool
done
                  then Pat
WildP Pat -> [Pat] -> [Pat]
forall a. a -> [a] -> [a]
: Bool -> [Pat]
rest Bool
done
                  else case Maybe a
v' of
                    Nothing -> Pat
WildP Pat -> [Pat] -> [Pat]
forall a. a -> [a] -> [a]
: Bool -> [Pat]
rest Bool
done
                    Just _ -> Name -> Pat
VarP Name
v Pat -> [Pat] -> [Pat]
forall a. a -> [a] -> [a]
: Bool -> [Pat]
rest Bool
True
                pat :: [Pat]
pat = (Maybe Name -> (Bool -> [Pat]) -> Bool -> [Pat])
-> (Bool -> [Pat]) -> [Maybe Name] -> Bool -> [Pat]
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Maybe Name -> (Bool -> [Pat]) -> Bool -> [Pat]
forall a. Maybe a -> (Bool -> [Pat]) -> Bool -> [Pat]
patf ([Pat] -> Bool -> [Pat]
forall a b. a -> b -> a
const []) [Maybe Name]
ps Bool
False
            in [Pat -> Body -> [Dec] -> Match
Match (Name -> [Pat] -> Pat
ConP Name
name [Pat]
pat) (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Exp -> Exp -> Exp
AppE (Name -> Exp
VarE Name
argDictName) (Name -> Exp
VarE Name
v)) []]
      ForallC _ _ (GadtC [name :: Name
name] _ _) -> [Match] -> Q [Match]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Match] -> Q [Match]) -> [Match] -> Q [Match]
forall a b. (a -> b) -> a -> b
$
        [Pat -> Body -> [Dec] -> Match
Match (Name -> [FieldPat] -> Pat
RecP Name
name []) (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE 'Dict) []]
      a :: Con
a -> String -> Q [Match]
forall a. HasCallStack => String -> a
error (String -> Q [Match]) -> String -> Q [Match]
forall a b. (a -> b) -> a -> b
$ "deriveArgDict matches: Unmatched 'Dec': " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Con -> String
forall a. Show a => a -> String
show Con
a
    a :: Info
a -> String -> Q [Match]
forall a. HasCallStack => String -> a
error (String -> Q [Match]) -> String -> Q [Match]
forall a b. (a -> b) -> a -> b
$ "deriveArgDict matches: Unmatched 'Info': " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Info -> String
forall a. Show a => a -> String
show Info
a

kindArity :: Kind -> Int
kindArity :: Type -> Int
kindArity = \case
  ForallT _ _ t :: Type
t -> Type -> Int
kindArity Type
t
  AppT (AppT ArrowT _) t :: Type
t -> 1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Type -> Int
kindArity Type
t
  SigT t :: Type
t _ -> Type -> Int
kindArity Type
t
  ParensT t :: Type
t -> Type -> Int
kindArity Type
t
  _ -> 0

tyConArity :: Name -> Q Int
tyConArity :: Name -> Q Int
tyConArity n :: Name
n = Name -> Q Info
reify Name
n Q Info -> (Info -> Q Int) -> Q Int
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> Q Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Q Int) -> (Info -> Int) -> Info -> Q Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. \case
   TyConI (DataD _ _ ts :: [TyVarBndr]
ts mk :: Maybe Type
mk _ _) -> Int -> Maybe Int -> Int
forall a. a -> Maybe a -> a
fromMaybe 0 ((Type -> Int) -> Maybe Type -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Int
kindArity Maybe Type
mk) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [TyVarBndr] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TyVarBndr]
ts
   _ -> String -> Int
forall a. HasCallStack => String -> a
error (String -> Int) -> String -> Int
forall a b. (a -> b) -> a -> b
$ "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

gadtIndices :: Name -> Name -> Q [Either Type Type]
gadtIndices :: Name -> Name -> Q [Either Type Type]
gadtIndices c :: Name
c n :: Name
n = Name -> Q Info
reify Name
n Q Info -> (Info -> Q [Either Type Type]) -> Q [Either Type Type]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  TyConI (DataD _ _ _ _ constrs :: [Con]
constrs _) -> ([[Either Type Type]] -> [Either Type Type])
-> Q [[Either Type Type]] -> Q [Either Type Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Either Type Type]] -> [Either Type Type]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Q [[Either Type Type]] -> Q [Either Type Type])
-> Q [[Either Type Type]] -> Q [Either Type Type]
forall a b. (a -> b) -> a -> b
$ [Con] -> (Con -> Q [Either Type Type]) -> Q [[Either Type Type]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Con]
constrs ((Con -> Q [Either Type Type]) -> Q [[Either Type Type]])
-> (Con -> Q [Either Type Type]) -> Q [[Either Type Type]]
forall a b. (a -> b) -> a -> b
$ \case
    GadtC _ _ (AppT _ typ :: Type
typ) -> [Either Type Type] -> Q [Either Type Type]
forall (m :: * -> *) a. Monad m => a -> m a
return [Type -> Either Type Type
forall a b. b -> Either a b
Right Type
typ]
    ForallC _ _ (GadtC _ bts :: [BangType]
bts (AppT _ (VarT _))) -> ([[Either Type Type]] -> [Either Type Type])
-> Q [[Either Type Type]] -> Q [Either Type Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Either Type Type]] -> [Either Type Type]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Q [[Either Type Type]] -> Q [Either Type Type])
-> Q [[Either Type Type]] -> Q [Either Type Type]
forall a b. (a -> b) -> a -> b
$ [BangType]
-> (BangType -> Q [Either Type Type]) -> Q [[Either Type Type]]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [BangType]
bts ((BangType -> Q [Either Type Type]) -> Q [[Either Type Type]])
-> (BangType -> Q [Either Type Type]) -> Q [[Either Type Type]]
forall a b. (a -> b) -> a -> b
$ \case
      (_, AppT t :: Type
t (VarT _)) -> do
        Bool
hasArgDictInstance <- ([Dec] -> Bool) -> Q [Dec] -> Q Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool -> Bool
not (Bool -> Bool) -> ([Dec] -> Bool) -> [Dec] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Dec] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) (Q [Dec] -> Q Bool) -> Q [Dec] -> Q Bool
forall a b. (a -> b) -> a -> b
$ Name -> [Type] -> Q [Dec]
reifyInstances ''ArgDict [Name -> Type
VarT Name
c, Type
t]
        [Either Type Type] -> Q [Either Type Type]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Either Type Type] -> Q [Either Type Type])
-> [Either Type Type] -> Q [Either Type Type]
forall a b. (a -> b) -> a -> b
$ if Bool
hasArgDictInstance then [Type -> Either Type Type
forall a b. a -> Either a b
Left Type
t] else []
      _ -> [Either Type Type] -> Q [Either Type Type]
forall (m :: * -> *) a. Monad m => a -> m a
return []
    ForallC _ _ (GadtC _ _ (AppT _ typ :: Type
typ)) -> [Either Type Type] -> Q [Either Type Type]
forall (m :: * -> *) a. Monad m => a -> m a
return [Type -> Either Type Type
forall a b. b -> Either a b
Right Type
typ]
    _ -> [Either Type Type] -> Q [Either Type Type]
forall (m :: * -> *) a. Monad m => a -> m a
return []
  a :: Info
a -> String -> Q [Either Type Type]
forall a. HasCallStack => String -> a
error (String -> Q [Either Type Type]) -> String -> Q [Either Type Type]
forall a b. (a -> b) -> a -> b
$ "gadtIndices: Unmatched 'Info': " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Info -> String
forall a. Show a => a -> String
show Info
a