{-# LANGUAGE PatternSynonyms #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Disco.AST.Desugared
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Typed abstract syntax trees representing the typechecked, desugared
-- Disco language.
--
-----------------------------------------------------------------------------

module Disco.AST.Desugared
       ( -- * Desugared, type-annotated terms
       DTerm
       , pattern DTVar
       , pattern DTPrim
       , pattern DTUnit
       , pattern DTBool
       , pattern DTChar
       , pattern DTNat
       , pattern DTRat
       , pattern DTAbs
       , pattern DTApp
       , pattern DTPair
       , pattern DTCase
       , pattern DTTyOp
       , pattern DTNil
       , pattern DTTest

       , Container(..)
       , DBinding
       , pattern DBinding
         -- * Branches and guards
       , DBranch

       , DGuard
       , pattern DGPat

       , DPattern
       , pattern DPVar
       , pattern DPWild
       , pattern DPUnit
       , pattern DPPair
       , pattern DPInj

       , DProperty
       )
       where

import           GHC.Generics

import           Data.Void
import           Unbound.Generics.LocallyNameless

import           Disco.AST.Generic
import           Disco.Names                      (QName (..))
import           Disco.Syntax.Operators
import           Disco.Syntax.Prims
import           Disco.Types

data DS

type DProperty = Property_ DS

-- | A @DTerm@ is a term which has been typechecked and desugared, so
--   it has fewer constructors and complex features than 'ATerm', but
--   still retains typing information.

type DTerm = Term_ DS

type instance X_Binder DS         = Name DTerm

type instance X_TVar DS           = Void -- names are qualified
type instance X_TPrim DS          = Type
type instance X_TLet DS           = Void -- Let gets translated to lambda
type instance X_TUnit DS          = ()
type instance X_TBool DS          = Type
type instance X_TChar DS          = ()
type instance X_TString DS        = Void
type instance X_TNat DS           = Type
type instance X_TRat DS           = ()
type instance X_TAbs DS           = Type -- For lambas this is the function type but
                                         -- for forall/exists it's the argument type
type instance X_TApp DS           = Type
type instance X_TCase DS          = Type
type instance X_TChain DS         = Void -- Chains are translated into conjunctions of
                                         -- binary comparisons
type instance X_TTyOp DS          = Type
type instance X_TContainer DS     = Void -- Literal containers are desugared into
                                         -- conversion functions applied to list literals

type instance X_TContainerComp DS = Void -- Container comprehensions are translated
                                         -- into monadic chains

type instance X_TAscr DS          = Void -- No type ascriptions
type instance X_TTup DS           = Void -- No tuples, only pairs
type instance X_TParens DS        = Void -- No explicit parens

-- Extra constructors
type instance X_Term DS = X_DTerm

data X_DTerm
  = DTPair_ Type DTerm DTerm
  | DTNil_ Type
  | DTTest_ [(String, Type, Name DTerm)] DTerm
  | DTVar_ Type (QName DTerm)
  deriving (Show, Generic)

instance Subst Type X_DTerm
instance Alpha X_DTerm

pattern DTVar :: Type -> QName DTerm -> DTerm
pattern DTVar ty qname = XTerm_ (DTVar_ ty qname)

pattern DTPrim :: Type -> Prim -> DTerm
pattern DTPrim ty name = TPrim_ ty name

pattern DTUnit :: DTerm
pattern DTUnit = TUnit_ ()

pattern DTBool :: Type -> Bool -> DTerm
pattern DTBool ty bool = TBool_ ty bool

pattern DTNat  :: Type -> Integer -> DTerm
pattern DTNat ty int = TNat_ ty int

pattern DTRat :: Rational -> DTerm
pattern DTRat rat = TRat_ () rat

pattern DTChar :: Char -> DTerm
pattern DTChar c = TChar_ () c

pattern DTAbs :: Quantifier -> Type -> Bind (Name DTerm) DTerm -> DTerm
pattern DTAbs q ty lam = TAbs_ q ty lam

pattern DTApp  :: Type -> DTerm -> DTerm -> DTerm
pattern DTApp ty term1 term2 = TApp_ ty term1 term2

pattern DTPair :: Type -> DTerm -> DTerm -> DTerm
pattern DTPair ty t1 t2 = XTerm_ (DTPair_ ty t1 t2)

pattern DTCase :: Type -> [DBranch] -> DTerm
pattern DTCase ty branch = TCase_ ty branch

pattern DTTyOp :: Type -> TyOp -> Type -> DTerm
pattern DTTyOp ty1 tyop ty2 = TTyOp_ ty1 tyop ty2

pattern DTNil :: Type -> DTerm
pattern DTNil ty = XTerm_ (DTNil_ ty)

-- | A test frame, recording a collection of variables with their types and
--   their original user-facing names. Used for legible reporting of test
--   failures inside the enclosed term.
pattern DTTest :: [(String, Type, Name DTerm)] -> DTerm -> DTerm
pattern DTTest ns t = XTerm_ (DTTest_ ns t)

{-# COMPLETE DTVar, DTPrim, DTUnit, DTBool, DTChar, DTNat, DTRat,
             DTAbs, DTApp, DTPair, DTCase, DTTyOp,
             DTNil, DTTest #-}

type instance X_TLink DS = Void

type DBinding = Binding_ DS

pattern DBinding :: Maybe (Embed PolyType) -> Name DTerm -> Embed DTerm -> DBinding
pattern DBinding m b n = Binding_ m b n

{-# COMPLETE DBinding #-}

type DBranch = Bind (Telescope DGuard) DTerm

type DGuard = Guard_ DS

type instance X_GBool DS = Void   -- Boolean guards get desugared to pattern-matching
type instance X_GPat  DS = ()
type instance X_GLet  DS = Void   -- Let gets desugared to 'when' with a variable

pattern DGPat :: Embed DTerm -> DPattern -> DGuard
pattern DGPat embedt pat = GPat_ () embedt pat

{-# COMPLETE DGPat #-}

type DPattern = Pattern_ DS

type instance X_PVar     DS = Embed Type
type instance X_PWild    DS = Embed Type
type instance X_PAscr    DS = Void
type instance X_PUnit    DS = ()
type instance X_PBool    DS = Void
type instance X_PChar    DS = Void
type instance X_PString  DS = Void
type instance X_PTup     DS = Void
type instance X_PInj     DS = Void
type instance X_PNat     DS = Void
type instance X_PCons    DS = Void
type instance X_PList    DS = Void
type instance X_PAdd     DS = Void
type instance X_PMul     DS = Void
type instance X_PSub     DS = Void
type instance X_PNeg     DS = Void
type instance X_PFrac    DS = Void

-- In the desugared language, constructor patterns (DPPair, DPInj) can
-- only contain variables, not nested patterns.  This means that the
-- desugaring phase has to make explicit the order of matching by
-- exploding nested patterns into sequential guards, which makes the
-- interpreter simpler.

type instance X_Pattern  DS =
  Either
    (Embed Type, Name DTerm, Name DTerm)     -- DPPair
    (Embed Type, Side, Name DTerm)           -- DPInj

pattern DPVar :: Type -> Name DTerm -> DPattern
pattern DPVar ty name <- PVar_ (unembed -> ty) name
  where
    DPVar ty name = PVar_ (embed ty) name

pattern DPWild :: Type -> DPattern
pattern DPWild ty <- PWild_ (unembed -> ty)
  where
    DPWild ty = PWild_ (embed ty)

pattern DPUnit :: DPattern
pattern DPUnit = PUnit_ ()

pattern DPPair  :: Type -> Name DTerm -> Name DTerm -> DPattern
pattern DPPair ty x1 x2 <- XPattern_ (Left (unembed -> ty, x1, x2))
  where
    DPPair ty x1 x2 = XPattern_ (Left (embed ty, x1, x2))

pattern DPInj  :: Type -> Side -> Name DTerm -> DPattern
pattern DPInj ty s x <- XPattern_ (Right (unembed -> ty, s, x))
  where
    DPInj ty s x = XPattern_ (Right (embed ty, s, x))

{-# COMPLETE DPVar, DPWild, DPUnit, DPPair, DPInj #-}

type instance X_QBind  DS = Void
type instance X_QGuard DS = Void

------------------------------------------------------------
-- getType
------------------------------------------------------------

instance HasType DTerm where
  getType (DTVar ty _)     = ty
  getType (DTPrim ty _)    = ty
  getType DTUnit           = TyUnit
  getType (DTBool ty _)    = ty
  getType (DTChar _)       = TyC
  getType (DTNat ty _)     = ty
  getType (DTRat _)        = TyF
  getType (DTAbs Lam ty _) = ty
  getType DTAbs{}          = TyProp
  getType (DTApp ty _ _)   = ty
  getType (DTPair ty _ _)  = ty
  getType (DTCase ty _)    = ty
  getType (DTTyOp ty _ _)  = ty
  getType (DTNil ty)       = ty
  getType (DTTest _ _)     = TyProp

instance HasType DPattern where
  getType (DPVar ty _)    = ty
  getType (DPWild ty)     = ty
  getType DPUnit          = TyUnit
  getType (DPPair ty _ _) = ty
  getType (DPInj ty _ _)  = ty