    {-# LANGUAGE CPP #-}
    {-# 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 Name
    n = do
      (typeHead, constrs) <- Name -> Q (Type, [Con])
    getDeclInfo Name
    n
      c <- newName "c"
      ts <- gadtIndices c constrs
      let constraints = ((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 Type
    t -> Type -> Type -> Type
    AppT (Type -> Type -> Type
    AppT (Name -> Type
    ConT ''Has) (Name -> Type
    VarT Name
    c)) Type
    t
            Right Type
    t -> (Type -> Type -> Type
    AppT (Name -> Type
    VarT Name
    c) Type
    t)
      ms <- matches c constrs 'argDict
      return
        [ InstanceD Nothing constraints (AppT (AppT (ConT ''Has) (VarT c)) typeHead)
          [ ValD (VarP 'argDict) (NormalB (LamCaseE ms)) [] ]
        ]

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

    matches :: Name -> [Con] -> Name -> Q [Match]
    matches :: Name -> [Con] -> Name -> Q [Match]
    matches Name
    c [Con]
    constrs Name
    argDictName = do
      x <- String -> Q Name
    forall (m :: * -> *). Quote m => String -> m Name
    newName String
    "x"
      fmap concat . forM constrs $ \case
        GadtC [Name
    name] [BangType]
    _ Type
    _ -> [Match] -> Q [Match]
    forall a. a -> Q a
    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 [TyVarBndr Specificity]
    _ [Type]
    _ (GadtC [Name
    name] [BangType]
    bts (AppT Type
    _ (VarT Name
    b))) -> do
          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
            (Bang
    _, AppT Type
    t (VarT Name
    b')) | Name
    b Name -> Name -> Bool
    forall a. Eq a => a -> a -> Bool
    == Name
    b' -> do
              hasArgDictInstance <- Bool -> Bool
    not (Bool -> Bool) -> ([Dec] -> Bool) -> [Dec] -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [Dec] -> Bool
    forall a. [a] -> 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 ''Has [Name -> Type
    VarT Name
    c, Type
    t]
              return $ if hasArgDictInstance
                then Just x
                else Nothing
            BangType
    _ -> Maybe Name -> Q (Maybe Name)
    forall a. a -> Q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return Maybe Name
    forall a. Maybe a
    Nothing
          return $ case catMaybes 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) []]
            (Name
    v:[Name]
    _) ->
              let patf :: Maybe a -> (Bool -> [Pat]) -> Bool -> [Pat]
    patf = \Maybe a
    v' Bool -> [Pat]
    rest 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
                      Maybe a
    Nothing -> Pat
    WildP Pat -> [Pat] -> [Pat]
    forall a. a -> [a] -> [a]
    : Bool -> [Pat]
    rest Bool
    done
                      Just a
    _ -> 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 a b. (a -> b -> b) -> b -> [a] -> b
    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
    conPCompat 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 [TyVarBndr Specificity]
    _ [Type]
    _ (GadtC [Name
    name] [BangType]
    _ Type
    _) -> [Match] -> Q [Match]
    forall a. a -> Q a
    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) []]
        Con
    a -> String -> Q [Match]
    forall a. HasCallStack => String -> a
    error (String -> Q [Match]) -> String -> Q [Match]
    forall a b. (a -> b) -> a -> b
    $ String
    "deriveArgDict matches: Unmatched 'Dec': " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Con -> String
    forall a. Show a => a -> String
    show Con
    a

    conPCompat :: Name -> [Pat] -> Pat
    conPCompat :: Name -> [Pat] -> Pat
    conPCompat Name
    name =
      Name -> [Type] -> [Pat] -> Pat
    ConP
        Name
    name
    #if MIN_VERSION_template_haskell(2, 18, 0)
        []
    #endif

    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

    getDeclInfo :: Name -> Q (Type, [Con])
    getDeclInfo :: Name -> Q (Type, [Con])
    getDeclInfo Name
    n = Name -> Q Info
    reify Name
    n Q Info -> (Info -> Q (Type, [Con])) -> Q (Type, [Con])
    forall a b. Q a -> (a -> Q b) -> Q b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= \case
      TyConI (DataD [Type]
    _ Name
    _ [TyVarBndr BndrVis]
    ts Maybe Type
    mk [Con]
    constrs [DerivClause]
    _) -> do
        let arity :: Int
    arity = Int -> Maybe Int -> Int
    forall a. a -> Maybe a -> a
    fromMaybe Int
    0 ((Type -> Int) -> Maybe Type -> Maybe Int
    forall a b. (a -> b) -> Maybe a -> Maybe b
    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 BndrVis] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [TyVarBndr BndrVis]
    ts
        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
    - Int
    1) (String -> Q Name
    forall (m :: * -> *). Quote m => String -> m Name
    newName String
    "a")
        let typeHead = (Name -> Type -> Type) -> Type -> [Name] -> Type
    forall a b. (a -> b -> b) -> b -> [a] -> b
    forall (t :: * -> *) a b.
    Foldable t =>
    (a -> b -> b) -> b -> t a -> b
    foldr (\Name
    v Type
    x -> Type -> Type -> Type
    AppT Type
    x (Name -> Type
    VarT Name
    v)) (Name -> Type
    ConT Name
    n) [Name]
    tyVars
        return (typeHead, constrs)
      DataConI Name
    _ (AppT Type
    typeHead Type
    _) Name
    parent -> Type -> Name -> Q (Type, [Con])
    forall {a}. a -> Name -> Q (a, [Con])
    handleParent Type
    typeHead Name
    parent
      DataConI Name
    _ (ForallT [TyVarBndr Specificity]
    _ [Type]
    _ (AppT Type
    typeHead Type
    _)) Name
    parent -> Type -> Name -> Q (Type, [Con])
    forall {a}. a -> Name -> Q (a, [Con])
    handleParent Type
    typeHead Name
    parent
      Info
    a -> String -> Q (Type, [Con])
    forall a. HasCallStack => String -> a
    error (String -> Q (Type, [Con])) -> String -> Q (Type, [Con])
    forall a b. (a -> b) -> a -> b
    $ String
    "getDeclInfo: Unmatched 'Info': " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Info -> String
    forall a. Show a => a -> String
    show Info
    a
      where
        handleParent :: a -> Name -> Q (a, [Con])
    handleParent a
    typeHead Name
    parent = Name -> Q Info
    reify Name
    parent Q Info -> (Info -> Q (a, [Con])) -> Q (a, [Con])
    forall a b. Q a -> (a -> Q b) -> Q b
    forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
    >>= \case
          FamilyI Dec
    _ [Dec]
    instances -> do
            let instCons :: InstanceDec -> [Con]
                instCons :: Dec -> [Con]
    instCons = \case
                  DataInstD [Type]
    _ Maybe [TyVarBndr ()]
    _ Type
    _ Maybe Type
    _ [Con]
    cons [DerivClause]
    _ -> [Con]
    cons
                  NewtypeInstD [Type]
    _ Maybe [TyVarBndr ()]
    _ Type
    _ Maybe Type
    _ Con
    con [DerivClause]
    _ -> [Con
    con]
                  Dec
    _ -> String -> [Con]
    forall a. HasCallStack => String -> a
    error (String -> [Con]) -> String -> [Con]
    forall a b. (a -> b) -> a -> b
    $ String
    "getDeclInfo: Expected a data or newtype family instance"
                conNames :: Con -> [Name]
                conNames :: Con -> [Name]
    conNames = \case
                  NormalC Name
    other [BangType]
    _ -> [Name
    other]
                  RecC Name
    other [VarBangType]
    _ -> [Name
    other]
                  InfixC BangType
    _ Name
    other BangType
    _ -> [Name
    other]
                  ForallC [TyVarBndr Specificity]
    _ [Type]
    _ Con
    con -> Con -> [Name]
    conNames Con
    con
                  GadtC [Name]
    others [BangType]
    _ Type
    _ -> [Name]
    others
                  RecGadtC [Name]
    others [VarBangType]
    _ Type
    _ -> [Name]
    others
                instHasThisConstructor :: Dec -> Bool
    instHasThisConstructor Dec
    i = (Name -> Bool) -> [Name] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (Name -> Name -> Bool
    forall a. Eq a => a -> a -> Bool
    == Name
    n) ([Name] -> Bool) -> [Name] -> Bool
    forall a b. (a -> b) -> a -> b
    $ Con -> [Name]
    conNames (Con -> [Name]) -> [Con] -> [Name]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Dec -> [Con]
    instCons Dec
    i
            case (Dec -> Bool) -> [Dec] -> [Dec]
    forall a. (a -> Bool) -> [a] -> [a]
    filter Dec -> Bool
    instHasThisConstructor [Dec]
    instances of
              [] -> String -> Q (a, [Con])
    forall a. HasCallStack => String -> a
    error (String -> Q (a, [Con])) -> String -> Q (a, [Con])
    forall a b. (a -> b) -> a -> b
    $ String
    "getDeclInfo: Couldn't find data family instance for constructor " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Name -> String
    forall a. Show a => a -> String
    show Name
    n
              l :: [Dec]
    l@(Dec
    _:Dec
    _:[Dec]
    _) -> String -> Q (a, [Con])
    forall a. HasCallStack => String -> a
    error (String -> Q (a, [Con])) -> String -> Q (a, [Con])
    forall a b. (a -> b) -> a -> b
    $ String
    "getDeclInfo: Expected one data family instance for constructor " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Name -> String
    forall a. Show a => a -> String
    show Name
    n String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ String
    " but found multiple: " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ [Dec] -> String
    forall a. Show a => a -> String
    show [Dec]
    l
              [Dec
    i] -> (a, [Con]) -> Q (a, [Con])
    forall a. a -> Q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return (a
    typeHead, Dec -> [Con]
    instCons Dec
    i)
          Info
    a -> String -> Q (a, [Con])
    forall a. HasCallStack => String -> a
    error (String -> Q (a, [Con])) -> String -> Q (a, [Con])
    forall a b. (a -> b) -> a -> b
    $ String
    "getDeclInfo: Unmatched parent of data family instance: " String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Info -> String
    forall a. Show a => a -> String
    show Info
    a

    gadtIndices :: Name -> [Con] -> Q [Either Type Type]
    gadtIndices :: Name -> [Con] -> Q [Either Type Type]
    gadtIndices Name
    c [Con]
    constrs = ([[Either Type Type]] -> [Either Type Type])
    -> Q [[Either Type Type]] -> Q [Either Type Type]
    forall a b. (a -> b) -> Q a -> Q b
    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 [Name]
    _ [BangType]
    _ (AppT Type
    _ Type
    typ) -> [Either Type Type] -> Q [Either Type Type]
    forall a. a -> Q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return [Type -> Either Type Type
    forall a b. b -> Either a b
    Right Type
    typ]
      ForallC [TyVarBndr Specificity]
    _ [Type]
    _ (GadtC [Name]
    _ [BangType]
    bts (AppT Type
    _ (VarT Name
    _))) -> ([[Either Type Type]] -> [Either Type Type])
    -> Q [[Either Type Type]] -> Q [Either Type Type]
    forall a b. (a -> b) -> Q a -> Q b
    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
        (Bang
    _, AppT Type
    t (VarT Name
    _)) -> do
          hasArgDictInstance <- ([Dec] -> Bool) -> Q [Dec] -> Q Bool
    forall a b. (a -> b) -> Q a -> Q b
    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 a. [a] -> 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 ''Has [Name -> Type
    VarT Name
    c, Type
    t]
          return $ if hasArgDictInstance then [Left t] else []
        BangType
    _ -> [Either Type Type] -> Q [Either Type Type]
    forall a. a -> Q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return []
      ForallC [TyVarBndr Specificity]
    _ [Type]
    _ (GadtC [Name]
    _ [BangType]
    _ (AppT Type
    _ Type
    typ)) -> [Either Type Type] -> Q [Either Type Type]
    forall a. a -> Q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return [Type -> Either Type Type
    forall a b. b -> Either a b
    Right Type
    typ]
      Con
    _ -> [Either Type Type] -> Q [Either Type Type]
    forall a. a -> Q a
    forall (m :: * -> *) a. Monad m => a -> m a
    return []