Typed.hs

{-# 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))