-----------------------------------------------------------------------------
-- |
-- Module      :  Disco.Typecheck.Erase
-- Copyright   :  (c) 2016 disco team (see LICENSE)
-- License     :  BSD-style (see LICENSE)
-- Maintainer  :  byorgey@gmail.com
--
-- Typecheck the Disco surface language and transform it into a
-- type-annotated AST.
--
-----------------------------------------------------------------------------

module Disco.Typecheck.Erase where

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

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

import           Disco.AST.Desugared
import           Disco.AST.Surface
import           Disco.AST.Typed
import           Disco.Names                             (QName (..))

-- | Erase all the type annotations from a term.
erase :: ATerm -> Term
erase (ATVar _ (QName _ x)) = TVar (coerce x)
erase (ATPrim _ x)          = TPrim x
erase (ATLet _ bs)          = TLet $ bind (mapTelescope eraseBinding tel) (erase at)
  where (tel,at) = unsafeUnbind bs
erase ATUnit                = TUnit
erase (ATBool _ b)          = TBool b
erase (ATChar c)            = TChar c
erase (ATString s)          = TString s
erase (ATNat _ i)           = TNat i
erase (ATRat r)             = TRat r
erase (ATAbs q _ b)         = TAbs q $ bind (map erasePattern x) (erase at)
  where (x,at) = unsafeUnbind b
erase (ATApp _ t1 t2)       = TApp (erase t1) (erase t2)
erase (ATTup _ ats)         = TTup (map erase ats)
erase (ATCase _ brs)        = TCase (map eraseBranch brs)
erase (ATChain _ at lnks)   = TChain (erase at) (map eraseLink lnks)
erase (ATTyOp _ op ty)      = TTyOp op ty
erase (ATContainer _ c ats aell)   = TContainer c (map (erase *** fmap erase) ats) ((fmap . fmap) erase aell)
erase (ATContainerComp _ c b)      = TContainerComp c $ bind (mapTelescope eraseQual tel) (erase at)
  where (tel,at) = unsafeUnbind b
erase (ATTest _ x)          = erase x

eraseBinding :: ABinding -> Binding
eraseBinding (ABinding mty x (unembed -> at)) = Binding mty (coerce x) (embed (erase at))

erasePattern :: APattern -> Pattern
erasePattern (APVar _ n)        = PVar (coerce n)
erasePattern (APWild _)         = PWild
erasePattern APUnit             = PUnit
erasePattern (APBool b)         = PBool b
erasePattern (APChar c)         = PChar c
erasePattern (APString s)       = PString s
erasePattern (APTup _ alp)      = PTup $ map erasePattern alp
erasePattern (APInj _ s apt)    = PInj s (erasePattern apt)
erasePattern (APNat _ n)        = PNat n
erasePattern (APCons _ ap1 ap2) = PCons (erasePattern ap1) (erasePattern ap2)
erasePattern (APList _ alp)     = PList $ map erasePattern alp
erasePattern (APAdd _ s p t)    = PAdd s (erasePattern p) (erase t)
erasePattern (APMul _ s p t)    = PMul s (erasePattern p) (erase t)
erasePattern (APSub _ p t)      = PSub (erasePattern p) (erase t)
erasePattern (APNeg _ p)        = PNeg (erasePattern p)
erasePattern (APFrac _ p1 p2)   = PFrac (erasePattern p1) (erasePattern p2)

eraseBranch :: ABranch -> Branch
eraseBranch b = bind (mapTelescope eraseGuard tel) (erase at)
  where (tel,at) = unsafeUnbind b

eraseGuard :: AGuard -> Guard
eraseGuard (AGBool (unembed -> at))  = GBool (embed (erase at))
eraseGuard (AGPat (unembed -> at) p) = GPat (embed (erase at)) (erasePattern p)
eraseGuard (AGLet b)                 = GLet (eraseBinding b)

eraseLink :: ALink -> Link
eraseLink (ATLink bop at) = TLink bop (erase at)

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

eraseProperty :: AProperty -> Property
eraseProperty = erase

------------------------------------------------------------
-- DTerm erasure

eraseDTerm :: DTerm -> Term
eraseDTerm (DTVar _ (QName _ x)) = TVar (coerce x)
eraseDTerm (DTPrim _ x)     = TPrim x
eraseDTerm DTUnit           = TUnit
eraseDTerm (DTBool _ b)     = TBool b
eraseDTerm (DTChar c)       = TChar c
eraseDTerm (DTNat _ n)      = TNat n
eraseDTerm (DTRat r)        = TRat r
eraseDTerm (DTAbs q _ b)    = TAbs q $ bind [PVar . coerce $ x] (eraseDTerm dt)
  where (x, dt) = unsafeUnbind b
eraseDTerm (DTApp _ d1 d2)  = TApp (eraseDTerm d1) (eraseDTerm d2)
eraseDTerm (DTPair _ d1 d2) = TTup [eraseDTerm d1, eraseDTerm d2]
eraseDTerm (DTCase _ bs)    = TCase (map eraseDBranch bs)
eraseDTerm (DTTyOp _ op ty) = TTyOp op ty
eraseDTerm (DTNil _)        = TList [] Nothing
eraseDTerm (DTTest _ x)     = eraseDTerm x

eraseDBranch :: DBranch -> Branch
eraseDBranch b = bind (mapTelescope eraseDGuard tel) (eraseDTerm d)
  where
    (tel, d) = unsafeUnbind b

eraseDGuard :: DGuard -> Guard
eraseDGuard (DGPat (unembed -> d) p) = GPat (embed (eraseDTerm d)) (eraseDPattern p)

eraseDPattern :: DPattern -> Pattern
eraseDPattern (DPVar _ x)      = PVar (coerce x)
eraseDPattern (DPWild _)       = PWild
eraseDPattern DPUnit           = PUnit
eraseDPattern (DPPair _ x1 x2) = PTup (map (PVar . coerce) [x1,x2])
eraseDPattern (DPInj _ s x)    = PInj s (PVar (coerce x))