    {- Data/Singletons/Syntax.hs

    (c) Richard Eisenberg 2014
    rae@cs.brynmawr.edu

    Converts a list of DLetDecs into a LetDecEnv for easier processing,
    and contains various other AST definitions.
    -}

    {-# LANGUAGE DataKinds, TypeFamilies, PolyKinds, DeriveDataTypeable,
                 FlexibleInstances, ConstraintKinds #-}

    module Data.Singletons.Syntax where

    import Prelude hiding ( exp )
    import Data.Kind (Constraint, Type)
    import Language.Haskell.TH.Syntax hiding (Type)
    import Language.Haskell.TH.Desugar
    import qualified Language.Haskell.TH.Desugar.OMap.Strict as OMap
    import Language.Haskell.TH.Desugar.OMap.Strict (OMap)
    import Language.Haskell.TH.Desugar.OSet (OSet)

    type VarPromotions = [(Name, Name)] -- from term-level name to type-level name

    -- Information that is accumulated when promoting patterns.
    data PromDPatInfos = PromDPatInfos
      { PromDPatInfos -> VarPromotions
    prom_dpat_vars    :: VarPromotions
          -- Maps term-level pattern variables to their promoted, type-level counterparts.
      , PromDPatInfos -> OSet Name
    prom_dpat_sig_kvs :: OSet Name
          -- Kind variables bound by DSigPas.
          -- See Note [Explicitly binding kind variables] in Data.Singletons.Promote.Monad
      }

    instance Semigroup PromDPatInfos where
      PromDPatInfos vars1 :: VarPromotions
    vars1 sig_kvs1 :: OSet Name
    sig_kvs1 <> :: PromDPatInfos -> PromDPatInfos -> PromDPatInfos
    <> PromDPatInfos vars2 :: VarPromotions
    vars2 sig_kvs2 :: OSet Name
    sig_kvs2
        = VarPromotions -> OSet Name -> PromDPatInfos
    PromDPatInfos (VarPromotions
    vars1 VarPromotions -> VarPromotions -> VarPromotions
    forall a. Semigroup a => a -> a -> a
    <> VarPromotions
    vars2) (OSet Name
    sig_kvs1 OSet Name -> OSet Name -> OSet Name
    forall a. Semigroup a => a -> a -> a
    <> OSet Name
    sig_kvs2)

    instance Monoid PromDPatInfos where
      mempty :: PromDPatInfos
    mempty = VarPromotions -> OSet Name -> PromDPatInfos
    PromDPatInfos VarPromotions
    forall a. Monoid a => a
    mempty OSet Name
    forall a. Monoid a => a
    mempty

    -- A list of 'SingDSigPaInfos' is produced when singling pattern signatures, as we
    -- must case on the 'DExp's and match on them using the supplied 'DType's to
    -- bring the necessary singleton equality constraints into scope.
    -- See @Note [Singling pattern signatures]@.
    type SingDSigPaInfos = [(DExp, DType)]

    -- The parts of data declarations that are relevant to singletons.
    data DataDecl = DataDecl Name [DTyVarBndr] [DCon]

    -- The parts of type synonyms that are relevant to singletons.
    data TySynDecl = TySynDecl Name [DTyVarBndr] DType

    -- The parts of open type families that are relevant to singletons.
    type OpenTypeFamilyDecl = TypeFamilyDecl 'Open

    -- The parts of closed type families that are relevant to singletons.
    type ClosedTypeFamilyDecl = TypeFamilyDecl 'Closed

    -- The parts of type families that are relevant to singletons.
    newtype TypeFamilyDecl (info :: FamilyInfo)
      = TypeFamilyDecl { TypeFamilyDecl info -> DTypeFamilyHead
    getTypeFamilyDecl :: DTypeFamilyHead }
    -- Whether a type family is open or closed.
    data FamilyInfo = Open | Closed

    data ClassDecl ann = ClassDecl { ClassDecl ann -> DCxt
    cd_cxt  :: DCxt
                                   , ClassDecl ann -> Name
    cd_name :: Name
                                   , ClassDecl ann -> [DTyVarBndr]
    cd_tvbs :: [DTyVarBndr]
                                   , ClassDecl ann -> [FunDep]
    cd_fds  :: [FunDep]
                                   , ClassDecl ann -> LetDecEnv ann
    cd_lde  :: LetDecEnv ann
                                   }

    data InstDecl  ann = InstDecl { InstDecl ann -> DCxt
    id_cxt     :: DCxt
                                  , InstDecl ann -> Name
    id_name    :: Name
                                  , InstDecl ann -> DCxt
    id_arg_tys :: [DType]
                                  , InstDecl ann -> OMap Name DType
    id_sigs    :: OMap Name DType
                                  , InstDecl ann -> [(Name, LetDecRHS ann)]
    id_meths   :: [(Name, LetDecRHS ann)] }

    type UClassDecl = ClassDecl Unannotated
    type UInstDecl  = InstDecl Unannotated

    type AClassDecl = ClassDecl Annotated
    type AInstDecl  = InstDecl Annotated

    {-
    We see below several datatypes beginning with "A". These are annotated structures,
    necessary for Promote to communicate key things to Single. In particular, promotion
    of expressions is *not* deterministic, due to the necessity to create unique names
    for lets, cases, and lambdas. So, we put these promotions into an annotated AST
    so that Single can use the right promotions.
    -}

    -- A DExp with let, lambda, and type-signature nodes annotated with their
    -- type-level equivalents
    data ADExp = ADVarE Name
               | ADConE Name
               | ADLitE Lit
               | ADAppE ADExp ADExp
               | ADLamE [Name]         -- type-level names corresponding to term-level ones
                        DType          -- the promoted lambda
                        [Name] ADExp
               | ADCaseE ADExp [ADMatch] DType
                   -- the type is the return type
               | ADLetE ALetDecEnv ADExp
               | ADSigE DType          -- the promoted expression
                        ADExp DType

    -- A DPat with a pattern-signature node annotated with its type-level equivalent
    data ADPat = ADLitP Lit
               | ADVarP Name
               | ADConP Name [ADPat]
               | ADTildeP ADPat
               | ADBangP ADPat
               | ADSigP DType -- The promoted pattern. Will not contain any wildcards,
                              -- as per Note [Singling pattern signatures]
                        ADPat DType
               | ADWildP

    data ADMatch = ADMatch VarPromotions ADPat ADExp
    data ADClause = ADClause VarPromotions
                             [ADPat] ADExp

    data AnnotationFlag = Annotated | Unannotated

    -- These are used at the type-level exclusively
    type Annotated   = 'Annotated
    type Unannotated = 'Unannotated

    type family IfAnn (ann :: AnnotationFlag) (yes :: k) (no :: k) :: k where
      IfAnn Annotated   yes no = yes
      IfAnn Unannotated yes no = no

    data family LetDecRHS :: AnnotationFlag -> Type
    data instance LetDecRHS Annotated
      = AFunction DType  -- promote function (unapplied)
        Int    -- number of arrows in type
        [ADClause]
      | AValue DType -- promoted exp
        Int   -- number of arrows in type
        ADExp
    data instance LetDecRHS Unannotated = UFunction [DClause]
                                        | UValue DExp

    type ALetDecRHS = LetDecRHS Annotated
    type ULetDecRHS = LetDecRHS Unannotated

    data LetDecEnv ann = LetDecEnv
                       { LetDecEnv ann -> OMap Name (LetDecRHS ann)
    lde_defns :: OMap Name (LetDecRHS ann)
                       , LetDecEnv ann -> OMap Name DType
    lde_types :: OMap Name DType  -- type signatures
                       , LetDecEnv ann -> OMap Name Fixity
    lde_infix :: OMap Name Fixity -- infix declarations
                       , LetDecEnv ann -> IfAnn ann (OMap Name DType) ()
    lde_proms :: IfAnn ann (OMap Name DType) () -- possibly, promotions
                       , LetDecEnv ann -> IfAnn ann (OMap Name (OSet Name)) ()
    lde_bound_kvs :: IfAnn ann (OMap Name (OSet Name)) ()
                         -- The set of bound variables in scope.
                         -- See Note [Explicitly binding kind variables]
                         -- in Data.Singletons.Promote.Monad
                       }
    type ALetDecEnv = LetDecEnv Annotated
    type ULetDecEnv = LetDecEnv Unannotated

    instance Semigroup ULetDecEnv where
      LetDecEnv defns1 :: OMap Name (LetDecRHS Unannotated)
    defns1 types1 :: OMap Name DType
    types1 infx1 :: OMap Name Fixity
    infx1 _ _ <> :: ULetDecEnv -> ULetDecEnv -> ULetDecEnv
    <> LetDecEnv defns2 :: OMap Name (LetDecRHS Unannotated)
    defns2 types2 :: OMap Name DType
    types2 infx2 :: OMap Name Fixity
    infx2 _ _ =
        OMap Name (LetDecRHS Unannotated)
    -> OMap Name DType
    -> OMap Name Fixity
    -> IfAnn Unannotated (OMap Name DType) ()
    -> IfAnn Unannotated (OMap Name (OSet Name)) ()
    -> ULetDecEnv
    forall (ann :: AnnotationFlag).
    OMap Name (LetDecRHS ann)
    -> OMap Name DType
    -> OMap Name Fixity
    -> IfAnn ann (OMap Name DType) ()
    -> IfAnn ann (OMap Name (OSet Name)) ()
    -> LetDecEnv ann
    LetDecEnv (OMap Name (LetDecRHS Unannotated)
    defns1 OMap Name (LetDecRHS Unannotated)
    -> OMap Name (LetDecRHS Unannotated)
    -> OMap Name (LetDecRHS Unannotated)
    forall a. Semigroup a => a -> a -> a
    <> OMap Name (LetDecRHS Unannotated)
    defns2) (OMap Name DType
    types1 OMap Name DType -> OMap Name DType -> OMap Name DType
    forall a. Semigroup a => a -> a -> a
    <> OMap Name DType
    types2) (OMap Name Fixity
    infx1 OMap Name Fixity -> OMap Name Fixity -> OMap Name Fixity
    forall a. Semigroup a => a -> a -> a
    <> OMap Name Fixity
    infx2) () ()

    instance Monoid ULetDecEnv where
      mempty :: ULetDecEnv
    mempty = OMap Name (LetDecRHS Unannotated)
    -> OMap Name DType
    -> OMap Name Fixity
    -> IfAnn Unannotated (OMap Name DType) ()
    -> IfAnn Unannotated (OMap Name (OSet Name)) ()
    -> ULetDecEnv
    forall (ann :: AnnotationFlag).
    OMap Name (LetDecRHS ann)
    -> OMap Name DType
    -> OMap Name Fixity
    -> IfAnn ann (OMap Name DType) ()
    -> IfAnn ann (OMap Name (OSet Name)) ()
    -> LetDecEnv ann
    LetDecEnv OMap Name (LetDecRHS Unannotated)
    forall k v. OMap k v
    OMap.empty OMap Name DType
    forall k v. OMap k v
    OMap.empty OMap Name Fixity
    forall k v. OMap k v
    OMap.empty () ()

    valueBinding :: Name -> ULetDecRHS -> ULetDecEnv
    valueBinding :: Name -> LetDecRHS Unannotated -> ULetDecEnv
    valueBinding n :: Name
    n v :: LetDecRHS Unannotated
    v = ULetDecEnv
    emptyLetDecEnv { lde_defns :: OMap Name (LetDecRHS Unannotated)
    lde_defns = Name -> LetDecRHS Unannotated -> OMap Name (LetDecRHS Unannotated)
    forall k v. k -> v -> OMap k v
    OMap.singleton Name
    n LetDecRHS Unannotated
    v }

    typeBinding :: Name -> DType -> ULetDecEnv
    typeBinding :: Name -> DType -> ULetDecEnv
    typeBinding n :: Name
    n t :: DType
    t = ULetDecEnv
    emptyLetDecEnv { lde_types :: OMap Name DType
    lde_types = Name -> DType -> OMap Name DType
    forall k v. k -> v -> OMap k v
    OMap.singleton Name
    n DType
    t }

    infixDecl :: Fixity -> Name -> ULetDecEnv
    infixDecl :: Fixity -> Name -> ULetDecEnv
    infixDecl f :: Fixity
    f n :: Name
    n = ULetDecEnv
    emptyLetDecEnv { lde_infix :: OMap Name Fixity
    lde_infix = Name -> Fixity -> OMap Name Fixity
    forall k v. k -> v -> OMap k v
    OMap.singleton Name
    n Fixity
    f }

    emptyLetDecEnv :: ULetDecEnv
    emptyLetDecEnv :: ULetDecEnv
    emptyLetDecEnv = ULetDecEnv
    forall a. Monoid a => a
    mempty

    buildLetDecEnv :: Quasi q => [DLetDec] -> q ULetDecEnv
    buildLetDecEnv :: [DLetDec] -> q ULetDecEnv
    buildLetDecEnv = ULetDecEnv -> [DLetDec] -> q ULetDecEnv
    forall (m :: * -> *).
    Quasi m =>
    ULetDecEnv -> [DLetDec] -> m ULetDecEnv
    go ULetDecEnv
    emptyLetDecEnv
      where
        go :: ULetDecEnv -> [DLetDec] -> m ULetDecEnv
    go acc :: ULetDecEnv
    acc [] = ULetDecEnv -> m ULetDecEnv
    forall (m :: * -> *) a. Monad m => a -> m a
    return ULetDecEnv
    acc
        go acc :: ULetDecEnv
    acc (DFunD name :: Name
    name clauses :: [DClause]
    clauses : rest :: [DLetDec]
    rest) =
          ULetDecEnv -> [DLetDec] -> m ULetDecEnv
    go (Name -> LetDecRHS Unannotated -> ULetDecEnv
    valueBinding Name
    name ([DClause] -> LetDecRHS Unannotated
    UFunction [DClause]
    clauses) ULetDecEnv -> ULetDecEnv -> ULetDecEnv
    forall a. Semigroup a => a -> a -> a
    <> ULetDecEnv
    acc) [DLetDec]
    rest
        go acc :: ULetDecEnv
    acc (DValD (DVarP name :: Name
    name) exp :: DExp
    exp : rest :: [DLetDec]
    rest) =
          ULetDecEnv -> [DLetDec] -> m ULetDecEnv
    go (Name -> LetDecRHS Unannotated -> ULetDecEnv
    valueBinding Name
    name (DExp -> LetDecRHS Unannotated
    UValue DExp
    exp) ULetDecEnv -> ULetDecEnv -> ULetDecEnv
    forall a. Semigroup a => a -> a -> a
    <> ULetDecEnv
    acc) [DLetDec]
    rest
        go acc :: ULetDecEnv
    acc (dec :: DLetDec
    dec@(DValD {}) : rest :: [DLetDec]
    rest) = do
          [DLetDec]
    flattened <- DLetDec -> m [DLetDec]
    forall (q :: * -> *). Quasi q => DLetDec -> q [DLetDec]
    flattenDValD DLetDec
    dec
          ULetDecEnv -> [DLetDec] -> m ULetDecEnv
    go ULetDecEnv
    acc ([DLetDec]
    flattened [DLetDec] -> [DLetDec] -> [DLetDec]
    forall a. [a] -> [a] -> [a]
    ++ [DLetDec]
    rest)
        go acc :: ULetDecEnv
    acc (DSigD name :: Name
    name ty :: DType
    ty : rest :: [DLetDec]
    rest) =
          ULetDecEnv -> [DLetDec] -> m ULetDecEnv
    go (Name -> DType -> ULetDecEnv
    typeBinding Name
    name DType
    ty ULetDecEnv -> ULetDecEnv -> ULetDecEnv
    forall a. Semigroup a => a -> a -> a
    <> ULetDecEnv
    acc) [DLetDec]
    rest
        go acc :: ULetDecEnv
    acc (DInfixD f :: Fixity
    f n :: Name
    n : rest :: [DLetDec]
    rest) =
          ULetDecEnv -> [DLetDec] -> m ULetDecEnv
    go (Fixity -> Name -> ULetDecEnv
    infixDecl Fixity
    f Name
    n ULetDecEnv -> ULetDecEnv -> ULetDecEnv
    forall a. Semigroup a => a -> a -> a
    <> ULetDecEnv
    acc) [DLetDec]
    rest
        go acc :: ULetDecEnv
    acc (DPragmaD{} : rest :: [DLetDec]
    rest) = ULetDecEnv -> [DLetDec] -> m ULetDecEnv
    go ULetDecEnv
    acc [DLetDec]
    rest

    -- See Note [DerivedDecl]
    data DerivedDecl (cls :: Type -> Constraint) = DerivedDecl
      { DerivedDecl cls -> Maybe DCxt
    ded_mb_cxt     :: Maybe DCxt
      , DerivedDecl cls -> DType
    ded_type       :: DType
      , DerivedDecl cls -> Name
    ded_type_tycon :: Name
      , DerivedDecl cls -> DataDecl
    ded_decl       :: DataDecl
      }

    type DerivedEqDecl   = DerivedDecl Eq
    type DerivedShowDecl = DerivedDecl Show

    {- Note [DerivedDecl]
    ~~~~~~~~~~~~~~~~~~~~~
    Most derived instances are wholly handled in
    Data.Singletons.Partition.partitionDecs. There are two notable exceptions to
    this rule, however:

    * Eq instances (which are handled entirely outside of partitionDecs)
    * Show instances (which are partially handled outside of partitionDecs)

    For these instances, we use a DerivedDecl data type to encode just enough
    information to recreate the derived instance:

    1. Just the instance context, if it's standalone-derived, or Nothing if it's in
       a deriving clause (ded_mb_cxt)
    2. The datatype, applied to some number of type arguments, as in the
       instance declaration (ded_type)
    3. The datatype name (ded_type_tycon), cached for convenience
    4. The datatype's constructors (ded_cons)

    Why are these instances handled outside of partitionDecs?

    * Deriving Eq in singletons not only derives PEq/SEq instances, but it also
      derives SDecide, TestEquality, and TestCoercion instances. This additional
      complication makes Eq difficult to integrate with the other deriving
      machinery, so we handle it specially in Data.Singletons.Promote and
      Data.Singletons.Single (depending on the task at hand).
    * Deriving Show in singletons not only derives PShow/SShow instances, but it
      also derives Show instances for singletons types. To make this work,
      we let partitionDecs handle the PShow/SShow instances, but we also stick the
      relevant info into a DerivedDecl value for later use in
      Data.Singletons.Single, where we additionally generate Show
      instances.
    -}