{- 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 VarPromotions
vars1 OSet Name
sig_kvs1 <> :: PromDPatInfos -> PromDPatInfos -> PromDPatInfos
<> PromDPatInfos VarPromotions
vars2 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
              , ClassDecl ann -> [OpenTypeFamilyDecl]
cd_atfs :: [OpenTypeFamilyDecl]
                  -- Associated type families. Only recorded for
                  -- defunctionalization purposes.
                  -- See Note [Partitioning, type synonyms, and type families]
                  -- in D.S.Partition.
              }

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 OMap Name (LetDecRHS Unannotated)
defns1 OMap Name DType
types1 OMap Name Fixity
infx1 IfAnn Unannotated (OMap Name DType) ()
_ IfAnn Unannotated (OMap Name (OSet Name)) ()
_ <> :: ULetDecEnv -> ULetDecEnv -> ULetDecEnv
<> LetDecEnv OMap Name (LetDecRHS Unannotated)
defns2 OMap Name DType
types2 OMap Name Fixity
infx2 IfAnn Unannotated (OMap Name DType) ()
_ IfAnn Unannotated (OMap Name (OSet Name)) ()
_ =
    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 Name
n 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 Name
n 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 Fixity
f 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 ULetDecEnv
acc [] = ULetDecEnv -> m ULetDecEnv
forall (m :: * -> *) a. Monad m => a -> m a
return ULetDecEnv
acc
    go ULetDecEnv
acc (DFunD Name
name [DClause]
clauses : [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 ULetDecEnv
acc (DValD (DVarP Name
name) DExp
exp : [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 ULetDecEnv
acc (dec :: DLetDec
dec@(DValD {}) : [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 ULetDecEnv
acc (DSigD Name
name DType
ty : [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 ULetDecEnv
acc (DInfixD Fixity
f Name
n : [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 ULetDecEnv
acc (DPragmaD{} : [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.
-}