amelia
/
setoid


{-# LANGUAGE NamedFieldPuns #-}

{-# LANGUAGE OverloadedStrings #-}

{-# LANGUAGE FlexibleContexts #-}

{-# LANGUAGE DerivingVia #-}

module Elaboration where


import Control.Monad.Except

import Control.Monad.Reader

import Control.Concurrent


import qualified Data.HashMap.Strict as HashMap

import qualified Data.IntMap.Strict as IntMap

import Data.Text (Text)


import Elaboration.Monad


import Evaluate


import Presyntax


import Syntax


import System.IO.Unsafe ( unsafeDupablePerformIO )


import Value


elabNext :: MVar Int

elabNext = unsafeDupablePerformIO (newMVar 0)

{-# NOINLINE elabNext #-}


freshMeta :: Value -> ElabM Term

freshMeta expected = do

 ctx <- ask

 names <- getNames

 thisMeta <- liftIO $ do

 m <- modifyMVar elabNext $ \x -> pure (x + 1, x)

 modifyMVar_ elabMetas $ pure . IntMap.insert m (Unsolved names expected)

 pure m

 pure $ NewMeta (MV thisMeta) (elabBound ctx)


insert :: Term -> VTy -> ElabM (Term, VTy)

insert f (VPi Im _ d r) = do

 t <- freshMeta d

 t_nf <- asks (flip evaluate t . elabEnv)

 insert (App Im f t) (r $$ t_nf)

insert f x = pure (f, x)


insert' :: Term -> VTy -> ElabM (Term, VTy)

insert' t@(Lam Im _ _) ty = pure (t, ty)

insert' t ty = insert t ty


infer :: RawExpr -> ElabM (Term, VTy)

infer (RSrcPos start end expr) = local (\st -> st { elabSourcePos = (start, end) }) (infer expr)


infer (Rvar name) = ask >>= lookup where

 lookup ElabState{elabNames, elabConstrs, elabLevel} =

 case HashMap.lookup name elabNames of

 Just (l, t) -> pure (Bv (lvl2Ix elabLevel l), t)

 Nothing ->

 case HashMap.lookup name elabConstrs of

 Just t -> pure (Con name, t)

 Nothing -> typeError (NotInScope name)


infer (Rapp p x y) = do

 (x, x_ty) <-

 infer x >>= \(x, x_ty) ->

 case p of

 Ex -> insert x x_ty

 _ -> pure (x, x_ty)


 (_, d, r) <- isPiType p x_ty

 y <- check y d

 y_nf <- asks (flip evaluate y . elabEnv)


 pure (App p x y, r $$ y_nf)


infer (Rpi e v d r) = do

 d <- check d VType

 d_nf <- asks (flip evaluate d . elabEnv)

 assumeLocal v d_nf $ do

 r <- check r VType

 pure (Pi e v d r, VType)


infer (Rsigma v d r) = do

 d <- check d VType

 d_nf <- asks (flip evaluate d . elabEnv)

 assumeLocal v d_nf $ do

 r <- check r VType

 pure (Sigma v d r, VType)


infer (Rlet v t d b) = do

 t <- check t VType

 t_nf <- asks (flip evaluate t . elabEnv)


 d <- check d t_nf

 d_nf <- asks (flip evaluate d . elabEnv)


 defineLocal v t_nf d_nf $ do

 (b, ty) <- infer b

 pure (Let v t d b, ty)


infer Rtype = pure (Type, VType)


infer Rhole = do

 ty <- freshMeta VType

 ty_nf <- asks (flip evaluate ty . elabEnv)

 tm <- freshMeta ty_nf

 pure (tm, ty_nf)


infer (Rlam p v t) = do

 env <- asks elabEnv

 lvl <- asks elabLevel

 dom <- freshMeta VType

 let dom_nf = evaluate env dom

 assumeLocal v dom_nf $ do

 (b, rng) <- infer t

 pure (Lam p v b, VPi p v dom_nf (Closure env (quote (succ lvl) rng)))


infer Rtop = pure (Top, VType)

infer Runit = pure (Unit, VTop)


infer (Req a b) = do

 t <- freshMeta VType

 t_nf <- asks (flip evaluate t . elabEnv)

 a <- check a t_nf

 b <- check b t_nf

 pure (Id t a b, VType)


infer Rrefl =

 pure (Refl, forAll Im "A" VType $ \a -> forAll Im "x" a $ \x -> VEq a x x)


infer Rcoe =

 pure ( Coe

 , forAll Im "A" VType $ \a ->

 forAll Im "B" VType $ \b ->

 forAll Ex "_" (VEq VType a b) $ \_ ->

 forAll Ex "_" a $ const b

 )


infer Rcong =

 pure ( Cong

 , forAll Im "A" VType $ \a ->

 forAll Im "B" VType $ \b ->

 forAll Im "x" a $ \x ->

 forAll Im "y" a $ \y ->

 forAll Ex "f" (forAll Ex "_" a (const b)) $ \f ->

 forAll Ex "p" (VEq a x y) $ \_ ->

 VEq b (vApp f Ex x) (vApp f Ex y)

 )


infer Rsym =

 pure ( Sym

 , forAll Im "A" VType $ \a -> forAll Im "x" a $ \x -> forAll Im "y" a $ \y -> forAll Ex "p" (VEq a x y) $ \_ -> VEq a y x

 )


infer (Rproj1 e) = do

 (t, ty) <- infer e

 (_, d, _) <- isSigmaType ty

 pure (Proj1 t, d)


infer (Rproj2 e) = do

 (t, ty) <- infer e

 t_nf <- asks (flip evaluate t . elabEnv)

 (_, _, r) <- isSigmaType ty

 pure (Proj2 t, r $$ vProj1 t_nf)


infer c = do

 t <- asks elabSwitches

 when (t >= 128) $

 error $ "Unhandled case in type checker, stack overflew etc: " ++ show c


 t <- freshMeta VType

 t_nf <- asks (flip evaluate t . elabEnv)


 c <- local (\e -> e { elabSwitches = succ (elabSwitches e)}) $

 check c t_nf


 pure (c, t_nf)


check :: RawExpr -> VTy -> ElabM Term

check (RSrcPos start end expr) ty = local (\st -> st { elabSourcePos = (start, end) }) (check expr ty)


check (Rlam e v t) (VPi e' _ d r) | e == e' = do

 level <- asks (unLvl . elabLevel)

 assumeLocal v d $

 Lam e v <$> check t (r $$ vVar (Bound level))


check t (VPi Im x d r) = do

 level <- asks (unLvl . elabLevel)

 assumeLocal x d $

 Lam Im x <$> check t (r $$ vVar (Bound level))


check (Rlam e v t) ty = do

 (_, d, r) <- isPiType e ty

 level <- asks (unLvl . elabLevel)

 assumeLocal v d $

 Lam e v <$> check t (r $$ vVar (Bound level))


check (Rlet v t d b) ty = do

 t <- check t VType

 t_nf <- asks (flip evaluate t . elabEnv)


 d <- check d t_nf

 d_nf <- asks (flip evaluate d . elabEnv)


 defineLocal v t_nf d_nf $ do

 b <- check b ty

 pure (Let v t d b)


check (Rpair a b) ty = do

 (_, d, r) <- isSigmaType ty

 a <- check a d

 a_nf <- asks (flip evaluate a . elabEnv)

 b <- check b (r $$ a_nf)

 pure (Pair a b)


check e ty = do

 (new, e_ty) <- uncurry insert =<< infer e

 unify e_ty ty

 `catchError` \_ -> do

 l <- asks elabLevel

 names <- getNames

 typeError (NotEqual names (quote l (zonk ty)) (quote l (zonk e_ty)))

 pure new


isPiType :: Plicity -> VTy -> ElabM (Text, VTy, Closure)

isPiType i = go . force where

 go (VPi i' a b c)

 | i == i' = pure (a, b, c)

 go ty | not (flexible ty) = do

 l <- asks elabLevel

 names <- getNames

 typeError (NotFunction names (quote l ty))

 go ty = do

 env <- asks elabEnv

 t <- freshMeta VType

 let t_nf = evaluate env t

 assumeLocal "α" t_nf $ do

 r <- freshMeta VType

 unify ty (VPi i "α" t_nf (Closure env r))

 pure ("α", t_nf, Closure env r)


isSigmaType :: VTy -> ElabM (Text, VTy, Closure)

isSigmaType = go . force where

 go (VSigma a b c) = pure (a, b, c)

 go ty = do

 env <- asks elabEnv

 t <- freshMeta VType

 let t_nf = evaluate env t

 assumeLocal "α" t_nf $ do

 r <- freshMeta VType

 unify ty (VSigma "α" t_nf (Closure env r))

 pure ("α", t_nf, Closure env r)