Polish up the type checker

3 years ago · 9be5402444
--- a/cubical.cabal
+++ b/cubical.cabal
@ -25,12 +25,16 @@ executable cubical
 , containers ^>= 0.6
 , bytestring ^>= 0.10

 , prettyprinter ^>= 1.7
 , prettyprinter-ansi-terminal ^>= 1.1

 other-modules: Presyntax.Lexer
 , Presyntax.Parser
 , Presyntax.Tokens
 , Presyntax.Presyntax

 , Syntax
 , Syntax.Pretty
 
 , Elab
 , Elab.Eval
--- a/src/Elab.hs
+++ b/src/Elab.hs
@ -1,15 +1,36 @@
 {-# LANGUAGE TupleSections, OverloadedStrings #-}
 {-# LANGUAGE DeriveAnyClass #-}
 module Elab where

 import Control.Monad.Reader
 import Control.Exception

 import qualified Data.Map.Strict as Map
 import Data.Typeable

 import Elab.Monad
 import Elab.Eval

 import qualified Presyntax.Presyntax as P

 import Syntax
 import Elab.Eval

 infer :: P.Expr -> ElabM (Term, NFType)
 infer (P.Var t) = (Ref (Bound t),) <$> getNfType (Bound t)
 infer (P.Span ex a b) = do
 env <- ask
 liftIO $
 runElab (infer ex) env
 `catches` [ Handler $ \e@WhileChecking{} -> throwIO e
 , Handler $ \e -> throwIO (WhileChecking a b e)
 ]

 infer (P.Var t) = do
 name <- getNameFor t
 case name of
 Builtin _ wi -> elabWiredIn wi name
 _ -> do
 nft <- getNfType name
 pure (Ref name, nft)

 infer (P.App p f x) = do
 (f, f_ty) <- infer f
@ -32,12 +53,31 @@ infer (P.Sigma s d r) = do
 r <- check r VType
 pure (Sigma s d r, VType)

 infer (P.Proj1 x) = do
 (tm, ty) <- infer x
 (d, _, wp) <- isSigmaType ty
 pure (Proj1 (wp tm), d)

 infer (P.Proj2 x) = do
 (tm, ty) <- infer x
 tm_nf <- eval tm
 (_, r, wp) <- isSigmaType ty
 pure (Proj2 (wp tm), r (vProj1 tm_nf))

 infer exp = do
 t <- newMeta VType
 tm <- check exp t
 tm <- switch $ check exp t
 pure (tm, t)

 check :: P.Expr -> NFType -> ElabM Term
 check (P.Span ex a b) ty = do
 env <- ask
 liftIO $
 runElab (check ex ty) env
 `catches` [ Handler $ \e@WhileChecking{} -> throwIO e
 , Handler $ \e -> throwIO (WhileChecking a b e)
 ]

 check (P.Lam p var body) (VPi p' dom (Closure _ rng)) | p == p' =
 assume (Bound var) dom $
 Lam p var <$> check body (rng (VVar (Bound var)))
@ -59,10 +99,13 @@ check (P.Pair a b) ty = do
 pure (wp (Pair a b))

 check exp ty = do
 (tm, has) <- infer exp
 (tm, has) <- switch $ infer exp
 unify has ty
 pure tm

 elabWiredIn :: WiredIn -> Name -> ElabM (Term, NFType)
 elabWiredIn WiType _ = pure (Type, VType)

 isPiType :: P.Plicity -> NFType -> ElabM (Value, NFType -> NFType, Term -> Term)
 isPiType p (VPi p' d (Closure _ k)) | p == p' = pure (d, k, id)
 isPiType p t = do
@ -84,4 +127,35 @@ isSigmaType t = do
 pure (dom, const rng, wp)

 identityTy :: NFType
 identityTy = VPi P.Im VType (Closure "A" $ \t -> VPi P.Ex t (Closure "_" (const t)))
 identityTy = VPi P.Im VType (Closure "A" $ \t -> VPi P.Ex t (Closure "_" (const t)))

 checkStatement :: P.Statement -> ElabM a -> ElabM a
 checkStatement (P.Decl name ty) k = do
 ty <- check ty VType
 ty_nf <- eval ty
 assume (Defined name) ty_nf k

 checkStatement (P.Defn name rhs) k = do
 ty <- asks (Map.lookup (Defined name) . getEnv)
 case ty of
 Nothing -> do
 (tm, ty) <- infer rhs
 tm_nf <- eval tm
 define (Defined name) ty tm_nf k
 Just (ty_nf, nm) -> do
 unless (nm == VVar (Defined name)) . liftIO . throwIO $
 Redefinition (Defined name)

 rhs <- check rhs ty_nf
 rhs_nf <- eval rhs
 define (Defined name) ty_nf rhs_nf k

 checkProgram :: [P.Statement] -> ElabM ElabEnv
 checkProgram [] = ask
 checkProgram (st:sts) = checkStatement st $ checkProgram sts

 newtype Redefinition = Redefinition { getRedefName :: Name }
 deriving (Show, Typeable, Exception)

 data WhileChecking = WhileChecking { startPos :: P.Posn, endPos :: P.Posn, exc :: SomeException }
 deriving (Show, Typeable, Exception)
--- a/src/Elab/Eval.hs
+++ b/src/Elab/Eval.hs
@ -42,6 +42,30 @@ applProj fun PProj2 = vProj2 fun
 force :: Value -> Value
 force = unsafePerformIO . forceIO

 -- everywhere force
 zonkIO :: Value -> IO Value
 zonkIO (VNe hd sp) = do
 sp' <- traverse zonkSp sp
 case hd of
 HMeta (MV _ cell) -> do
 solved <- liftIO $ readIORef cell
 case solved of
 Just vl -> zonkIO $ foldl applProj vl (reverse sp')
 Nothing -> pure $ VNe hd sp'
 hd -> pure $ VNe hd sp'
 where
 zonkSp (PApp p x) = PApp p <$> zonkIO x
 zonkSp PProj1 = pure PProj1
 zonkSp PProj2 = pure PProj2
 zonkIO (VLam p (Closure s k)) = pure $ VLam p (Closure s (zonk . k))
 zonkIO (VPi p d (Closure s k)) = VPi p <$> zonkIO d <*> pure (Closure s (zonk . k))
 zonkIO (VSigma d (Closure s k)) = VSigma <$> zonkIO d <*> pure (Closure s (zonk . k))
 zonkIO (VPair a b) = VPair <$> zonkIO a <*> zonkIO b
 zonkIO VType = pure VType

 zonk :: Value -> Value
 zonk = unsafePerformIO . zonkIO

 evalWithEnv :: ElabEnv -> Term -> Value
 evalWithEnv env (Ref x) =
 case Map.lookup x (getEnv env) of
@ -51,17 +75,17 @@ evalWithEnv env (App p f x) = vApp p (evalWithEnv env f) (evalWithEnv env x)

 evalWithEnv env (Lam p s t) =
 VLam p $ Closure s $ \a ->
 evalWithEnv (ElabEnv (Map.insert (Bound s) (error "type of abs", a) (getEnv env))) t
 evalWithEnv env { getEnv = Map.insert (Bound s) (error "type of abs", a) (getEnv env) } t

 evalWithEnv env (Pi p s d t) =
 VPi p (evalWithEnv env d) $ Closure s $ \a ->
 evalWithEnv (ElabEnv (Map.insert (Bound s) (error "type of abs", a) (getEnv env))) t
 evalWithEnv env { getEnv = (Map.insert (Bound s) (error "type of abs", a) (getEnv env))} t

 evalWithEnv _ (Meta m) = VNe (HMeta m) []

 evalWithEnv env (Sigma s d t) =
 VSigma (evalWithEnv env d) $ Closure s $ \a ->
 evalWithEnv (ElabEnv (Map.insert (Bound s) (error "type of abs", a) (getEnv env))) t
 evalWithEnv env { getEnv = Map.insert (Bound s) (error "type of abs", a) (getEnv env) } t

 evalWithEnv e (Pair a b) = VPair (evalWithEnv e a) (evalWithEnv e b) 

@ -91,6 +115,7 @@ data NotEqual = NotEqual Value Value
 unify :: Value -> Value -> ElabM ()
 unify topa topb = join $ go <$> forceIO topa <*> forceIO topb where
 go (VNe (HMeta mv) sp) rhs = solveMeta mv sp rhs
 go rhs (VNe (HMeta mv) sp) = solveMeta mv sp rhs

 go (VNe x a) (VNe x' a')
 | x == x', length a == length a' =
@ -118,12 +143,18 @@ unify topa topb = join $ go <$> forceIO topa <*> forceIO topb where
 unify d d'
 unify (k t) (k' t)

 go VType VType = pure ()

 go _ _ = fail

 fail = liftIO . throwIO $ NotEqual topa topb

 unifySpine (PApp a v) (PApp a' v')
 | a == a' = unify v v'

 unifySpine PProj1 PProj1 = pure ()
 unifySpine PProj2 PProj2 = pure ()

 unifySpine _ _ = fail

 isConvertibleTo :: Value -> Value -> ElabM (Term -> Term)
@ -143,9 +174,9 @@ newMeta _dom = do

 env <- asks getEnv

 t <- for (Map.toList env) $ \(n, (_, c)) -> pure $
 case c of
 VVar n' | n == n' -> Just (PApp Ex (VVar n'))
 t <- for (Map.toList env) $ \(n, _) -> pure $
 case n of
 Bound{} -> Just (PApp Ex (VVar n))
 _ -> Nothing

 pure (VNe (HMeta m) (catMaybes t))
@ -161,11 +192,12 @@ _nameCounter = unsafePerformIO $ newIORef 0

 solveMeta :: MV -> [Projection] -> Value -> ElabM ()
 solveMeta m@(MV _ cell) sp rhs = do
 env <- ask
 liftIO $ print (m, sp, rhs)
 names <- checkSpine Set.empty sp
 checkScope (Set.fromList (Bound <$> names)) rhs
 let tm = quote rhs
 lam = evalWithEnv emptyEnv $ foldr (Lam Ex) tm names
 lam = evalWithEnv env $ foldr (Lam Ex) tm names
 liftIO . atomicModifyIORef' cell $ \case
 Just _ -> error "filled cell in solvedMeta"
 Nothing -> (Just lam, ())
@ -174,9 +206,10 @@ checkScope :: Set Name -> Value -> ElabM ()
 checkScope scope (VNe h sp) =
 do
 case h of
 HVar v ->
 HVar v@Bound{} ->
 unless (v `Set.member` scope) . liftIO . throwIO $
 NotInScope v
 HVar{} -> pure ()
 HMeta{} -> pure ()
 traverse_ checkProj sp
 where
--- a/src/Elab/Monad.hs
+++ b/src/Elab/Monad.hs
@ -1,6 +1,7 @@
 {-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE DeriveAnyClass #-}
 module Elab.Monad where

 import Control.Monad.Reader
@ -8,30 +9,37 @@ import Control.Exception

 import qualified Data.Map.Strict as Map
 import Data.Map.Strict (Map)
 import Data.Text (Text)
 import Data.Typeable

 import Syntax
 import qualified Data.Text as T

 newtype ElabEnv = ElabEnv { getEnv :: Map Name (NFType, Value) }
 data ElabEnv = ElabEnv { getEnv :: Map Name (NFType, Value), nameMap :: Map Text Name, pingPong :: Int }

 newtype ElabM a = ElabM { runElab :: ElabEnv -> IO a }
 deriving (Functor, Applicative, Monad, MonadReader ElabEnv, MonadIO)
 via ReaderT ElabEnv IO

 newtype NotInScope = NotInScope { getName :: Name }
 newtype NotInScope = NotInScope { nameNotInScope :: Name }
 deriving (Show, Typeable)
 deriving anyclass (Exception)

 emptyEnv :: ElabEnv
 emptyEnv = ElabEnv mempty
 emptyEnv = ElabEnv mempty (Map.singleton (T.pack "Type") (Builtin (T.pack "Type") WiType)) 0

 assume :: Name -> Value -> ElabM a -> ElabM a
 assume nm ty = local go where
 go = ElabEnv . Map.insert nm (ty, VVar nm) . getEnv
 go x = x { getEnv = Map.insert nm (ty, VVar nm) (getEnv x), nameMap = Map.insert (getNameText nm) nm (nameMap x) }

 getNameText :: Name -> Text
 getNameText (Bound x) = x
 getNameText (Defined x) = x
 getNameText (Builtin x _) = x

 define :: Name -> Value -> Value -> ElabM a -> ElabM a
 define nm ty vl = local go where
 go = ElabEnv . Map.insert nm (ty, vl) . getEnv
 go x = x { getEnv = Map.insert nm (ty, vl) (getEnv x), nameMap = Map.insert (getNameText nm) nm (nameMap x) }

 getValue :: Name -> ElabM Value
 getValue nm = do
@ -45,4 +53,19 @@ getNfType nm = do
 vl <- asks (Map.lookup nm . getEnv)
 case vl of
 Just v -> pure (fst v)
 Nothing -> liftIO . throwIO $ NotInScope nm
 Nothing -> liftIO . throwIO $ NotInScope nm

 getNameFor :: Text -> ElabM Name
 getNameFor x = do
 vl <- asks (Map.lookup x . nameMap)
 case vl of
 Just v -> pure v
 Nothing -> liftIO . throwIO $ NotInScope (Bound x)

 switch :: ElabM a -> ElabM a
 switch k =
 do
 depth <- asks pingPong
 when (depth >= 128) $ liftIO $ throwIO StackOverflow
 local go k
 where go e = e { pingPong = pingPong e + 1 }
--- a/src/Main.hs
+++ b/src/Main.hs
@ -1,14 +1,65 @@
 {-# LANGUAGE BlockArguments #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 module Main where

 import Control.Exception

 import qualified Data.ByteString.Lazy as Bsl
 import qualified Data.Text.Encoding as T
 import qualified Data.Map.Strict as Map
 import qualified Data.Text.IO as T
 import qualified Data.Text as T
 import Data.Text ( Text )
 import Data.Foldable

 import Elab.Monad
 import Elab.Eval
 import Elab

 import Presyntax.Presyntax (Posn(Posn))
 import Presyntax.Parser
 import Presyntax.Tokens
 import Presyntax.Lexer

 import System.Exit
 import Syntax.Pretty

 main :: IO ()
 main = do
 t <- Bsl.readFile "test.tt"
 let Right tks = runAlex t parseProg
 traverse_ print tks
 env <- runElab (checkProgram tks) emptyEnv `catch` \e -> displayAndDie (T.decodeUtf8 (Bsl.toStrict t)) (e :: SomeException)
 for_ (Map.toList (getEnv env)) $ \(n, x) -> putStrLn $ show n ++ " : " ++ showValue (zonk (fst x))

 displayAndDie :: Exception e => Text -> e -> IO a
 displayAndDie lines e = do
 () <- throwIO e `catches` displayExceptions lines
 exitFailure

 displayExceptions :: Text -> [Handler ()]
 displayExceptions lines =
 [ Handler \(WhileChecking a b e) -> do
 T.putStrLn $ squiggleUnder a b lines
 displayExceptions' lines e
 , Handler \(NotEqual ta tb) -> do
 putStrLn . unlines $
 [ "\x1b[1;31merror\x1b[0m: Mismatch between actual and expected types:"
 , " * \x1b[1mActual\x1b[0m: " ++ showValue (zonk ta)
 , " * \x1b[1mExpected\x1b[0m: " ++ showValue (zonk tb)
 ]
 ]

 displayExceptions' :: Exception e => Text -> e -> IO ()
 displayExceptions' lines e = displayAndDie lines e `catch` \(_ :: ExitCode) -> pure ()

 squiggleUnder :: Posn -> Posn -> Text -> Text
 squiggleUnder (Posn l c) (Posn l' c') file
 | l == l' =
 let
 line = T.pack (show l) <> " | " <> T.lines file !! (l - 1)
 padding = T.replicate (length (show l)) (T.singleton ' ') <> " |"
 squiggle = T.replicate c " " <> T.pack "\x1b[1;31m" <> T.replicate (c' - c) "~" <> T.pack "\x1b[0m"

 in T.unlines [ padding, line, padding <> squiggle ]
 | otherwise = T.pack (show (Posn l c, Posn l' c'))
--- a/src/Presyntax/Lexer.x
+++ b/src/Presyntax/Lexer.x
@ -37,7 +37,7 @@ tokens :-
 <0> \} { always TokCBrace }
 <0> \; { always TokSemi }

 <0> \n { begin newline }
 <0> \n { just $ pushStartCode newline }

 -- newline: emit a semicolon when de-denting
 <newline> {
--- a/src/Presyntax/Parser.y
+++ b/src/Presyntax/Parser.y
@ -8,6 +8,8 @@ import Presyntax.Presyntax
 import Presyntax.Tokens
 import Presyntax.Lexer

 import Prelude hiding (span)

 }

 %name parseExp Exp
@ -23,7 +25,7 @@ import Presyntax.Lexer
 %error { parseError }

 %token
 var { Token (TokVar $$) _ _ }
 var { $$@(Token (TokVar _) _ _) }
 '(' { Token TokOParen _ _ }
 ')' { Token TokCParen _ _ }

@ -46,51 +48,55 @@ import Presyntax.Lexer

 Exp :: { Expr }
 Exp
 : Exp ExpProj { App Ex $1 $2 }
 | Exp '{' Exp '}' { App Im $1 $3 }
 : Exp ExpProj { span $1 $2 $ App Ex $1 $2 }
 | Exp '{' Exp '}' { span $1 $4 $ App Im $1 $3 }

 | '\\' LambdaList '->' Exp { makeLams $2 $4 }
 | '(' VarList ':' Exp ')' ProdTail { makePis Ex $2 $4 $6 }
 | '{' VarList ':' Exp '}' ProdTail { makePis Im $2 $4 $6 }
 | ExpProj '->' Exp { Pi Ex (T.singleton '_') $1 $3 }
 | '\\' LambdaList '->' Exp { span $1 $4 $ makeLams $2 $4 }
 | '(' VarList ':' Exp ')' ProdTail { span $1 $6 $ makePis Ex $2 $4 $6 }
 | '{' VarList ':' Exp '}' ProdTail { span $1 $6 $ makePis Im $2 $4 $6 }
 | ExpProj '->' Exp { span $1 $3 $ Pi Ex (T.singleton '_') $1 $3 }

 | '(' VarList ':' Exp ')' '*' Exp { makeSigmas $2 $4 $7 }
 | ExpProj '*'  Exp { Sigma (T.singleton '_') $1 $3 }
 | '(' VarList ':' Exp ')' '*' Exp { span $1 $7 $ makeSigmas $2 $4 $7 }
 | ExpProj '*' Exp  { span $1 $3 $ Sigma (T.singleton '_') $1 $3 }

 | ExpProj { $1 }

 ProdTail :: { Expr }
 : '(' VarList ':' Exp ')' ProdTail { makePis Ex $2 $4 $6 }
 | '{' VarList ':' Exp '}' ProdTail { makePis Im $2 $4 $6 }
 | '->' Exp { $2 }
 : '(' VarList ':' Exp ')' ProdTail { span $1 $6 $ makePis Ex $2 $4 $6 }
 | '{' VarList ':' Exp '}' ProdTail { span $1 $6 $ makePis Im $2 $4 $6 }
 | '->' Exp { span $2 $2 $ $2 }

 LambdaList :: { [(Plicity, Text)] }
 : var { [(Ex, $1)] }
 | var LambdaList { (Ex, $1):$2 }
 : var { [(Ex, getVar $1)] }
 | var LambdaList { (Ex, getVar $1):$2 }

 | '{'var'}' { [(Im, getVar $2)] }
 | '{'var'}' LambdaList { (Im, getVar $2):$4 }

 | '{'var'}' { [(Im, $2)] }
 | '{'var'}' LambdaList { (Im, $2):$4 }
 LhsList :: { [(Plicity, Text)] }
 : { [] }
 | LambdaList { $1 }

 VarList :: { [Text] }
 : var { [$1] }
 | var VarList { $1:$2 }
 : var { [getVar $1] }
 | var VarList { getVar $1:$2 }

 ExpProj :: { Expr }
 : ExpProj '.1' { Proj1 $1 }
 | ExpProj '.2' { Proj2 $1 }
 : ExpProj '.1' { span $1 $2 $ Proj1 $1 }
 | ExpProj '.2' { span $1 $2 $ Proj2 $1 }
 | Atom { $1 }

 Atom :: { Expr }
 : var { Var $1 }
 | '(' Tuple ')' { $2 }
 : var { span $1 $1 $ Var (getVar $1) }
 | '(' Tuple ')' { span $1 $3 $ $2 }

 Tuple :: { Expr }
 : Exp { $1 }
 | Exp ',' Tuple { Pair $1 $3 }
 | Exp ',' Tuple { span $1 $3 $ Pair $1 $3 }

 Statement :: { Statement }
 : var ':' Exp { Decl $1 $3 }
 | var '=' Exp { Defn $1 $3 }
 : var ':' Exp  { Decl (getVar $1) $3 }
 | var LhsList '=' Exp { Defn (getVar $1) (makeLams $2 $4) }

 Program :: { [Statement] }
 : Statement { [$1] }
@ -104,4 +110,24 @@ parseError x = alexError (show x)
 makeLams xs b = foldr (uncurry Lam) b xs
 makePis p xs t b = foldr (flip (Pi p) t) b xs
 makeSigmas xs t b = foldr (flip Sigma t) b xs

 class HasPosn a where
 startPosn :: a -> Posn
 endPosn :: a -> Posn

 instance HasPosn Token where
 startPosn (Token _ l c) = Posn l c
 endPosn (Token t l c) = Posn l (c + tokSize t)

 instance HasPosn Expr where
 startPosn (Span _ s _) = s
 startPosn _ = error "no start posn in parsed expression?"

 endPosn (Span _ _ e) = e
 endPosn _ = error "no end posn in parsed expression?"

 span s e ex = Span ex (startPosn s) (endPosn e)

 getVar (Token (TokVar s) _ _) = s
 getVar _ = error "getVar non-var"
 }
--- a/src/Presyntax/Presyntax.hs
+++ b/src/Presyntax/Presyntax.hs
@ -17,9 +17,17 @@ data Expr
 | Pair Expr Expr
 | Proj1 Expr
 | Proj2 Expr

 | Span Expr Posn Posn
 deriving (Eq, Show, Ord)

 data Statement
 = Decl Text Expr
 | Defn Text Expr
 deriving (Eq, Show, Ord)

 data Posn
 = Posn { posnLine :: {-# UNPACK #-} !Int
 , posnColm :: {-# UNPACK #-} !Int
 }
 deriving (Eq, Show, Ord)
--- a/src/Presyntax/Tokens.hs
+++ b/src/Presyntax/Tokens.hs
@ -1,6 +1,7 @@
 module Presyntax.Tokens where

 import Data.Text (Text)
 import qualified Data.Text as T

 data TokenClass
 = TokVar Text
@ -25,6 +26,23 @@ data TokenClass
 | TokSemi
 deriving (Eq, Show, Ord)

 tokSize :: TokenClass -> Int
 tokSize (TokVar x) = T.length x
 tokSize TokEof = 0
 tokSize TokLambda = 1
 tokSize TokOParen = 1
 tokSize TokOBrace = 1
 tokSize TokCBrace = 1
 tokSize TokCParen = 1
 tokSize TokStar = 1
 tokSize TokColon = 1
 tokSize TokEqual = 1
 tokSize TokComma = 1
 tokSize TokSemi = 1
 tokSize TokArrow = 2
 tokSize TokPi1 = 2
 tokSize TokPi2 = 2

 data Token
 = Token { tokenClass :: TokenClass
 , tokStartLine :: !Int
--- a/src/Syntax.hs
+++ b/src/Syntax.hs
@ -35,7 +35,13 @@ instance Show MV where
 show = ('?':) . T.unpack . mvName

 data Name
 = Bound Text
 = Bound Text
 | Defined Text
 | Builtin Text WiredIn
 deriving (Eq, Show, Ord)

 data WiredIn
 = WiType
 deriving (Eq, Show, Ord)

 type NFType = Value