module A where
  
import GHC.Generics
import Unbound.Generics.LocallyNameless hiding (prec,empty,Data,Refl,Val)
import Unbound.Generics.LocallyNameless.Alpha

import Control.Monad
import Control.Monad.Except

import Data.Monoid (Monoid(..))

import qualified Data.List as List
import Data.Map (Map)
import qualified Data.Map as Map


import Util
import Text.PrettyPrint as PP


-- System A

type TyName = Name Ty
type ValName = Name Val

data Flag = Un | Init
  deriving (Eq, Ord, Show, Generic)

data Ty = TyVar TyName
        | TyInt
        | All (Bind [TyName] [Ty])
        | TyProd [(Ty, Flag)]  -- new
        | Exists (Bind TyName Ty) 
   deriving (Show, Generic)

data Val = TmInt Int
        | TmVar ValName
        | TApp (Ann Val) Ty  
        | Pack Ty (Ann Val)  
   deriving (Show, Generic)       
            
data Ann v = Ann v Ty
   deriving (Show, Generic)
            
data Decl   = 
    DeclVar     ValName (Embed (Ann Val))
  | DeclPrj     Int ValName (Embed (Ann Val))
  | DeclPrim    ValName (Embed (Ann Val, Prim, Ann Val))
  | DeclUnpack  TyName ValName (Embed (Ann Val))  
  | DeclMalloc  ValName (Embed [Ty]) -- new
  | DeclAssign  ValName (Embed (Ann Val, Int, Ann Val)) --new
     -- x = v1 [i] <- v2
    deriving (Show, Generic)
             
data Tm = 
    Let   (Bind Decl Tm)
  | App   (Ann Val) [Ann Val] 
  | TmIf0 (Ann Val) Tm Tm
  | Halt  Ty (Ann Val)    
    deriving (Show, Generic)

data HeapVal = 
    Tuple [Ann Val]
  | Code (Bind [TyName] (Bind [ValName] Tm))
    deriving (Show, Generic)

newtype Heap = Heap (Map ValName (Ann HeapVal)) 
  deriving (Semigroup, Monoid)

------------------------------------------------------
instance Alpha Flag
instance Alpha Ty 
instance Alpha Val 
instance Alpha a => Alpha (Ann a)
instance Alpha Decl
instance Alpha Tm

instance Subst Ty Ty where
  isvar (TyVar x) = Just (SubstName x)
  isvar _ = Nothing
instance Subst Ty Prim
instance Subst Ty Tm
instance Subst Ty (Ann Val)
instance Subst Ty Decl
instance Subst Ty Val
instance Subst Ty Flag

instance Subst Val Flag
instance Subst Val Prim
instance Subst Val Ty
instance Subst Val (Ann Val)
instance Subst Val Decl
instance Subst Val Tm
instance Subst Val Val where
  isvar (TmVar x) = Just (SubstName x)
  isvar _  = Nothing
  
------------------------------------------------------
-- Helper functions
------------------------------------------------------

-- Tag all error messages as from this module
throwErrorA :: MonadError String m => String -> m a
throwErrorA s = throwError ("A:" ++ s)

mkTyApp :: (MonadError String m, Fresh m) => Ann Val -> [Ty] -> m (Ann Val)
mkTyApp av [] = return av
mkTyApp av@(Ann _ (All bnd)) (ty:tys) = do
    (as, atys) <- unbind bnd               
    case as of 
      a:as' -> 
        let atys' = subst a ty atys in
        mkTyApp (Ann (TApp av ty) (All (bind as' atys'))) tys
      _ -> throwErrorA "mkTyApp: not a polymorphic All"
mkTyApp (Ann _ ty) _ = throwErrorA "type error: not an All"

lets :: [Decl] -> Tm -> Tm
lets [] tm = tm 
lets (d:ds) tm = Let (bind d (lets ds tm))

-----------------------------------------------------------------
  Typechecker
-----------------------------------------------------------------

type Delta = [ TyName ]
type Gamma = [ (ValName, Ty) ]

data Ctx = Ctx { getDelta :: Delta , getGamma :: Gamma }
  deriving Show
  
emptyCtx :: Ctx
emptyCtx = Ctx { getDelta = [], getGamma = [] }

checkTyVar :: Ctx -> TyName -> M ()
checkTyVar g v = do
    unless (v `List.elem` getDelta g) $ 
      throwErrorA $ "Type variable not found " ++ show v ++ "\n"
            ++ "in context: " ++ pp g

lookupTmVar :: Ctx -> ValName -> M Ty
lookupTmVar g v = do
    case lookup v (getGamma g) of
      Just s -> return s
      Nothing -> throwErrorA $ "Term variable not found " ++ show v ++ "\n"
        ++ "in context: " ++ pp g

extendTy :: TyName -> Ctx -> Ctx
extendTy n ctx = ctx { getDelta =  n : getDelta ctx }

extendTys :: [TyName] -> Ctx -> Ctx
extendTys ns ctx = foldr extendTy ctx ns

extendTm :: ValName -> Ty -> Ctx -> Ctx
extendTm n ty ctx = ctx { getGamma = (n, ty) : getGamma ctx }

extendTms :: [ValName] -> [Ty] -> Ctx -> Ctx
extendTms [] [] ctx = ctx
extendTms (n:ns) (ty:tys) ctx = extendTm n ty (extendTms ns tys ctx)
extendTms _ _ _ = error "wrong number"


tcty :: Ctx -> Ty -> M ()
tcty g  (TyVar x) =
   checkTyVar g x
tcty g  (All b) = do
   (xs, tys) <- unbind b
   let g' = extendTys xs g
   mapM_ (tcty g') tys
tcty g TyInt =  return ()
tcty g (TyProd tys) = do
   mapM_ (tcty g . fst) tys
tcty g (Exists b) = do 
  (a, ty) <- unbind b
  tcty (extendTy a g) ty


typecheckHeapVal :: Ctx -> Ann HeapVal -> M Ty
typecheckHeapVal g (Ann (Code bnd) (All bnd')) = do  
  mb  <- unbind2 bnd bnd' -- may fail
  case mb of 
    Just (as, bnd2, _, tys) -> do
      (xs, e) <- unbind bnd2
      let g' = extendTys as g
      mapM_ (tcty g') tys
      typecheck (extendTms xs tys g') e
      return (All bnd')
    Nothing -> throwErrorA "wrong # of type variables"
typecheckHeapVal g (Ann (Code bnd) _ ) = 
    throwErrorA "code must have forall type"
  
typecheckHeapVal g (Ann (Tuple es) ty) = do 
  tys <- mapM (typecheckAnnVal g) es
  let ty' = TyProd $ map (,Un) tys 
  if ty `aeq` ty' 
    then return ty
    else throwErrorA "incorrect annotation on tuple"

typecheckVal :: Ctx -> Val -> M Ty
typecheckVal g (TmVar x) = lookupTmVar g x
typecheckVal g (TmInt i)    = return TyInt
typecheckVal g (TApp av ty) = do
  tcty g ty
  ty' <- typecheckAnnVal g av
  case ty' of 
    All bnd -> do 
      (as, bs) <- unbind bnd
      case as of 
        [] -> throwErrorA "can't instantiate non-polymorphic function"
        (a:as') -> do
          let bs' = subst a ty bs
          return (All (bind as' bs'))
    _ -> throwErrorA "type error"
typecheckVal g (Pack _ _) = throwErrorA "pack must be annotated"

typecheckAnnVal :: Ctx -> Ann Val -> M Ty
typecheckAnnVal g (Ann (Pack ty1 av) ty) = do
  case ty of 
    Exists bnd -> do 
      (a, ty2) <- unbind bnd
      tcty g ty1
      ty' <- typecheckAnnVal g av
      if not (ty' `aeq` subst a ty1 ty2)
         then throwErrorA "type error"
         else return ty 
    _ -> throwErrorA "must be exists for pack"    
typecheckAnnVal g (Ann v ty) = do  
  tcty g ty
  ty' <- typecheckVal g v 
  if ty `aeq` ty'
     then return ty
     else throwErrorA $ "wrong annotation on: " ++ pp v ++ "\nInferred: " ++ pp ty' ++ "\nAnnotated: " ++ pp ty 

typecheckDecl :: Ctx -> Decl -> M Ctx
typecheckDecl g (DeclVar x (Embed av)) = do
  ty <- typecheckAnnVal g av
  return $ extendTm x ty g
typecheckDecl g (DeclPrj i x (Embed av@(Ann v _))) = do
  ty <- typecheckAnnVal g av
  case ty of 
    TyProd tys | i < length tys -> 
      return $ extendTm x (fst (tys !! i)) g
    _ -> throwErrorA "cannot project"
typecheckDecl g (DeclPrim x (Embed (av1, _, av2))) = do
  ty1 <- typecheckAnnVal g av1
  ty2 <- typecheckAnnVal g av2
  case (ty1 , ty2) of 
    (TyInt, TyInt) -> return $ extendTm x TyInt g
    _ -> throwErrorA "TypeError"
typecheckDecl g (DeclUnpack a x (Embed av)) = do
  tya <- typecheckAnnVal g av
  case tya of 
    Exists bnd -> do 
      let ty = patUnbind a bnd 
      return $ extendTy a (extendTm x ty g)
    _ -> throwErrorA "TypeError"
typecheckDecl g (DeclMalloc x (Embed tys)) = do                
  mapM_ (tcty g) tys
  return $ extendTm x (TyProd (map (,Un) tys)) g      
typecheckDecl g (DeclAssign x (Embed (av1@(Ann v1 _), i, av2))) = do
  ty1 <- typecheckAnnVal g av1 
  ty2 <- typecheckAnnVal g av2
  case ty1 of 
    TyProd tys | i < length tys -> 
      let (xs,(ty,_):ys) = splitAt i tys in
      if ty `aeq` ty2 
        then return $ extendTm x (TyProd (xs ++ (ty,Init) : ys)) g
        else throwErrorA "TypeError"
    _ -> throwErrorA "TypeError"
         
typecheck :: Ctx -> Tm -> M ()
typecheck g (Let bnd) = do
  (d,e) <- unbind bnd
  g' <- typecheckDecl g d
  typecheck g' e
typecheck g (App av es) = do
  ty <- typecheckAnnVal g av
  case ty of
   (All bnd) -> do
     (as, argtys) <- unbind bnd
     argtys' <- mapM (typecheckAnnVal g) es
     if not (null as)  
       then throwErrorA "must use type application"
       else 
         if length argtys /= length argtys'
           then throwErrorA "incorrect args"
           else unless (and (zipWith aeq argtys argtys')) $ 
              throwErrorA "arg mismatch"
   _ -> throwErrorA "need forall type"
typecheck g (TmIf0 av e1 e2) = do
  ty0 <- typecheckAnnVal g av
  typecheck g e1
  typecheck g e2
  if ty0 `aeq` TyInt then 
    return ()
  else   
    throwErrorA "TypeError"
typecheck g (Halt ty av) = do
  ty' <- typecheckAnnVal g av
  unless (ty `aeq` ty') $
    throwErrorA "type error"

         
         
progcheck :: (Tm, Heap) -> M ()
progcheck (tm, Heap m) = do
  let g = 
        Map.foldlWithKey (\ctx x (Ann _ ty) -> extendTm x ty ctx) 
        emptyCtx m
  mapM_ (typecheckHeapVal g) (Map.elems m)
  typecheck g tm

-----------------------------------------------------------------
-- Pretty-printer
-----------------------------------------------------------------

instance Display Ty where
  display (TyVar n)     = display n
  display TyInt       = return $ text "Int"
  display (All bnd) = lunbind bnd $ \ (as,tys) -> do
    da <- displayList as
    dt <- displayList tys
    if null as 
      then return $ parens dt <+> text "-> void"
      else prefix "forall" (brackets da PP.<> text "." <+> parens dt <+> text "-> void")
  display (TyProd tys) = displayTuple tys
  display (Exists bnd) = lunbind bnd $ \ (a,ty) -> do
    da <- display a 
    dt <- display ty
    prefix "exists" (da PP.<> text "." <+> dt)
    
instance Display (Ty, Flag) where    
  display (ty, fl) = do
    dty <- display ty
    let f = case fl of { Un -> "0" ; Init -> "1" }
    return $ dty PP.<> text "^" PP.<> text f
    
instance Display (ValName,Embed Ty) where                         
  display (n, Embed ty) = do
    dn <- display n
    dt <- display ty
    return $ dn PP.<> colon PP.<> dt
    
instance Display Val where                         
  display (TmInt i) = return $ int i
  display (TmVar n) = display n
  display (Pack ty e) = do 
    dty <- display ty
    de  <- display e 
    prefix "pack" (brackets (dty PP.<> comma PP.<> de))
  display (TApp av ty) = do
    dv <- display av
    dt <- display ty
    return $ dv <+> brackets dt

instance Display HeapVal where
  display (Code bnd) = lunbind bnd $ \(as, bnd2) -> lunbind bnd2 $ \(xtys, e) -> do
    ds    <- displayList as  
    dargs <- displayList xtys
    de    <- withPrec (precedence "code") $ display e
    let tyArgs = if null as then empty else brackets ds
    let tmArgs = if null xtys then empty else parens dargs
    prefix "code"  (tyArgs PP.<> tmArgs PP.<> text "." $$ de)
    
  display (Tuple es) = displayTuple es
  

instance Display a => Display (Ann a) where
  display (Ann av _) = display av

instance Display Tm where
  display (App av args) = do
    da    <- display av
    dargs <- displayList args
    let tmArgs = if null args then empty else space PP.<> parens dargs
    return $ da PP.<> tmArgs
  display (Halt ty v) = do 
    dv <- display v
    --dt <- display ty
    return $ text "halt" <+> dv -- <+> text ":" <+> dt
  display (Let bnd) = lunbind bnd $ \(d, e) -> do
    dd <- display d
    de <- display e
    return (text "let" <+> dd <+> text "in" $$ de)
  display (TmIf0 e0 e1 e2) = do
    d0 <- display e0
    d1 <- display e1
    d2 <- display e2
    prefix "if0" $ parens $ sep [d0 PP.<> comma , d1 PP.<> comma, d2]

instance Display Decl where
  display (DeclVar x (Embed av)) = do
    dx <- display x
    dv <- display av
    return $ dx <+> text "=" <+> dv
  display (DeclPrj i x (Embed av)) = do
    dx <- display x
    dv <- display av
    return $ dx <+> text "=" <+> text "pi" PP.<> int i <+> dv
  display (DeclPrim x (Embed (e1, p, e2))) = do
    dx <- display x
    let str = show p
    d1 <- display e1 
    d2 <- display e2 
    return $ dx <+> text "=" <+> d1 <+> text str <+> d2
  display (DeclUnpack a x (Embed av)) = do
    da <- display a
    dx <- display x
    dav <- display av
    return $ brackets (da PP.<> comma PP.<> dx) <+> text "=" <+> dav
  display (DeclMalloc x (Embed tys)) = do
    dx <- display x
    dtys <- displayTuple tys
    return $ dx <+> text "= malloc" PP.<> dtys
  display (DeclAssign x (Embed (av1, i, av2))) = do
    dx <- display x
    dav1 <- display av1
    dav2 <- display av2
    return $ dx <+> text "=" <+> dav1 <+> brackets (text (show i)) 
      <+>  text "<-" <+> dav2

instance Display Heap where
  display (Heap m) = do
    fcns <- mapM (\(d,v) -> do 
                     dn <- display d
                     dv <- display v
                     return (dn, dv)) (Map.toList m)
    return $ hang (text "letrec") 2 $ 
      vcat [ n <+> text "=" <+> dv | (n,dv) <- fcns ]

instance Display (Tm, Heap) where
  display (tm,h) = do
    dh <- display h
    dt <- display tm
    return $ dh $$ text "in" <+> dt

instance Display (ValName, Ty) where
  display (v, ty) = do
    dt <- display ty
    return $ text (show v) <+> colon <+> dt

instance Display Ctx where
  display (Ctx delta gamma) = do
    tyvars <- mapM display delta
    tmvars <- mapM display gamma
    return $ vcat [text "delta:" <+> hsep tyvars,
                   text "gamma:" <+> sep tmvars ]
 System A
 new
 new
new
 x = v1 [i] <- v2
----------------------------------------------------
----------------------------------------------------
 Helper functions
----------------------------------------------------
 Tag all error messages as from this module
---------------------------------------------------------------
---------------------------------------------------------------
 may fail
---------------------------------------------------------------
 Pretty-printer
---------------------------------------------------------------
dt <- display ty
 <+> text ":" <+> dt
module A where
  
import GHC.Generics
import Unbound.Generics.LocallyNameless hiding (prec,empty,Data,Refl,Val)
import Unbound.Generics.LocallyNameless.Alpha

import Control.Monad
import Control.Monad.Except

import Data.Monoid (Monoid(..))

import qualified Data.List as List
import Data.Map (Map)
import qualified Data.Map as Map


import Util
import Text.PrettyPrint as PP



type TyName = Name Ty
type ValName = Name Val

data Flag = Un | Init
  deriving (Eq, Ord, Show, Generic)

data Ty = TyVar TyName
        | TyInt
        | All (Bind [TyName] [Ty])
        | Exists (Bind TyName Ty) 
   deriving (Show, Generic)

data Val = TmInt Int
        | TmVar ValName
        | TApp (Ann Val) Ty  
        | Pack Ty (Ann Val)  
   deriving (Show, Generic)       
            
data Ann v = Ann v Ty
   deriving (Show, Generic)
            
data Decl   = 
    DeclVar     ValName (Embed (Ann Val))
  | DeclPrj     Int ValName (Embed (Ann Val))
  | DeclPrim    ValName (Embed (Ann Val, Prim, Ann Val))
  | DeclUnpack  TyName ValName (Embed (Ann Val))  
    deriving (Show, Generic)
             
data Tm = 
    Let   (Bind Decl Tm)
  | App   (Ann Val) [Ann Val] 
  | TmIf0 (Ann Val) Tm Tm
  | Halt  Ty (Ann Val)    
    deriving (Show, Generic)

data HeapVal = 
    Tuple [Ann Val]
  | Code (Bind [TyName] (Bind [ValName] Tm))
    deriving (Show, Generic)

newtype Heap = Heap (Map ValName (Ann HeapVal)) 
  deriving (Semigroup, Monoid)

instance Alpha Flag
instance Alpha Ty 
instance Alpha Val 
instance Alpha a => Alpha (Ann a)
instance Alpha Decl
instance Alpha Tm

instance Subst Ty Ty where
  isvar (TyVar x) = Just (SubstName x)
  isvar _ = Nothing
instance Subst Ty Prim
instance Subst Ty Tm
instance Subst Ty (Ann Val)
instance Subst Ty Decl
instance Subst Ty Val
instance Subst Ty Flag

instance Subst Val Flag
instance Subst Val Prim
instance Subst Val Ty
instance Subst Val (Ann Val)
instance Subst Val Decl
instance Subst Val Tm
instance Subst Val Val where
  isvar (TmVar x) = Just (SubstName x)
  isvar _  = Nothing
  

throwErrorA :: MonadError String m => String -> m a
throwErrorA s = throwError ("A:" ++ s)

mkTyApp :: (MonadError String m, Fresh m) => Ann Val -> [Ty] -> m (Ann Val)
mkTyApp av [] = return av
mkTyApp av@(Ann _ (All bnd)) (ty:tys) = do
    (as, atys) <- unbind bnd               
    case as of 
      a:as' -> 
        let atys' = subst a ty atys in
        mkTyApp (Ann (TApp av ty) (All (bind as' atys'))) tys
      _ -> throwErrorA "mkTyApp: not a polymorphic All"
mkTyApp (Ann _ ty) _ = throwErrorA "type error: not an All"

lets :: [Decl] -> Tm -> Tm
lets [] tm = tm 
lets (d:ds) tm = Let (bind d (lets ds tm))

  Typechecker

type Delta = [ TyName ]
type Gamma = [ (ValName, Ty) ]

data Ctx = Ctx { getDelta :: Delta , getGamma :: Gamma }
  deriving Show
  
emptyCtx :: Ctx
emptyCtx = Ctx { getDelta = [], getGamma = [] }

checkTyVar :: Ctx -> TyName -> M ()
checkTyVar g v = do
    unless (v `List.elem` getDelta g) $ 
      throwErrorA $ "Type variable not found " ++ show v ++ "\n"
            ++ "in context: " ++ pp g

lookupTmVar :: Ctx -> ValName -> M Ty
lookupTmVar g v = do
    case lookup v (getGamma g) of
      Just s -> return s
      Nothing -> throwErrorA $ "Term variable not found " ++ show v ++ "\n"
        ++ "in context: " ++ pp g

extendTy :: TyName -> Ctx -> Ctx
extendTy n ctx = ctx { getDelta =  n : getDelta ctx }

extendTys :: [TyName] -> Ctx -> Ctx
extendTys ns ctx = foldr extendTy ctx ns

extendTm :: ValName -> Ty -> Ctx -> Ctx
extendTm n ty ctx = ctx { getGamma = (n, ty) : getGamma ctx }

extendTms :: [ValName] -> [Ty] -> Ctx -> Ctx
extendTms [] [] ctx = ctx
extendTms (n:ns) (ty:tys) ctx = extendTm n ty (extendTms ns tys ctx)
extendTms _ _ _ = error "wrong number"


tcty :: Ctx -> Ty -> M ()
tcty g  (TyVar x) =
   checkTyVar g x
tcty g  (All b) = do
   (xs, tys) <- unbind b
   let g' = extendTys xs g
   mapM_ (tcty g') tys
tcty g TyInt =  return ()
tcty g (TyProd tys) = do
   mapM_ (tcty g . fst) tys
tcty g (Exists b) = do 
  (a, ty) <- unbind b
  tcty (extendTy a g) ty


typecheckHeapVal :: Ctx -> Ann HeapVal -> M Ty
typecheckHeapVal g (Ann (Code bnd) (All bnd')) = do  
  case mb of 
    Just (as, bnd2, _, tys) -> do
      (xs, e) <- unbind bnd2
      let g' = extendTys as g
      mapM_ (tcty g') tys
      typecheck (extendTms xs tys g') e
      return (All bnd')
    Nothing -> throwErrorA "wrong # of type variables"
typecheckHeapVal g (Ann (Code bnd) _ ) = 
    throwErrorA "code must have forall type"
  
typecheckHeapVal g (Ann (Tuple es) ty) = do 
  tys <- mapM (typecheckAnnVal g) es
  let ty' = TyProd $ map (,Un) tys 
  if ty `aeq` ty' 
    then return ty
    else throwErrorA "incorrect annotation on tuple"

typecheckVal :: Ctx -> Val -> M Ty
typecheckVal g (TmVar x) = lookupTmVar g x
typecheckVal g (TmInt i)    = return TyInt
typecheckVal g (TApp av ty) = do
  tcty g ty
  ty' <- typecheckAnnVal g av
  case ty' of 
    All bnd -> do 
      (as, bs) <- unbind bnd
      case as of 
        [] -> throwErrorA "can't instantiate non-polymorphic function"
        (a:as') -> do
          let bs' = subst a ty bs
          return (All (bind as' bs'))
    _ -> throwErrorA "type error"
typecheckVal g (Pack _ _) = throwErrorA "pack must be annotated"

typecheckAnnVal :: Ctx -> Ann Val -> M Ty
typecheckAnnVal g (Ann (Pack ty1 av) ty) = do
  case ty of 
    Exists bnd -> do 
      (a, ty2) <- unbind bnd
      tcty g ty1
      ty' <- typecheckAnnVal g av
      if not (ty' `aeq` subst a ty1 ty2)
         then throwErrorA "type error"
         else return ty 
    _ -> throwErrorA "must be exists for pack"    
typecheckAnnVal g (Ann v ty) = do  
  tcty g ty
  ty' <- typecheckVal g v 
  if ty `aeq` ty'
     then return ty
     else throwErrorA $ "wrong annotation on: " ++ pp v ++ "\nInferred: " ++ pp ty' ++ "\nAnnotated: " ++ pp ty 

typecheckDecl :: Ctx -> Decl -> M Ctx
typecheckDecl g (DeclVar x (Embed av)) = do
  ty <- typecheckAnnVal g av
  return $ extendTm x ty g
typecheckDecl g (DeclPrj i x (Embed av@(Ann v _))) = do
  ty <- typecheckAnnVal g av
  case ty of 
    TyProd tys | i < length tys -> 
      return $ extendTm x (fst (tys !! i)) g
    _ -> throwErrorA "cannot project"
typecheckDecl g (DeclPrim x (Embed (av1, _, av2))) = do
  ty1 <- typecheckAnnVal g av1
  ty2 <- typecheckAnnVal g av2
  case (ty1 , ty2) of 
    (TyInt, TyInt) -> return $ extendTm x TyInt g
    _ -> throwErrorA "TypeError"
typecheckDecl g (DeclUnpack a x (Embed av)) = do
  tya <- typecheckAnnVal g av
  case tya of 
    Exists bnd -> do 
      let ty = patUnbind a bnd 
      return $ extendTy a (extendTm x ty g)
    _ -> throwErrorA "TypeError"
typecheckDecl g (DeclMalloc x (Embed tys)) = do                
  mapM_ (tcty g) tys
  return $ extendTm x (TyProd (map (,Un) tys)) g      
typecheckDecl g (DeclAssign x (Embed (av1@(Ann v1 _), i, av2))) = do
  ty1 <- typecheckAnnVal g av1 
  ty2 <- typecheckAnnVal g av2
  case ty1 of 
    TyProd tys | i < length tys -> 
      let (xs,(ty,_):ys) = splitAt i tys in
      if ty `aeq` ty2 
        then return $ extendTm x (TyProd (xs ++ (ty,Init) : ys)) g
        else throwErrorA "TypeError"
    _ -> throwErrorA "TypeError"
         
typecheck :: Ctx -> Tm -> M ()
typecheck g (Let bnd) = do
  (d,e) <- unbind bnd
  g' <- typecheckDecl g d
  typecheck g' e
typecheck g (App av es) = do
  ty <- typecheckAnnVal g av
  case ty of
   (All bnd) -> do
     (as, argtys) <- unbind bnd
     argtys' <- mapM (typecheckAnnVal g) es
     if not (null as)  
       then throwErrorA "must use type application"
       else 
         if length argtys /= length argtys'
           then throwErrorA "incorrect args"
           else unless (and (zipWith aeq argtys argtys')) $ 
              throwErrorA "arg mismatch"
   _ -> throwErrorA "need forall type"
typecheck g (TmIf0 av e1 e2) = do
  ty0 <- typecheckAnnVal g av
  typecheck g e1
  typecheck g e2
  if ty0 `aeq` TyInt then 
    return ()
  else   
    throwErrorA "TypeError"
typecheck g (Halt ty av) = do
  ty' <- typecheckAnnVal g av
  unless (ty `aeq` ty') $
    throwErrorA "type error"

         
         
progcheck :: (Tm, Heap) -> M ()
progcheck (tm, Heap m) = do
  let g = 
        Map.foldlWithKey (\ctx x (Ann _ ty) -> extendTm x ty ctx) 
        emptyCtx m
  mapM_ (typecheckHeapVal g) (Map.elems m)
  typecheck g tm


instance Display Ty where
  display (TyVar n)     = display n
  display TyInt       = return $ text "Int"
  display (All bnd) = lunbind bnd $ \ (as,tys) -> do
    da <- displayList as
    dt <- displayList tys
    if null as 
      then return $ parens dt <+> text "-> void"
      else prefix "forall" (brackets da PP.<> text "." <+> parens dt <+> text "-> void")
  display (TyProd tys) = displayTuple tys
  display (Exists bnd) = lunbind bnd $ \ (a,ty) -> do
    da <- display a 
    dt <- display ty
    prefix "exists" (da PP.<> text "." <+> dt)
    
instance Display (Ty, Flag) where    
  display (ty, fl) = do
    dty <- display ty
    let f = case fl of { Un -> "0" ; Init -> "1" }
    return $ dty PP.<> text "^" PP.<> text f
    
instance Display (ValName,Embed Ty) where                         
  display (n, Embed ty) = do
    dn <- display n
    dt <- display ty
    return $ dn PP.<> colon PP.<> dt
    
instance Display Val where                         
  display (TmInt i) = return $ int i
  display (TmVar n) = display n
  display (Pack ty e) = do 
    dty <- display ty
    de  <- display e 
    prefix "pack" (brackets (dty PP.<> comma PP.<> de))
  display (TApp av ty) = do
    dv <- display av
    dt <- display ty
    return $ dv <+> brackets dt

instance Display HeapVal where
  display (Code bnd) = lunbind bnd $ \(as, bnd2) -> lunbind bnd2 $ \(xtys, e) -> do
    ds    <- displayList as  
    dargs <- displayList xtys
    de    <- withPrec (precedence "code") $ display e
    let tyArgs = if null as then empty else brackets ds
    let tmArgs = if null xtys then empty else parens dargs
    prefix "code"  (tyArgs PP.<> tmArgs PP.<> text "." $$ de)
    
  display (Tuple es) = displayTuple es
  

instance Display a => Display (Ann a) where
  display (Ann av _) = display av

instance Display Tm where
  display (App av args) = do
    da    <- display av
    dargs <- displayList args
    let tmArgs = if null args then empty else space PP.<> parens dargs
    return $ da PP.<> tmArgs
  display (Halt ty v) = do 
    dv <- display v
  display (Let bnd) = lunbind bnd $ \(d, e) -> do
    dd <- display d
    de <- display e
    return (text "let" <+> dd <+> text "in" $$ de)
  display (TmIf0 e0 e1 e2) = do
    d0 <- display e0
    d1 <- display e1
    d2 <- display e2
    prefix "if0" $ parens $ sep [d0 PP.<> comma , d1 PP.<> comma, d2]

instance Display Decl where
  display (DeclVar x (Embed av)) = do
    dx <- display x
    dv <- display av
    return $ dx <+> text "=" <+> dv
  display (DeclPrj i x (Embed av)) = do
    dx <- display x
    dv <- display av
    return $ dx <+> text "=" <+> text "pi" PP.<> int i <+> dv
  display (DeclPrim x (Embed (e1, p, e2))) = do
    dx <- display x
    let str = show p
    d1 <- display e1 
    d2 <- display e2 
    return $ dx <+> text "=" <+> d1 <+> text str <+> d2
  display (DeclUnpack a x (Embed av)) = do
    da <- display a
    dx <- display x
    dav <- display av
    return $ brackets (da PP.<> comma PP.<> dx) <+> text "=" <+> dav
  display (DeclMalloc x (Embed tys)) = do
    dx <- display x
    dtys <- displayTuple tys
    return $ dx <+> text "= malloc" PP.<> dtys
  display (DeclAssign x (Embed (av1, i, av2))) = do
    dx <- display x
    dav1 <- display av1
    dav2 <- display av2
    return $ dx <+> text "=" <+> dav1 <+> brackets (text (show i)) 
      <+>  text "<-" <+> dav2

instance Display Heap where
  display (Heap m) = do
    fcns <- mapM (\(d,v) -> do 
                     dn <- display d
                     dv <- display v
                     return (dn, dv)) (Map.toList m)
    return $ hang (text "letrec") 2 $ 
      vcat [ n <+> text "=" <+> dv | (n,dv) <- fcns ]

instance Display (Tm, Heap) where
  display (tm,h) = do
    dh <- display h
    dt <- display tm
    return $ dh $$ text "in" <+> dt

instance Display (ValName, Ty) where
  display (v, ty) = do
    dt <- display ty
    return $ text (show v) <+> colon <+> dt

instance Display Ctx where
  display (Ctx delta gamma) = do
    tyvars <- mapM display delta
    tmvars <- mapM display gamma
    return $ vcat [text "delta:" <+> hsep tyvars,
                   text "gamma:" <+> sep tmvars ]