amelia
/
CCHM


{-# LANGUAGE ViewPatterns #-}

{-# LANGUAGE LambdaCase #-}

{-# LANGUAGE BlockArguments #-}

{-# LANGUAGE TupleSections #-}

{-# LANGUAGE DeriveAnyClass #-}

module Elab where


import Control.Exception

import Control.Monad


import qualified Data.Map.Strict as Map

import qualified Data.Set as Set

import Data.Traversable

import Data.Set (Set)

import Data.Typeable

import Data.Foldable


import Eval


import qualified Presyntax as P


import Syntax


import Systems


data TypeError

 = NotInScope String

 | UnifyError UnifyError

 | WrongFaces Value [([Value], Value, Elab.TypeError)]

 | InSpan (Int, Int) (Int, Int) Elab.TypeError

 | IncompleteSystem P.Formula P.Formula

 | IncompatibleFaces (P.Formula, Term) (P.Formula, Term) Elab.TypeError

 | InvalidSystem (Set Face) (Set Face)

 deriving (Show, Typeable, Exception)


check :: Env -> P.Exp -> Value -> IO Term

check env (P.Span s e exp) wants =

 check env exp wants

 `catch` \case

 InSpan s e err -> throwIO $ InSpan s e err

 err -> throwIO $ InSpan s e err

 `catch` \er -> throwIO $ InSpan s e (UnifyError er)


check env exp (VSub a fm@(toFormula -> Just phi) el) = do

 tm <- check env exp a

 for (addFormula phi env) \env -> do

 let tm' = eval env tm

 unifyTC env tm' el

 pure (InclSub (quote a) (quote fm) (quote el) tm)


check env (P.Let v t d b) expected = do

 ty <- check env t VTypeω

 let ty' = eval env ty

 d <- check env d ty'

 let d' = eval env d

 b <- check env { names = Map.insert v (ty', d') (names env) } b expected

 pure (Let v (quote ty') d b)


check env (P.Pair a b) expected = do

 (_, fst, snd) <- isSigmaType expected

 a <- check env a fst

 let a' = eval env a

 b <- check env b (snd a')

 pure (Pair a b)


check env (P.Partial fs) ty = do

 let formula = orFormula (map fst fs)

 (extent, ty) <- isPartialType formula ty

 let ourFaces = Systems.faces formula

 extentFaces = Systems.faces extent


 unless (toDNF formula == toDNF extent) $

 throwIO $ IncompleteSystem formula extent


 let range = formulaToTm $ toDNF formula


 rangeTm <- check env range VI

 let rangeTy = eval env rangeTm


 ts <- for fs $ \(fn, tn) -> do

 tms <- for (addFormula fn env) \env -> check env tn ty

 pure (fn, head tms)


 fmap (System . FMap . Map.fromList) $ for ts \(fn, tn) -> do

 for ts \(fm, tm) -> do

 when (fn /= fm && possible (fn `P.And` fm)) do

 for_ (addFormula (fn `P.And` fm) env) $ \env ->

 unifyTC (env) (eval env tn) (eval env tm)

 `catch` \e -> throwIO (IncompatibleFaces (fn, tn) (fm, tm) e)

 pure (fn, tn)


check env exp (VPartial fm@(toFormula -> Just phi) a) = do

 case addFormula phi env of

 [] -> pure $ System (FMap mempty)

 (x:xs) -> do

 tm <- check x exp a

 for_ xs $ \e -> check e exp a

 pure tm


check env (P.Lam s b) expected = do

 expc <- isPiOrPathType expected

 case expc of

 Left (_, d, r) -> do -- function

 bd <- check env { names = Map.insert s (makeValueGluingSub d s) (names env) } b (r (VVar s))

 pure (Lam s (quote d) bd)

 Right (a, x, y) -> do

 bd <- check env { names = Map.insert s (VI, VVar s) (names env) } b (a @@ VVar s)


 let t = Lam s I bd

 t' = eval env t


 checkBoundary env [s] t'

 [ ([VI0], x)

 , ([VI1], y)

 ]


 pure (PathI (quote a) (quote x) (quote y) s bd)


check env exp expected = do

 (term, actual) <- infer env exp

 unifyTC env actual expected

 pure term


makeValueGluingSub :: Value -> String -> (Value, Value)

makeValueGluingSub ty@(VSub a phi a0) s = (ty, VOfSub a phi a0 (VVar s))

makeValueGluingSub ty s = (ty, VVar s)


addFormula :: P.Formula -> Env -> [Env]

addFormula (P.And x y) = addFormula x >=> addFormula y

addFormula (P.Or x y) = (++) <$> addFormula x <*> addFormula y

addFormula P.Top = pure

addFormula P.Bot = const []

addFormula (P.Is0 x) = \env -> pure env{ names = Map.insert x (VI, VI0) (names env) }

addFormula (P.Is1 x) = \env -> pure env{ names = Map.insert x (VI, VI1) (names env) }


unifyTC :: Env -> Value -> Value -> IO ()

unifyTC env a b = unify env a b `catch` \e -> const (throwIO (UnifyError (Mismatch a b))) (e :: UnifyError)


checkBoundary :: Env -> [String] -> Value -> [([Value], Value)] -> IO ()

checkBoundary env ns f = finish <=< go where

 go :: [([Value], Value)] -> IO [([Value], Value, Elab.TypeError)]

 go [] = pure []

 go ((ixs, vl):faces) = do

 let env' = foldr (\(x, t) env -> env { names = Map.insert x t (names env) }) env (zip ns (zip (repeat VI) ixs))

 t <- try $ unifyTC env' (foldl (@@) f ixs) vl

 case t of

 Right _ -> go faces

 Left e -> ((ixs, vl, e):) <$> go faces


 finish [] = pure ()

 finish xs = throwIO $ WrongFaces f xs


infer :: Env -> P.Exp -> IO (Term, Value)

infer env (P.Span s e exp) =

 infer env exp

 `catch` \case

 InSpan s e err -> throwIO $ InSpan s e err

 err -> throwIO $ InSpan s e err

 `catch` \er -> throwIO $ InSpan s e (UnifyError er)


infer env (P.Var s) =

 case Map.lookup s (names env) of

 Just (t, _) -> pure (Var s, t)

 Nothing -> throwIO (NotInScope s)


infer env (P.App f x) = do

 (fun, ty) <- infer env f

 funt <- isPiOrPathType ty

 case funt of

 Left (_, dom, rng) -> do

 arg <- check env x dom

 let arg' = eval env arg

 pure (App fun arg, rng arg')

 Right (a, ai0, ai1) -> do

 arg <- check env x VI

 let arg' = eval env arg

 pure (PathP (quote a) (quote ai0) (quote ai1) fun arg, a @@ arg')


infer env (P.Pi s d r) = do

 (dom, ty) <- infer env d

 case ty of

 VType -> pure VType

 VTypeω -> pure VTypeω

 _ -> throwIO . UnifyError $ NotSort ty

 let d' = eval env dom

 (rng, rng_t) <-

 infer env { names = Map.insert s (makeValueGluingSub d' s) (names env) } r

 case ty of

 VType -> pure VType

 VTypeω -> pure VTypeω

 _ -> throwIO . UnifyError $ NotSort ty

 pure (Pi s dom rng, rng_t)


infer env (P.Sigma s d r) = do

 (dom, ty) <- infer env d

 rng_t <-

 case ty of

 VType -> pure VType

 VTypeω -> pure VTypeω

 _ -> throwIO . UnifyError $ NotSort ty

 let d' = eval env dom

 rng <- check env { names = Map.insert s (d', VVar s) (names env) } r rng_t

 pure (Sigma s dom rng, rng_t)


infer env P.Type = pure (Type, VType)

infer env P.Typeω = pure (Typeω, VTypeω)

infer env P.I = pure (I, VTypeω)

infer env P.I0 = pure (I0, VI)

infer env P.I1 = pure (I1, VI)


infer env (P.Cut e t) = do

 t <- check env t VType

 let t' = eval env t

 (, t') <$> check env e t'


infer env (P.IAnd x y) = do

 x <- check env x VI

 y <- check env y VI

 pure (IAnd x y, VI)


infer env (P.IOr x y) = do

 x <- check env x VI

 y <- check env y VI

 pure (IOr x y, VI)


infer env P.Path =

 pure

 ( Lam "A" (quote index_t) $

 Lam "x" (App (Var "A") I0) $

 Lam "y" (App (Var "A") I1) $

 Path (Var "A") (Var "x") (Var "y")

 , VPi "A" index_t \a ->

 VPi "x" (a @@ VI0) \_ ->

 VPi "y" (a @@ VI1) (const VType))


infer env P.PartialT =

 pure

 ( Lam "r" I $

 Lam "A" Type $

 Partial (Var "r") (Var "A")

 , VPi "I" VI \i ->

 VPi "A" VType (const VTypeω))


infer env P.PartialP =

 pure

 ( Lam "r" I $

 Lam "A" (Partial (Var "r") Typeω) $

 PartialP (Var "r") (Var "A")

 , VPi "phi" VI \i ->

 VPi "A" (VPartial i VTypeω) (const VTypeω))


infer env P.Comp = do

 let

 u_t a r = VPi "i" VI \i -> VPartial r (a @@ i)

 index_t :: Value

 index_t = VPi "_" VI (const VType)


 pure

 ( Lam "A" (quote index_t) $

 Lam "phi" I $

 Lam "u" (quote (u_t (VVar "A") (VVar "r"))) $

 Lam "a0" (Sub (App (Var "A") I0) (Var "phi") (App (Var "u") I0)) $

 Comp (Var "A") (Var "phi") (Var "u") (Var "a0")

 , VPi "A" index_t \a ->

 VPi "phi" VI \phi ->

 VPi "u" (u_t a phi) \u ->

 VPi "_" (VSub (a @@ VI0) phi (u @@ VI0)) \_ ->

 a @@ VI1

 )


infer env P.SubT =

 pure

 ( Lam "A" Type $

 Lam "phi" I $

 Lam "u" (Partial (Var "phi") (Var "A")) $

 Sub (Var "A") (Var "phi") (Var "u")

 , VPi "A" VType \a ->

 VPi "phi" VI \phi ->

 VPi "_" (VPartial phi a) (const VType)

 )


infer env P.GlueTy =

 pure

 ( Lam "A" Type $

 Lam "phi" I $

 Lam "T" (Partial (Var "phi") Type) $

 Lam "e" (PartialP (Var "phi") (quote (equiv (VVar "T") (VVar "A")))) $

 GlueTy (Var "A") (Var "phi") (Var "T") (Var "e")

 , VPi "A" VType \a ->

 VPi "phi" VI \phi ->

 VPi "T" (VPartial phi VType) \t ->

 VPi "e" (VPartialP phi (equiv t a)) \_ ->

 VType

 )


infer env P.Glue =

 pure

 ( Lam "A" Type $

 Lam "phi" I $

 Lam "T" (Partial (Var "phi") Type) $

 Lam "e" (PartialP (Var "phi") (quote (equiv (VVar "T") (VVar "A")))) $

 Lam "t" (PartialP (Var "phi") (Var "T")) $

 Lam "a" (Var "A") $

 Glue (Var "A") (Var "phi") (Var "T") (Var "e") (Var "t") (Var "a")

 , VPi "A" VType \a ->

 VPi "phi" VI \phi ->

 VPi "T" (VPartial phi VType) \t ->

 VPi "e" (VPartialP phi (equiv t a)) \_ ->

 VPi "_" (VPartialP phi t) \_ ->

 VPi "_" a \_ -> VType

 )


infer env P.Unglue =

 pure

 ( Lam "A" Type $

 Lam "phi" I $

 Lam "T" (Partial (Var "phi") Type) $

 Lam "e" (PartialP (Var "phi") (quote (equiv (VVar "T") (VVar "A")))) $

 Lam "g" (GlueTy (Var "A") (Var "phi") (Var "T") (Var "e")) $

 Unglue (Var "A") (Var "phi") (Var "T") (Var "e") (Var "g")

 , VPi "A" VType \a ->

 VPi "phi" VI \phi ->

 VPi "T" (VPartial phi VType) \t ->

 VPi "e" (VPartialP phi (equiv t a)) \e ->

 VPi "_" (VGlue a phi t e) \_ -> a

 )


infer env P.Bool = pure (Bool, VType)

infer env P.Tt = pure (Tt, VBool)

infer env P.Ff = pure (Ff, VBool)


infer env P.If =

 pure

 (Lam "P" (Pi "_" Bool Type) $

 Lam "x" (App (Var "P") Tt) $

 Lam "y" (App (Var "P") Ff) $

 Lam "b" Bool $

 If (Var "P") (Var "x") (Var "y") (Var "b")

 , VPi "P" (VPi "_" VBool (const VType)) \p ->

 VPi "_" (p @@ VTrue) \_ ->

 VPi "_" (p @@ VFalse) \_ ->

 VPi "x" VBool \x -> p @@ x)


infer env (P.INot x) = (, VI) . INot <$> check env x VI

infer env P.Lam{} = error "can't infer type for lambda"


infer env (P.Let v t d b) = do

 ty <- check env t VTypeω

 let ty' = eval env ty

 d <- check env d ty'

 let d' = eval env d

 (b, t) <- infer env{ names = Map.insert v (ty', d') (names env) } b

 pure (Let v ty d b, t)


infer env (P.Proj1 x) = do

 (t, ty) <- infer env x

 (_, d, _) <- isSigmaType ty

 pure (Proj1 t, d)


infer env (P.Proj2 x) = do

 (t, ty) <- infer env x

 let t' = eval env t

 (_, _, r) <- isSigmaType ty

 pure (Proj2 t, r (proj1 t'))


formulaToTm :: P.Formula -> P.Exp

formulaToTm (P.Is0 i) = P.INot (P.Var i)

formulaToTm (P.Is1 i) = P.Var i

formulaToTm (P.And x y) = P.IAnd (formulaToTm x) (formulaToTm y)

formulaToTm (P.Or x y) = P.IOr (formulaToTm x) (formulaToTm y)

formulaToTm P.Top = P.I1

formulaToTm P.Bot = P.I0


checkFormula :: Env -> P.Formula -> IO ()

checkFormula env P.Top = pure ()

checkFormula env P.Bot = pure ()

checkFormula env (P.And x y) = checkFormula env x *> checkFormula env y

checkFormula env (P.Or x y) = checkFormula env x *> checkFormula env y

checkFormula env (P.Is0 x) =

 case Map.lookup x (names env) of

 Just (ty, _) -> unifyTC env ty VI

 Nothing -> throwIO (NotInScope x)

checkFormula env (P.Is1 x) =

 case Map.lookup x (names env) of

 Just (ty, _) -> unifyTC env ty VI

 Nothing -> throwIO (NotInScope x)


index_t :: Value

index_t = VPi "_" VI (const VTypeω)