amelia
/
setoid


{-# LANGUAGE FlexibleContexts #-}

{-# LANGUAGE LambdaCase #-}

{-# LANGUAGE ViewPatterns #-}

{-# LANGUAGE BlockArguments #-}

{-# LANGUAGE OverloadedStrings #-}

module Evaluate where


import qualified Control.Exception as Exc

import Control.Monad.Except

import Control.Monad.Reader

import Control.Concurrent


import qualified Data.IntMap.Strict as IntMap

import qualified Data.Sequence as Seq

import qualified Data.Text as T


import Elaboration.Monad


import GHC.Stack (HasCallStack)


import Generics.SYB (mkT, everywhere)


import Syntax


import System.IO.Unsafe


import Value

import Data.Foldable

import Syntax.Pretty (showWithPrec)


evaluate :: HasCallStack => Env -> Term -> Value

evaluate env (Var (Bound i)) =

 case Seq.lookup i (locals env) of

 Just x -> x

 Nothing -> error $ "Variable of index " ++ show i ++ " not in scope"


evaluate _ (Con t) = VGlued (HCon t) mempty Nothing


evaluate _ Type = VType


evaluate env (Pi p t d r) = VPi p t (evaluate env d) (Closure env r)

evaluate env (Lam p t b) = VLam p t (Closure env b)

evaluate env (App p f x) = vApp (evaluate env f) p (evaluate env x)


evaluate env (Sigma t d r) = VSigma t (evaluate env d) (Closure env r)

evaluate env (Pair a b) = VPair (evaluate env a) (evaluate env b)

evaluate env (Proj1 a) = vProj1 (evaluate env a)

evaluate env (Proj2 a) = vProj2 (evaluate env a)


evaluate _ (Meta m) = VGlued (HMeta m) mempty Nothing

evaluate env (NewMeta m mask) = VGlued (HMeta m) (getVals (locals env) mask) Nothing where

 getVals Seq.Empty [] = Seq.Empty

 getVals (v Seq.:<| seq) (BDBound _:bds) = AppEx v Seq.:<| getVals seq bds

 getVals (_ Seq.:<| seq) (BDDefined _:bds) = getVals seq bds


evaluate _ Top = VTop

evaluate _ Unit = VUnit


evaluate _ Refl = VGlued HRefl mempty Nothing

evaluate _ Coe =

 function Im (T.pack "A") $ \a ->

 function Im (T.pack "B") $ \b ->

 function Ex (T.pack "p") $ \p ->

 function Ex (T.pack "x") $ \x ->

 vCoe a b p x


evaluate _ Cong =

 function Im (T.pack "A") $ \a ->

 function Im (T.pack "B") $ \b ->

 function Im (T.pack "x") $ \x ->

 function Im (T.pack "y") $ \y ->

 function Ex (T.pack "f") $ \f ->

 function Ex (T.pack "p") $ \p ->

 vCong a b x y f p


evaluate _ Sym =

 function Im (T.pack "A") $ \a ->

 function Im (T.pack "x") $ \x ->

 function Im (T.pack "y") $ \y ->

 function Ex (T.pack "p") $ \p ->

 vSym a x y p


evaluate e (Let _ _ c d) = evaluate e' d where

 e' = e { locals = evaluate e c Seq.:<| locals e }


evaluate env (Id a b c) = vId (evaluate env a) (evaluate env b) (evaluate env c)


vId :: Value -> Value -> Value -> Value

vId kind a b =

 let stuck = VEq kind a b

 solve = VEqG kind a b

 never = solve vBottom

 always = solve VTop

 in

 case force kind of

 VType ->

 case (a, b) of

 (VTop, VTop) -> always


 (VType, VType) -> always


 (VEqG _ _ _ a, b) -> vId VType a b

 (a, VEqG _ _ _ b) -> vId VType a b


 (VEq a _ _, VEq b _ _) -> vId VType a b


 (VPi i _ d r, VPi i' _ d' r')

 | i == i' ->

 solve $

 exists "p" (vId VType d d') $ \p ->

 forAll Ex "x" d $ \x ->

 vId VType (r $$ x) (r' $$ vCoe d d' p x)

 | otherwise -> never


 (VSigma _ d r, VSigma _ d' r') ->

 solve $

 exists "p" (vId VType d d') $ \p ->

 forAll Ex "x" d $ \x ->

 vId VType (r $$ x) (r' $$ vCoe d d' p x)


 (VNe _ _, _) -> stuck

 (_, VNe _ _) -> stuck


 _ -> never


 VTop -> always


 VPi i t dom cod ->

 solve $ forAll i t dom \vl -> vId (cod $$ vl) (vApp a i vl) (vApp b i vl)


 VSigma t dom cod ->

 -- a = (x, p)

 -- b = (y, q)

 -- (a, b) ≡ (c, d) : (x : A) * P x

 -- ~> (path : a == c) * coe (cong A Type P path) b == d

 let x = vProj1 a

 y = vProj1 b

 p = vProj2 a

 q = vProj2 b

 in solve $

 exists t (vId dom x y) $ \pr ->

 vId (cod $$ y) (vCoe (cod $$ x) (cod $$ y) (vCong dom VType x y (function Ex (T.pack "x") (cod $$)) pr) p) q


 VEq{} -> solve VTop


 _ -> stuck


vBottom :: Value

vBottom = forAll Im "A" VType id


vCoe :: VTy -> VTy -> Value -> Value -> Value

-- vCoe _ _ (VGlued HRefl _ _) element = element

vCoe (VPi i _ d r) ty2@(VPi i' _ d' r') p f

 | i /= i' = vApp p Ex ty2

 | otherwise =

 function i "x" $ \x ->

 let p1 = vProj1 p -- d == d'

 p2 = vProj2 p -- (x : A) -> r x == r' (coe p1 x)

 x0 = vCoe d' d (vSym VType d d' p1) x

 in vCoe (r $$ x0) (r' $$ x) (vApp p2 Ex x0) (vApp f i x0)


vCoe tyA tyB proof element = splitFibration tyA tyB $ VGlued HCoe (Seq.fromList spine) Nothing where

 spine = [AppIm tyA, AppIm tyB, AppEx proof, AppEx element]


 -- Types are split fibrations

 -- coe {A} {A} p x = x even when p /= refl

 splitFibration tA tB vstuck = unsafeDupablePerformIO $ do

 old <- readMVar elabMetas

 t <- runElab (unify tA tB) emptyElabState

 case t of

 Left _ -> do

 swapMVar elabMetas old

 pure vstuck

 Right _ -> pure element


vCong :: Value -> Value -> Value -> Value -> Value -> Value -> Value

vCong _a c _x _y g (force -> VGlued HCong (toList -> [AppIm a, AppIm _b, AppIm x, AppIm y, AppEx f, AppEx p]) _) =

 VGlued HCong (Seq.fromList [AppIm a, AppIm c, AppIm x, AppIm y, AppEx (function Ex "x" (vApp g Ex . vApp f Ex)), AppEx p]) Nothing

vCong _a b _x _ f (force -> VGlued HRefl (toList -> [AppIm _, AppIm x]) _) =

 VGlued HRefl (Seq.fromList [AppIm b, AppIm (vApp f Ex x)]) Nothing

vCong a b x y f p =

 VGlued HCong (Seq.fromList [AppIm a, AppIm b, AppIm x, AppIm y, AppEx f, AppEx p]) Nothing


vSym :: Value -> Value -> Value -> Value -> Value

vSym a _ y (VGlued HRefl _ Nothing) =

 VGlued HRefl (Seq.fromList [AppIm a, AppIm y]) Nothing

vSym _ _ _ (VGlued HSym (toList -> [_a, _y, _x, AppEx p]) Nothing) = p

vSym a x y p = VGlued HSym (Seq.fromList [AppIm a, AppIm x, AppIm y, AppEx p]) Nothing


vApp :: HasCallStack => Value -> Plicity -> Value -> Value

vApp (VGlued x s v) p r = VGlued x (s Seq.|> thing) (fmap (\v -> vApp v p r) v) where

 thing =

 case p of

 Ex -> AppEx r

 Im -> AppIm r

vApp (VLam _ _ c) _ a = c $$ a

vApp _fun _plic _arg = error "invalid application"


vProj1 :: Value -> Value

vProj1 (VPair a _) = a

vProj1 x = VProj1 x


vProj2 :: Value -> Value

vProj2 (VPair _ a) = a

vProj2 x = VProj2 x


($$) :: HasCallStack => Closure -> Value -> Value

Closure e t $$ v = evaluate e' t where

 e' = e { locals = v Seq.:<| locals e }

ClMeta (MetaFun f) $$ v = f v


forAll :: Plicity -> T.Text -> Value -> (Value -> Value) -> Value

forAll i t d = VPi i t d . ClMeta . MetaFun


exists :: T.Text -> Value -> (Value -> Value) -> Value

exists t d = VSigma t d . ClMeta . MetaFun


function :: Plicity -> T.Text -> (Value -> Value) -> Value

function i t = VLam i t . ClMeta . MetaFun


quote :: HasCallStack => Level -> Value -> Term

quote _ VType = Type


quote l (VPi p t d r) = Pi p t (quote l d) (quote (succ l) (r $$ vVar (Bound (unLvl l))))

quote l (VLam p t b) = Lam p t (quote (succ l) (b $$ vVar (Bound (unLvl l))))


quote l (VSigma t d r) = Sigma t (quote l d) (quote (succ l) (r $$ vVar (Bound (unLvl l))))

quote l (VPair a b) = Pair (quote l a) (quote l b)

quote l (VProj1 a) = Proj1 (quote l a)

quote l (VProj2 a) = Proj2 (quote l a)


quote _ VTop = Top

quote _ VUnit = Unit


quote l (VEq a b c) = Id (quote l a) (quote l b) (quote l c)

-- quote l (VEqG _ _ _ d) = quote l d

quote l (VEqG a b c _) = Id (quote l a) (quote l b) (quote l c)


quote l (VGlued v s _) = foldl app v' s where

 v' = case v of

 HVar (Bound i) -> Bv (lvl2Ix l (Lvl i))

 HCon t -> Con t

 HMeta m -> Meta m

 HRefl -> Refl

 HCoe -> Coe

 HCong -> Cong

 HSym -> Sym


 app f (AppEx t) = App Ex f (quote l t)

 app f (AppIm t) = App Im f (quote l t)

 app f SProj1 = Proj1 f

 app f SProj2 = Proj2 f


force :: Value -> Value

force = unsafeDupablePerformIO . go where

 go stuck@(VGlued (HMeta (MV m)) sp Nothing) = do

 t <- readMVar elabMetas

 case IntMap.lookup m t of

 Just (Solved vl) -> go $ foldl vAppSp vl sp

 _ -> pure stuck

 go (VGlued _ _ (Just vl)) = go vl

 go x = pure x


vAppSp :: Value -> SpineThing -> Value

vAppSp vl (AppEx f) = vApp vl Ex f

vAppSp vl (AppIm f) = vApp vl Im f

vAppSp vl SProj1 = vProj1 vl

vAppSp vl SProj2 = vProj2 vl


zonk :: Value -> Value

zonk (VLam vis var body) = VLam vis var (ClMeta (MetaFun (\v -> zonk (body $$ v))))

zonk (VPi vis var dom body) = VPi vis var (zonk dom) (ClMeta (MetaFun (\v -> zonk (body $$ v))))

zonk (VSigma var dom body) = VSigma var (zonk dom) (ClMeta (MetaFun (\v -> zonk (body $$ v))))

zonk t = everywhere (mkT force) t


unify :: VTy -> VTy -> ElabM ()

unify a b = asks elabLevel >>= flip go (a, b) where

 go, go' :: Level -> (VTy, VTy) -> ElabM ()

 go' l (VGlued h sp x, VGlued h' sp' y)

 | h == h' = goSpine l sp sp'

 | Just x <- x, Just y <- y = go l (x, y)


 -- flexible head (solve meta)

 go' l (VGlued (HMeta m) sp _, y) = solveMeta m sp y

 go' _ (x, VGlued (HMeta m) sp _) = solveMeta m sp x


 -- rigid heads (compare unfolding)

 go' l (VGlued _ _ (Just x), y) = go l (x, y)

 go' l (x, VGlued _ _ (Just y)) = go l (x, y)


 go' _ (VType, VType) = pure ()

 go' _ (VTop, VTop) = pure ()

 go' _ (VUnit, VUnit) = pure ()


 go' l (VPi i _ d r, VPi i' _ d' r') | i == i' = do

 go l (d, d')

 let i = unLvl l

 go (succ l) (r $$ vVar (Bound i), r' $$ vVar (Bound i))


 go' l (VSigma _ d r, VSigma _ d' r') = do

 go l (d, d')

 let i = unLvl l

 go (succ l) (r $$ vVar (Bound i), r' $$ vVar (Bound i))


 go' l (VLam i _ r, VLam i' _ r') | i == i' = do

 let i = unLvl l

 go (succ l) (r $$ vVar (Bound i), r' $$ vVar (Bound i))


 go' l (VLam p _ r, t) = do

 let i = unLvl l

 go (succ l) (r $$ vVar (Bound i), vApp t p (vVar (Bound i)))


 go' l (r, VLam p _ t) = do

 let i = unLvl l

 go (succ l) (vApp r p (vVar (Bound i)), t $$ vVar (Bound i))


 go' l (VEqG a b c _, VEqG a' b' c' _) = go l (a, a') *> go l (b, b') *> go l (c, c')

 go' l (VEq a b c, VEqG a' b' c' _) = go l (a, a') *> go l (b, b') *> go l (c, c')

 go' l (VEqG a b c _, VEq a' b' c') = go l (a, a') *> go l (b, b') *> go l (c, c')

 go' l (VEq a b c, VEq a' b' c') = go l (a, a') *> go l (b, b') *> go l (c, c')


 go' l (VEqG _ _ _ a, b) = go l (a, b)

 go' l (a, VEqG _ _ _ b) = go l (a, b)


 go' l (VProj1 a, VProj1 b) = go l (a, b)

 go' l (VProj2 a, VProj2 b) = go l (a, b)

 go' l (VPair a b, VPair a' b') = go l (a, a') >> go l (b, b')


 go' l (a, b) = do

 ns <- getNames

 typeError (NotEqual ns (quote l a) (quote l b))


 go l (a, b) = go' l (force a, force b)


 goSpine _ Seq.Empty Seq.Empty = pure ()

 goSpine l (AppEx x Seq.:<| xs) (AppEx y Seq.:<| ys) = do

 go l (x, y)

 goSpine l xs ys

 goSpine l (AppIm x Seq.:<| xs) (AppIm y Seq.:<| ys) = do

 go l (x, y)

 goSpine l xs ys

 goSpine l (x Seq.:<| xs) (y Seq.:<| ys) | x == y = goSpine l xs ys

 goSpine l _ _ = do

 ns <- getNames

 typeError (NotEqual ns (quote l a) (quote l b))


solveMeta :: HasCallStack => MetaVar -> Seq.Seq SpineThing -> VTy -> ElabM ()

solveMeta (MV meta) spine rhs =

 do

 level <- asks elabLevel


 pren <- invert level spine

 rhs <- rename (MV meta) pren rhs

 let solution = evaluate emptyEnv (lams (dom pren) rhs)


 -- need to deepSeq solutions here

 -- no deepSeq? no problem


 liftIO $ Exc.evaluate (length (show solution))


 liftIO . modifyMVar_ elabMetas $ pure . IntMap.insert meta (Solved solution)

 `catchError` \case

 [] -> do

 level <- asks elabLevel

 names <- getNames

 typeError (CantSolveMeta names (quote level (VGlued (HMeta (MV meta)) spine Nothing)) (quote level rhs))

 cs -> throwError cs


elabMetas :: MVar (IntMap.IntMap Meta)

elabMetas = unsafeDupablePerformIO (newMVar mempty)

{-# NOINLINE elabMetas #-}


lams :: Level -> Term -> Term

lams l = go (Lvl 0) where

 go x t | x == l = t

 go x t = Lam Ex (T.pack ("x" ++ show (unLvl x))) $ go (succ x) t


data PartialRen

 = PRen { dom :: {-# UNPACK #-} !Level

 , rng :: {-# UNPACK #-} !Level

 , sub :: IntMap.IntMap Level

 }

 deriving (Eq, Show, Ord)


liftRen :: PartialRen -> PartialRen

liftRen (PRen d r s) = PRen (succ d) (succ r) (IntMap.insert (unLvl r) d s)


invert :: Level -> Seq.Seq SpineThing -> ElabM PartialRen

invert gamma spine =

 do

 (dom, ren) <- go spine

 pure (PRen dom gamma ren)

 where

 go Seq.Empty = pure (Lvl 0, mempty)

 go (sp Seq.:|> AppEx t) = do

 (dom, ren) <- go sp

 case force t of

 VGlued (HVar (Bound l)) Seq.Empty _

 | IntMap.notMember l ren -> pure (succ dom, IntMap.insert l dom ren)

 _ -> throwError []

 go (_ Seq.:|> _) = throwError []


rename :: HasCallStack => MetaVar -> PartialRen -> Value -> ElabM Term

rename meta pren = go pren where

 go :: HasCallStack => PartialRen -> Value -> ElabM Term

 go pren (VGlued (HMeta m) sp _)

 | m == meta = throwError []

 | otherwise = goSp pren (Meta m) sp


 go pren (VGlued (HVar (Bound m)) sp _) =

 case IntMap.lookup m (sub pren) of

 Just v -> goSp pren (Bv (lvl2Ix (dom pren) v)) sp

 Nothing -> throwError []


 go pren (VGlued h sp _) = goHead h >>= \h -> goSp pren h sp where

 goHead HRefl = pure Refl

 goHead HCong = pure Cong

 goHead HCoe = pure Coe

 goHead HSym = pure Sym


 go pren (VPi p t d r) = Pi p t <$> go pren d <*> go (liftRen pren) (r $$ vVar (Bound (unLvl (rng pren))))

 go pren (VLam p t x) = Lam p t <$> go (liftRen pren) (x $$ vVar (Bound (unLvl (rng pren))))

 go _ VType = pure Type

 go _ VTop = pure Top

 go _ VUnit = pure Unit


 go pren (VSigma t d r) = Sigma t <$> go pren d <*> go (liftRen pren) (r $$ vVar (Bound (unLvl (rng pren))))

 go pren (VPair a b) = Pair <$> go pren a <*> go pren b

 go pren (VProj1 a) = Proj1 <$> go pren a

 go pren (VProj2 a) = Proj2 <$> go pren a


 go pren (VEq a b c) = Id <$> go pren a <*> go pren b <*> go pren c

 go pren (VEqG _ _ _ d) = go pren d


 -- go pren x = error (show x)


 goSp _ t Seq.Empty = pure t

 goSp pren t (sp Seq.:|> AppEx tm) = App Ex <$> goSp pren t sp <*> go pren tm

 goSp pren t (sp Seq.:|> AppIm tm) = App Im <$> goSp pren t sp <*> go pren tm

 goSp pren t (sp Seq.:|> SProj1) = Proj1 <$> goSp pren t sp

 goSp pren t (sp Seq.:|> SProj2) = Proj2 <$> goSp pren t sp