{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE PatternSynonyms    #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Disco.AST.Typed
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- SPDX-License-Identifier: BSD-3-Clause
--
-- Typed abstract syntax trees representing the typechecked surface
-- syntax of the Disco language.  Each tree node is annotated with the
-- type of its subtree.
--
-----------------------------------------------------------------------------

module Disco.AST.Typed
       ( -- * Type-annotated terms
         ATerm
       , pattern ATVar
       , pattern ATPrim
       , pattern ATLet
       , pattern ATUnit
       , pattern ATBool
       , pattern ATNat
       , pattern ATRat
       , pattern ATChar
       , pattern ATString
       , pattern ATAbs
       , pattern ATApp
       , pattern ATTup
       , pattern ATCase
       , pattern ATChain
       , pattern ATTyOp
       , pattern ATContainer
       , pattern ATContainerComp
       , pattern ATList
       , pattern ATListComp
       , pattern ATTest

       , ALink
       , pattern ATLink

       , Container(..)
       , ABinding
         -- * Branches and guards
       , ABranch

       , AGuard
       , pattern AGBool
       , pattern AGPat
       , pattern AGLet

       , AQual
       , pattern AQBind
       , pattern AQGuard

       , APattern
       , pattern APVar
       , pattern APWild
       , pattern APUnit
       , pattern APBool
       , pattern APTup
       , pattern APInj
       , pattern APNat
       , pattern APChar
       , pattern APString
       , pattern APCons
       , pattern APList
       , pattern APAdd
       , pattern APMul
       , pattern APSub
       , pattern APNeg
       , pattern APFrac

       , pattern ABinding
         -- * Utilities
       , varsBound
       , getType
       , setType
       , substQT

       , AProperty
       )
       where

import           Unbound.Generics.LocallyNameless
import           Unbound.Generics.LocallyNameless.Unsafe

import           Control.Arrow                           ((***))
import           Data.Coerce                             (coerce)
import           Data.Data                               (Data)
import           Data.Void

import           Control.Lens.Plated                     (transform)
import           Disco.AST.Generic
import           Disco.AST.Surface
import           Disco.Names
import           Disco.Pretty
import           Disco.Syntax.Operators
import           Disco.Syntax.Prims
import           Disco.Types

-- | The extension descriptor for Typed specific AST types.

data TY
  deriving Data

type AProperty = Property_ TY

------------------------------------------------------------

-- TODO: Should probably really do this with a 2-level/open recursion
-- approach, with a cofree comonad or whatever

-- | An @ATerm@ is a typechecked term where every node in the tree has
--   been annotated with the type of the subterm rooted at that node.

type ATerm = Term_ TY

type instance X_Binder          TY = [APattern]

type instance X_TVar            TY = Void -- Names are now qualified
type instance X_TPrim           TY = Type
type instance X_TLet            TY = Type
type instance X_TUnit           TY = ()
type instance X_TBool           TY = Type
type instance X_TNat            TY = Type
type instance X_TRat            TY = ()
type instance X_TChar           TY = ()
type instance X_TString         TY = ()
type instance X_TAbs            TY = Type
type instance X_TApp            TY = Type
type instance X_TCase           TY = Type
type instance X_TChain          TY = Type
type instance X_TTyOp           TY = Type
type instance X_TContainer      TY = Type
type instance X_TContainerComp  TY = Type
type instance X_TAscr           TY = Void -- No more type ascriptions in typechecked terms
type instance X_TTup            TY = Type
type instance X_TParens         TY = Void -- No more explicit parens

 -- A test frame for reporting counterexamples in a test. These don't appear
 -- in source programs, but because the deugarer manipulates partly-desugared
 -- terms it helps to be able to represent these in 'ATerm'.
type instance X_Term TY = Either ([(String, Type, Name ATerm)], ATerm) (Type, QName ATerm)

pattern ATVar :: Type -> QName ATerm -> ATerm
pattern ATVar ty qname = XTerm_ (Right (ty, qname))

pattern ATPrim :: Type -> Prim -> ATerm
pattern ATPrim ty name = TPrim_ ty name

pattern ATLet :: Type -> Bind (Telescope ABinding) ATerm -> ATerm
pattern ATLet ty bind = TLet_ ty bind

pattern ATUnit :: ATerm
pattern ATUnit = TUnit_ ()

pattern ATBool :: Type -> Bool -> ATerm
pattern ATBool ty bool = TBool_ ty bool

pattern ATNat  :: Type -> Integer -> ATerm
pattern ATNat ty int = TNat_ ty int

pattern ATRat :: Rational -> ATerm
pattern ATRat rat = TRat_ () rat

pattern ATChar :: Char -> ATerm
pattern ATChar c = TChar_ () c

pattern ATString :: String -> ATerm
pattern ATString s = TString_ () s

pattern ATAbs :: Quantifier -> Type -> Bind [APattern] ATerm -> ATerm
pattern ATAbs q ty bind = TAbs_ q ty bind

pattern ATApp  :: Type -> ATerm -> ATerm -> ATerm
pattern ATApp ty term1 term2 = TApp_ ty term1 term2

pattern ATTup :: Type -> [ATerm] -> ATerm
pattern ATTup ty termlist = TTup_ ty termlist

pattern ATCase :: Type -> [ABranch] -> ATerm
pattern ATCase ty branch = TCase_ ty branch

pattern ATChain :: Type -> ATerm -> [ALink] -> ATerm
pattern ATChain ty term linklist = TChain_ ty term linklist

pattern ATTyOp :: Type -> TyOp -> Type -> ATerm
pattern ATTyOp ty1 tyop ty2 = TTyOp_ ty1 tyop ty2

pattern ATContainer :: Type -> Container -> [(ATerm, Maybe ATerm)] -> Maybe (Ellipsis ATerm) -> ATerm
pattern ATContainer ty c tl mets = TContainer_ ty c tl mets

pattern ATContainerComp :: Type -> Container -> Bind (Telescope AQual) ATerm -> ATerm
pattern ATContainerComp ty c b = TContainerComp_ ty c b

pattern ATTest :: [(String, Type, Name ATerm)] -> ATerm -> ATerm
pattern ATTest ns t = XTerm_ (Left (ns, t))

{-# COMPLETE ATVar, ATPrim, ATLet, ATUnit, ATBool, ATNat, ATRat, ATChar,
             ATString, ATAbs, ATApp, ATTup, ATCase, ATChain, ATTyOp,
             ATContainer, ATContainerComp, ATTest #-}

pattern ATList :: Type -> [ATerm] -> Maybe (Ellipsis ATerm) -> ATerm
pattern ATList t xs e <- ATContainer t ListContainer (map fst -> xs) e
  where
    ATList t xs e = ATContainer t ListContainer (map (,Nothing) xs) e

pattern ATListComp :: Type -> Bind (Telescope AQual) ATerm -> ATerm
pattern ATListComp t b = ATContainerComp t ListContainer b

type ALink = Link_ TY

type instance X_TLink TY = ()

pattern ATLink :: BOp -> ATerm -> ALink
pattern ATLink bop term = TLink_ () bop term

{-# COMPLETE ATLink #-}


type AQual = Qual_ TY

type instance X_QBind TY = ()
type instance X_QGuard TY = ()


pattern AQBind :: Name ATerm -> Embed ATerm -> AQual
pattern AQBind namet embedt = QBind_ () namet embedt

pattern AQGuard :: Embed ATerm -> AQual
pattern AQGuard embedt = QGuard_ () embedt

{-# COMPLETE AQBind, AQGuard #-}

type ABinding = Binding_ TY

pattern ABinding :: Maybe (Embed PolyType) -> Name ATerm -> Embed ATerm -> ABinding
pattern ABinding m b n = Binding_ m b n

{-# COMPLETE ABinding #-}

type ABranch = Bind (Telescope AGuard) ATerm

type AGuard = Guard_ TY

type instance X_GBool TY = ()
type instance X_GPat  TY = ()
type instance X_GLet  TY = ()   -- ??? Type?

pattern AGBool :: Embed ATerm -> AGuard
pattern AGBool embedt = GBool_ () embedt

pattern AGPat :: Embed ATerm -> APattern -> AGuard
pattern AGPat embedt pat = GPat_ () embedt pat

pattern AGLet :: ABinding -> AGuard
pattern AGLet b = GLet_ () b

{-# COMPLETE AGBool, AGPat, AGLet #-}

type APattern = Pattern_ TY

-- We have to use Embed Type because we don't want any type variables
-- inside the types being treated as binders!

type instance X_PVar     TY = Embed Type
type instance X_PWild    TY = Embed Type
type instance X_PAscr    TY = Void -- No more ascriptions in typechecked patterns.
type instance X_PUnit    TY = ()
type instance X_PBool    TY = ()
type instance X_PChar    TY = ()
type instance X_PString  TY = ()
type instance X_PTup     TY = Embed Type
type instance X_PInj     TY = Embed Type
type instance X_PNat     TY = Embed Type
type instance X_PCons    TY = Embed Type
type instance X_PList    TY = Embed Type
type instance X_PAdd     TY = Embed Type
type instance X_PMul     TY = Embed Type
type instance X_PSub     TY = Embed Type
type instance X_PNeg     TY = Embed Type
type instance X_PFrac    TY = Embed Type

type instance X_Pattern  TY = ()

pattern APVar :: Type -> Name ATerm -> APattern
pattern APVar ty name <- PVar_ (unembed -> ty) name
  where
    APVar ty name = PVar_ (embed ty) name

pattern APWild :: Type -> APattern
pattern APWild ty <- PWild_ (unembed -> ty)
  where
    APWild ty = PWild_ (embed ty)

pattern APUnit :: APattern
pattern APUnit = PUnit_ ()

pattern APBool :: Bool -> APattern
pattern APBool  b = PBool_ () b

pattern APChar :: Char -> APattern
pattern APChar  c = PChar_ () c

pattern APString :: String -> APattern
pattern APString s = PString_ () s

pattern APTup  :: Type -> [APattern] -> APattern
pattern APTup ty lp <- PTup_ (unembed -> ty) lp
  where
    APTup ty lp = PTup_ (embed ty) lp

pattern APInj  :: Type -> Side -> APattern -> APattern
pattern APInj ty s p <- PInj_ (unembed -> ty) s p
  where
    APInj ty s p = PInj_ (embed ty) s p

pattern APNat  :: Type -> Integer -> APattern
pattern APNat ty n <- PNat_ (unembed -> ty) n
  where
    APNat ty n = PNat_ (embed ty) n

pattern APCons :: Type -> APattern -> APattern -> APattern
pattern APCons ty p1 p2 <- PCons_ (unembed -> ty) p1 p2
  where
    APCons ty p1 p2 = PCons_ (embed ty) p1 p2

pattern APList :: Type -> [APattern] -> APattern
pattern APList ty lp <- PList_ (unembed -> ty) lp
  where
    APList ty lp = PList_ (embed ty) lp

pattern APAdd :: Type -> Side -> APattern -> ATerm -> APattern
pattern APAdd ty s p t <- PAdd_ (unembed -> ty) s p t
  where
    APAdd ty s p t = PAdd_ (embed ty) s p t

pattern APMul :: Type -> Side -> APattern -> ATerm -> APattern
pattern APMul ty s p t <- PMul_ (unembed -> ty) s p t
  where
    APMul ty s p t = PMul_ (embed ty) s p t

pattern APSub :: Type -> APattern -> ATerm -> APattern
pattern APSub ty p t <- PSub_ (unembed -> ty) p t
  where
    APSub ty p t = PSub_ (embed ty) p t

pattern APNeg :: Type -> APattern -> APattern
pattern APNeg ty p <- PNeg_ (unembed -> ty) p
  where
    APNeg ty p = PNeg_ (embed ty) p

pattern APFrac :: Type -> APattern -> APattern -> APattern
pattern APFrac ty p1 p2 <- PFrac_ (unembed -> ty) p1 p2
  where
    APFrac ty p1 p2 = PFrac_ (embed ty) p1 p2

{-# COMPLETE APVar, APWild, APUnit, APBool, APChar, APString,
    APTup, APInj, APNat, APCons, APList, APAdd, APMul, APSub, APNeg, APFrac #-}

varsBound :: APattern -> [(Name ATerm, Type)]
varsBound (APVar ty n)    = [(n, ty)]
varsBound (APWild _)      = []
varsBound APUnit          = []
varsBound (APBool _)      = []
varsBound (APChar _)      = []
varsBound (APString _)    = []
varsBound (APTup _ ps)    = varsBound =<< ps
varsBound (APInj _ _ p)   = varsBound p
varsBound (APNat _ _)     = []
varsBound (APCons _ p q)  = varsBound p ++ varsBound q
varsBound (APList _ ps)   = varsBound =<< ps
varsBound (APAdd _ _ p _) = varsBound p
varsBound (APMul _ _ p _) = varsBound p
varsBound (APSub _ p _)   = varsBound p
varsBound (APNeg _ p)     = varsBound p
varsBound (APFrac _ p q)  = varsBound p ++ varsBound q

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

instance HasType ATerm where
  getType (ATVar ty _)             = ty
  getType (ATPrim ty _)            = ty
  getType ATUnit                   = TyUnit
  getType (ATBool ty _)            = ty
  getType (ATNat ty _)             = ty
  getType (ATRat _)                = TyF
  getType (ATChar _)               = TyC
  getType (ATString _)             = TyList TyC
  getType (ATAbs _ ty _)           = ty
  getType (ATApp ty _ _)           = ty
  getType (ATTup ty _)             = ty
  getType (ATTyOp ty _ _)          = ty
  getType (ATChain ty _ _)         = ty
  getType (ATContainer ty _ _ _)   = ty
  getType (ATContainerComp ty _ _) = ty
  getType (ATLet ty _)             = ty
  getType (ATCase ty _)            = ty
  getType (ATTest _ _ )            = TyProp

  setType ty (ATVar _ x      )       = ATVar ty x
  setType ty (ATPrim _ x     )       = ATPrim ty x
  setType _  ATUnit                  = ATUnit
  setType ty (ATBool _ b)            = ATBool ty b
  setType ty (ATNat _ x      )       = ATNat ty x
  setType _  (ATRat r)               = ATRat r
  setType _ (ATChar c)               = ATChar c
  setType _ (ATString cs)            = ATString cs
  setType ty (ATAbs q _ x    )       = ATAbs q ty x
  setType ty (ATApp _ x y    )       = ATApp ty x y
  setType ty (ATTup _ x      )       = ATTup ty x
  setType ty (ATTyOp _ x y   )       = ATTyOp ty x y
  setType ty (ATChain _ x y  )       = ATChain ty x y
  setType ty (ATContainer _ x y z)   = ATContainer ty x y z
  setType ty (ATContainerComp _ x y) = ATContainerComp ty x y
  setType ty (ATLet _ x      )       = ATLet ty x
  setType ty (ATCase _ x     )       = ATCase ty x
  setType _ (ATTest vs x)            = ATTest vs x

instance HasType APattern where
  getType (APVar ty _)     = ty
  getType (APWild ty)      = ty
  getType APUnit           = TyUnit
  getType (APBool _)       = TyBool
  getType (APChar _)       = TyC
  getType (APString _)     = TyList TyC
  getType (APTup ty _)     = ty
  getType (APInj ty _ _)   = ty
  getType (APNat ty _)     = ty
  getType (APCons ty _ _)  = ty
  getType (APList ty _)    = ty
  getType (APAdd ty _ _ _) = ty
  getType (APMul ty _ _ _) = ty
  getType (APSub ty _ _)   = ty
  getType (APNeg ty _)     = ty
  getType (APFrac ty _ _)  = ty

instance HasType ABranch where
  getType = getType . snd . unsafeUnbind

------------------------------------------------------------
-- subst
------------------------------------------------------------

substQT :: QName ATerm -> ATerm -> ATerm -> ATerm
substQT x s = transform $ \case
  t@(ATVar _ y)
    | x == y    -> s
    | otherwise -> t
  t -> t

------------------------------------------------------------
-- Exploding a typed term into a surface term with annotations
------------------------------------------------------------

instance Pretty ATerm where
  pretty = pretty . explode

explode :: ATerm -> Term
explode = \case
  ATVar ty x             -> TAscr (TVar (coerce (qname x))) (toPolyType ty)
  ATPrim ty x            -> TAscr (TPrim x) (toPolyType ty)
  ATLet ty tel           -> TAscr (TLet (explodeTelescope explodeBinding tel)) (toPolyType ty)
  ATUnit                 -> TUnit
  ATBool _ty b           -> TBool b
  ATNat ty x             -> TAscr (TNat x) (toPolyType ty)
  ATRat r                -> TRat r
  ATChar c               -> TChar c
  ATString cs            -> TString cs
  ATAbs q ty a           -> TAscr (TAbs q (explodeAbs a)) (toPolyType ty)
  ATApp ty x y           -> TAscr (TApp (explode x) (explode y)) (toPolyType ty)
  ATTup ty xs            -> TAscr (TTup (map explode xs)) (toPolyType ty)
  ATCase ty bs           -> TAscr (TCase (map explodeBranch bs)) (toPolyType ty)
  ATChain ty t ls        -> TAscr (TChain (explode t) (map explodeLink ls)) (toPolyType ty)
  ATTyOp ty x y          -> TAscr (TTyOp x y) (toPolyType ty)
  ATContainer ty c ts el ->
    TAscr
      (TContainer c (map (explode *** fmap explode) ts) (fmap (fmap explode) el))
      (toPolyType ty)
  ATContainerComp ty c b -> TAscr (TContainerComp c (explodeTelescope explodeQual b)) (toPolyType ty)
  ATTest _vs x           -> TAscr (explode x) (toPolyType TyProp)

explodeTelescope
  :: (Alpha a, Alpha b)
  => (a -> b) -> Bind (Telescope a) ATerm -> Bind (Telescope b) Term
explodeTelescope explodeBinder (unsafeUnbind -> (xs,at)) = bind (mapTelescope explodeBinder xs) (explode at)

explodeBinding :: ABinding -> Binding
explodeBinding (ABinding m b (unembed -> n)) = Binding m (coerce b) (embed (explode n))

explodeAbs :: Bind [APattern] ATerm -> Bind [Pattern] Term
explodeAbs (unsafeUnbind -> (aps, at)) = bind (map explodePattern aps) (explode at)

explodePattern :: APattern -> Pattern
explodePattern = \case
  APVar ty x      -> PAscr (PVar (coerce x)) ty
  APWild ty       -> PAscr PWild ty
  APUnit          -> PUnit
  APBool b        -> PBool b
  APChar c        -> PChar c
  APString s      -> PString s
  APTup ty ps     -> PAscr (PTup (map explodePattern ps)) ty
  APInj ty s p    -> PAscr (PInj s (explodePattern p)) ty
  APNat ty n      -> PAscr (PNat n) ty
  APCons ty p1 p2 -> PAscr (PCons (explodePattern p1) (explodePattern p2)) ty
  APList ty ps    -> PAscr (PList (map explodePattern ps)) ty
  APAdd ty s p t  -> PAscr (PAdd s (explodePattern p) (explode t)) ty
  APMul ty s p t  -> PAscr (PMul s (explodePattern p) (explode t)) ty
  APSub ty p t    -> PAscr (PSub (explodePattern p) (explode t)) ty
  APNeg ty p      -> PAscr (PNeg (explodePattern p)) ty
  APFrac ty p q   -> PAscr (PFrac (explodePattern p) (explodePattern q)) ty

explodeBranch :: ABranch -> Branch
explodeBranch = explodeTelescope explodeGuard

explodeGuard :: AGuard -> Guard
explodeGuard (AGBool (unembed -> at))   = GBool (embed (explode at))
explodeGuard (AGPat (unembed -> at) ap) = GPat (embed (explode at)) (explodePattern ap)
explodeGuard (AGLet ab)                 = GLet (explodeBinding ab)

explodeLink :: ALink -> Link
explodeLink (ATLink bop at) = TLink bop (explode at)

explodeQual :: AQual -> Qual
explodeQual (AQBind x (unembed -> at)) = QBind (coerce x) (embed (explode at))
explodeQual (AQGuard (unembed -> at))  = QGuard (embed (explode at))