amelia
/
CCHM


{-# LANGUAGE ViewPatterns #-}

{-# LANGUAGE LambdaCase #-}

{-# LANGUAGE BlockArguments #-}

{-# LANGUAGE TupleSections #-}

{-# LANGUAGE DeriveAnyClass #-}

module Elab where


import Control.Exception


import qualified Data.Map.Strict as Map

import Data.Typeable


import qualified Presyntax as P

import Syntax

import Eval

import Control.Monad

import Systems

import Data.Traversable

import qualified Data.Set as Set

import Data.Set (Set)

import Data.Foldable


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


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

 tm <- check env exp a

 for (addFormula phi env) \env ->

 let tm' = eval env tm

 in unifyTC env tm' el

 pure (InclSub (quote a) (quote fm) (quote el) 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 (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 (formula == 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 (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 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


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 <- check env d VType

 let d' = eval env dom

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

 pure (Pi s dom rng, VType)


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

 dom <- check env d VType

 let d' = eval env dom

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

 pure (Sigma s dom rng, 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 = do

 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 = do

 pure

 ( Lam "r" I $

 Lam "A" Type $

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

 , VPi "I" VI \i ->

 VPi "A" VType (const VType))


infer env P.Comp = do

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


 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 = do

 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.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)