PolyTyped.hs

module PolyTyped where

import Imports

-- We use Subst for type-level indices and substitutions and SubstTyped for
-- expression-level indices and substitution. To be clear which is which, we
-- qualify the imports (W for weak, S for strong).
import qualified Subst as W
import qualified SubstTyped as S

import SubstProperties 

-- We import the definition of System F types from the weakly-typed version of
-- the AST. This will allow us to write a type checker that shares these
-- types. (See TypeCheck.hs).

import Poly (Ty(..), STy(..))


-- NOTE: generated by singleton brackets in Poly

exTy :: Sing (IntTy :-> IntTy)
exTy = SIntTy :%-> SIntTy


-- | The well-typed AST for System F
data Exp :: [Ty] -> Ty -> Type where

  IntE   :: Int
         -> Exp g IntTy

  VarE   :: !(S.Idx g t)
         -> Exp g t

  LamE   :: Π (t1 :: Ty)          -- type of binder
         -> Exp (t1:g) t2         -- body of abstraction
         -> Exp g (t1 :-> t2)

  AppE   :: Exp g (t1 :-> t2)     -- function
         -> Exp g t1              -- argument
         -> Exp g t2

  TyLam  :: Exp (IncList g) t     -- bind a type variable
         -> Exp g (PolyTy t)

  TyApp  :: Exp g (PolyTy t1)     -- type function
         -> Π (t2 :: Ty)          -- type argument
         -> Exp g (W.Subst (W.SingleSub t2) t1)

-- Example, the polymorphic identity function
polyId :: Exp '[] (PolyTy (VarTy W.Z :-> VarTy W.Z))
polyId = TyLam (LamE (SVarTy W.SZ) (VarE S.Z))

-- Example, an instance of the polymorphic identity function
intId :: Exp '[] (IntTy :-> IntTy)
intId = TyApp polyId SIntTy


-- Example, an instance of the polymorphic identity function
alphaId :: Exp '[] (VarTy W.Z :-> VarTy W.Z)
alphaId = TyApp polyId (SVarTy W.SZ)




-- Although the types are *much* more explicit, the actual code of the
-- instance definition is unchanged from [Poly.hs](Poly.hs)
instance S.SubstDB Exp where

   var :: S.Idx g t -> Exp g t
   var = VarE

   subst :: S.Sub Exp g g' -> Exp g t -> Exp g' t
   subst s (IntE x)     = IntE x
   subst s (VarE x)     = S.applySub s x
   subst s (LamE ty e)  = LamE ty (S.subst (S.lift s) e)
   subst s (AppE e1 e2) = AppE (S.subst s e1) (S.subst s e2)
   subst s (TyLam e)    = TyLam (S.subst (incTy s) e)
   subst s (TyApp e t)  = TyApp (S.subst s e) t

-- | Apply a type substitution to an expression
substTy :: forall g s ty.
   Π (s :: W.Sub Ty) -> Exp g ty -> Exp (Map (W.SubstSym1 s) g) (W.Subst s ty)
substTy s (IntE x)     = IntE x
substTy s (VarE n)     = VarE $ S.mapIdx @(W.SubstSym1 s)  n
substTy s (LamE t e)   = LamE (W.sSubst s t)   (substTy s e)
substTy s (AppE e1 e2) = AppE (substTy s e1) (substTy s e2)
substTy s (TyLam (e :: Exp (IncList g) t))    
    | Refl <- axiom1 @g s 
     = TyLam (substTy (W.sLift s) e)
substTy s (TyApp (e :: Exp g (PolyTy t1)) (t :: STy t2))  
    | Refl <- axiom2 @t1 @t2 s
     = TyApp (substTy s e) (W.sSubst s t)

-- | Increment all types in a term substitution
-- Need to know that IncList ~ Map (SubstSym1 WeakSub)
incTy :: forall g g'. S.Sub Exp g g'
      -> S.Sub Exp (IncList g) (IncList g')
incTy (S.Inc (x :: S.IncBy g1)) 
  | Refl <- axiom_map1 @(W.SubstSym1 W.WeakSub) @g1 @g
  = S.Inc (S.mapInc @(W.SubstSym1 W.WeakSub)  x) 
incTy (e S.:< s1)   = substTy W.sWeakSub e S.:< incTy s1
incTy (s1 S.:<> s2) = incTy s1 S.:<> incTy s2


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

-- | Small-step evaluation of closed terms
step :: Exp '[] t -> Maybe (Exp '[] t)
step (IntE x)     = Nothing
--step (VarE n)     = error "Unbound variable"  
step (LamE t e)   = Nothing
step (AppE e1 e2) = Just $ stepApp e1 e2
step (TyLam e)    = Nothing
step (TyApp e t)  = Just $ stepTyApp e t

-- | Helper function for the AppE case. This function proves that we will
-- *always* take a step if a closed term is an application expression.
stepApp :: Exp '[] (t1 :-> t2) -> Exp '[] t1  -> Exp '[] t2
--stepApp (IntE x)       e2 = error "Type error"
--stepApp (VarE n)       e2 = error "Unbound variable"
stepApp (LamE t e1)    e2 = S.subst (S.singleSub e2) e1
stepApp (AppE e1' e2') e2 = AppE (stepApp e1' e2') e2
--stepApp (TyLam e)      e2 = error "Type error"
stepApp (TyApp e1 t1)  e2 = AppE (stepTyApp e1 t1) e2

-- | Helper function for the TyApp case. This function proves that we will
-- *always* take a step if a closed term is a type application expression.
stepTyApp :: Exp '[] (PolyTy t1) -> Π t -> Exp '[] (W.Subst (W.SingleSub t) t1)
--stepTyApp (IntE x)       e2 = error "Type error"
--stepTyApp (VarE n)       t1 = error "Unbound variable"
--stepTyApp (LamE t e1)    t1 = error "Type error"
stepTyApp (AppE e1' e2') t1 = TyApp (stepApp e1' e2') t1
stepTyApp (TyLam e1)     t1 = substTy (W.sSingleSub t1) e1
stepTyApp (TyApp e1 t2)  t1 = TyApp (stepTyApp e1 t2) t1


-- | Big-step evaluation of closed terms
-- To do this correctly, we need to define a separate type
-- for values. 
data Val :: [Ty] -> Ty -> Type where
  IntV :: Int -> Val g IntTy
  LamV :: Π (t1 :: Ty)          -- type of binder
        -> Exp (t1:g) t2        -- body of abstraction
        -> Val g (t1 :-> t2)
  TyLamV :: Exp (IncList g) t   -- bind a type variable
         -> Val g (PolyTy t)

eval :: Exp '[] t -> Val '[] t
eval (IntE x) = IntV x
-- eval (VarE n) = error "Unbound variable"
eval (LamE t e) = LamV t e
eval (AppE e1 e2) =
  case eval e1 of
    (LamV t e1') -> eval (S.subst (S.singleSub e2) e1')
eval (TyLam e) = TyLamV e
eval (TyApp e1 t) =
  case eval e1 of
    (TyLamV e1') -> eval (substTy (W.sSingleSub t) e1')



-- | Open, parallel reduction (i.e. reduce under lambda expressions)
-- This doesn't fully reduce the lambda term to normal form in one step
reduce :: forall g t. Exp g t -> Exp g t 
reduce (IntE x)     = IntE x
reduce (VarE n)     = VarE n
reduce (LamE t e)   = LamE t (reduce e)
reduce (TyLam e)    = TyLam (reduce e)
reduce (AppE e1 e2) = case reduce e1 of
  -- IntE x     -> error "type error" 
  -- TyLam e    -> error "type error"
  LamE t e1' -> S.subst (S.singleSub (reduce e2)) e1'
  e1'        -> AppE e1' (reduce e2)
reduce (TyApp e1 (t :: STy t1)) 
  | Refl <- axiom6 @t1 @g
  = case reduce e1 of
      -- IntE x    -> error "type error" 
      -- LamE t e1 -> error "type error" 
      TyLam e1' -> substTy (W.sSingleSub t) e1'
      e1'       -> TyApp e1' t