amelia
/
cubical


{-# LANGUAGE LambdaCase #-}

{-# LANGUAGE DeriveAnyClass #-}

{-# LANGUAGE ScopedTypeVariables #-}

module Elab.Eval where


import Control.Monad.Reader

import Control.Exception


import qualified Data.Map.Strict as Map

import qualified Data.Sequence as Seq

import qualified Data.Set as Set

import qualified Data.Text as T

import Data.Sequence (Seq)

import Data.Traversable

import Data.Set (Set)

import Data.Typeable

import Data.Foldable

import Data.IORef

import Data.Maybe


import Elab.Monad


import Presyntax.Presyntax (Plicity(..))


import Syntax.Pretty

import Syntax


import System.IO.Unsafe


import {-# SOURCE #-} Elab.WiredIn

import Prettyprinter


eval :: Term -> ElabM Value

eval t = asks (flip eval' t)


forceIO :: MonadIO m => Value -> m Value

forceIO vl@(VNe (HMeta (MV _ cell)) args) = do

 solved <- liftIO $ readIORef cell

 case solved of

 Just vl -> forceIO $ foldl applProj vl args

 Nothing -> pure vl

forceIO x = pure x


applProj :: Value -> Projection -> Value

applProj fun (PApp p arg) = vApp p fun arg

applProj fun (PIElim l x y i) = ielim l x y fun i

applProj fun PProj1 = vProj1 fun

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 sp'

 Nothing -> pure $ VNe hd sp'

 hd -> pure $ VNe hd sp'

 where

 zonkSp (PApp p x) = PApp p <$> zonkIO x

 zonkSp (PIElim l x y i) = PIElim <$> zonkIO l <*> zonkIO x <*> zonkIO y <*> zonkIO i

 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 (VPath line x y) = VPath <$> zonkIO line <*> zonkIO x <*> zonkIO y

zonkIO (VLine line f) = VLine <$> zonkIO line <*> zonkIO f


-- Sorts

zonkIO VType = pure VType

zonkIO VTypeω = pure VTypeω


zonkIO VI = pure VI

zonkIO VI0 = pure VI0

zonkIO VI1 = pure VI1


zonkIO (VIAnd x y) = iand <$> zonkIO x <*> zonkIO y

zonkIO (VIOr x y) = ior <$> zonkIO x <*> zonkIO y

zonkIO (VINot x) = inot <$> zonkIO x


zonk :: Value -> Value

zonk = unsafePerformIO . zonkIO


eval' :: ElabEnv -> Term -> Value

eval' env (Ref x) =

 case Map.lookup x (getEnv env) of

 Just (_, vl) -> vl

 _ -> error "variable not in scope when evaluating"

eval' env (App p f x) = vApp p (eval' env f) (eval' env x)


eval' env (Lam p s t) =

 VLam p $ Closure s $ \a ->

 eval' env { getEnv = Map.insert (Bound s) (error "type of abs", a) (getEnv env) } t


eval' env (Pi p s d t) =

 VPi p (eval' env d) $ Closure s $ \a ->

 eval' env { getEnv = (Map.insert (Bound s) (error "type of abs", a) (getEnv env))} t


eval' _ (Meta m) = VNe (HMeta m) mempty


eval' env (Sigma s d t) =

 VSigma (eval' env d) $ Closure s $ \a ->

 eval' env { getEnv = Map.insert (Bound s) (error "type of abs", a) (getEnv env) } t


eval' e (Pair a b) = VPair (eval' e a) (eval' e b)


eval' e (Proj1 a) = vProj1 (eval' e a)

eval' e (Proj2 a) = vProj2 (eval' e a)


eval' _ Type = VType

eval' _ Typeω = VTypeω

eval' _ I = VI

eval' _ I0 = VI0

eval' _ I1 = VI1


eval' e (IAnd x y) = iand (eval' e x) (eval' e y)

eval' e (IOr x y) = ior (eval' e x) (eval' e y)

eval' e (INot x) = inot (eval' e x)


eval' e (PathP l a b) = VPath (eval' e l) (eval' e a) (eval' e b)

eval' e (IElim l x y f i) = ielim (eval' e l) (eval' e x) (eval' e y) (eval' e f) (eval' e i)

eval' e (PathIntro p f) = VLine (eval' e p) (eval' e f)


vApp :: Plicity -> Value -> Value -> Value

vApp p (VLam p' k) arg

 | p == p' = clCont k arg

 | otherwise = error $ "wrong plicity " ++ show p ++ " vs " ++ show p' ++ " in app " ++ show (App p (quote (VLam p' k)) (quote arg))

vApp p (VNe h sp) arg = VNe h (sp Seq.:|> PApp p arg)

vApp _ x _ = error $ "can't apply " ++ show x


(@@) :: Value -> Value -> Value

(@@) = vApp Ex

infixl 9 @@


vProj1 :: Value -> Value

vProj1 (VPair a _) = a

vProj1 (VNe h sp) = VNe h (sp Seq.:|> PProj1)

vProj1 x = error $ "can't proj1 " ++ show x


vProj2 :: Value -> Value

vProj2 (VPair _ b) = b

vProj2 (VNe h sp) = VNe h (sp Seq.:|> PProj2)

vProj2 x = error $ "can't proj2 " ++ show x


data NotEqual = NotEqual Value Value

 deriving (Show, Typeable, Exception)


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

 traverse_ (uncurry unify'Spine) (Seq.zip a a')


 go (VNe _hd (_ Seq.:|> PIElim _l x y i)) rhs =

 case force i of

 VI0 -> unify' x rhs

 VI1 -> unify' y rhs

 _ -> fail


 go (VLam p (Closure _ k)) vl = do

 t <- VVar . Bound <$> newName

 unify' (k t) (vApp p vl t)


 go vl (VLam p (Closure _ k)) = do

 t <- VVar . Bound <$> newName

 unify' (vApp p vl t) (k t)


 go (VPair a b) vl = unify' a (vProj1 vl) *> unify' b (vProj2 vl)

 go vl (VPair a b) = unify' (vProj1 vl) a *> unify' (vProj2 vl) b


 go (VPi p d (Closure _ k)) (VPi p' d' (Closure _ k')) | p == p' = do

 t <- VVar . Bound <$> newName

 unify' d d'

 unify' (k t) (k' t)


 go (VSigma d (Closure _ k)) (VSigma d' (Closure _ k')) = do

 t <- VVar . Bound <$> newName

 unify' d d'

 unify' (k t) (k' t)


 go VType VType = pure ()

 go VTypeω VTypeω = pure ()


 go VI VI = pure ()

 go VI0 VI0 = pure ()

 go VI1 VI1 = pure ()


 go (VINot x) (VINot y) = unify' x y

 go (VIAnd x y) (VIAnd x' y') = unify' x x' *> unify' y y'

 go (VIOr x y) (VIOr x' y') = unify' x x' *> unify' y y'


 go (VPath l x y) (VPath l' x' y') = do

 unify' l l'

 unify' x x'

 unify' y y'


 go (VLine l p) p' = do

 n <- VVar . Bound <$> newName

 unify (p @@ n) (ielim l (l @@ VI0) (l @@ VI1) p' n)


 go p' (VLine l p) = do

 n <- VVar . Bound <$> newName

 unify (ielim l (l @@ VI0) (l @@ VI1) p' n) (p @@ n)


 go _ _ = fail


 fail = liftIO . throwIO $ NotEqual topa topb


 unify'Spine (PApp a v) (PApp a' v')

 | a == a' = unify' v v'


 unify'Spine PProj1 PProj1 = pure ()

 unify'Spine PProj2 PProj2 = pure ()


 unify'Spine (PIElim _ _ _ i) (PIElim _ _ _ j) = unify' i j


 unify'Spine _ _ = fail


unify :: Value -> Value -> ElabM ()

unify a b = unify' a b `catchElab` \(_ :: SomeException) -> liftIO $ throwIO (NotEqual a b)


isConvertibleTo :: Value -> Value -> ElabM (Term -> Term)

VPi Im d (Closure _v k) `isConvertibleTo` ty = do

 meta <- newMeta d

 cont <- k meta `isConvertibleTo` ty

 pure (\f -> cont (App Im f (quote meta)))

VType `isConvertibleTo` VTypeω = pure id


VPi p d (Closure _ k) `isConvertibleTo` VPi p' d' (Closure _ k') | p == p' = do

 wp <- d' `isConvertibleTo` d

 n <- newName

 wp_n <- eval (Lam Ex n (wp (Ref (Bound n))))


 wp' <- k (VVar (Bound n)) `isConvertibleTo` k' (wp_n @@ VVar (Bound n))

 pure (\f -> Lam p n (wp' (App p f (wp (Ref (Bound n))))))


isConvertibleTo a b = do

 unify' a b

 pure id


newMeta :: Value -> ElabM Value

newMeta _dom = do

 n <- newName

 c <- liftIO $ newIORef Nothing

 let m = MV n c


 env <- asks getEnv


 t <- for (Map.toList env) $ \(n, _) -> pure $

 case n of

 Bound{} -> Just (PApp Ex (VVar n))

 _ -> Nothing


 pure (VNe (HMeta m) (Seq.fromList (catMaybes t)))


newName :: MonadIO m => m T.Text

newName = liftIO $ do

 x <- atomicModifyIORef _nameCounter $ \x -> (x + 1, x + 1)

 pure (T.pack (show x))


_nameCounter :: IORef Int

_nameCounter = unsafePerformIO $ newIORef 0

{-# NOINLINE _nameCounter #-}


solveMeta :: MV -> Seq Projection -> Value -> ElabM ()

solveMeta m@(MV _ cell) sp rhs = do

 env <- ask

 names <- checkSpine Set.empty sp

 checkScope (Set.fromList (Bound <$> names)) rhs

 `withNote` hsep [prettyTm (quote (VNe (HMeta m) sp)), pretty '≡', prettyTm (quote rhs)]

 let tm = quote rhs

 lam = eval' env $ foldr (Lam Ex) tm names

 liftIO . atomicModifyIORef' cell $ \case

 Just _ -> error "filled cell in solvedMeta"

 Nothing -> (Just lam, ())


checkScope :: Set Name -> Value -> ElabM ()

checkScope scope (VNe h sp) =

 do

 case h of

 HVar v@Bound{} ->

 unless (v `Set.member` scope) . liftIO . throwIO $

 NotInScope v

 HVar{} -> pure ()

 HMeta{} -> pure ()

 traverse_ checkProj sp

 where

 checkProj (PApp _ t) = checkScope scope t

 checkProj (PIElim l x y i) = traverse_ (checkScope scope) [l, x, y, i]

 checkProj PProj1 = pure ()

 checkProj PProj2 = pure ()


checkScope scope (VLam _ (Closure n k)) =

 checkScope (Set.insert (Bound n) scope) (k (VVar (Bound n)))


checkScope scope (VPi _ d (Closure n k)) = do

 checkScope scope d

 checkScope (Set.insert (Bound n) scope) (k (VVar (Bound n)))


checkScope scope (VSigma d (Closure n k)) = do

 checkScope scope d

 checkScope (Set.insert (Bound n) scope) (k (VVar (Bound n)))


checkScope s (VPair a b) = traverse_ (checkScope s) [a, b]


checkScope _ VType = pure ()

checkScope _ VTypeω = pure ()


checkScope _ VI = pure ()

checkScope _ VI0 = pure ()

checkScope _ VI1 = pure ()


checkScope s (VIAnd x y) = traverse_ (checkScope s) [x, y]

checkScope s (VIOr x y) = traverse_ (checkScope s) [x, y]

checkScope s (VINot x) = checkScope s x


checkScope s (VPath line a b) = traverse_ (checkScope s) [line, a, b]

checkScope s (VLine _ line) = checkScope s line


checkSpine :: Set Name -> Seq Projection -> ElabM [T.Text]

checkSpine scope (PApp Ex (VVar n@(Bound t)) Seq.:<| xs)

 | n `Set.member` scope = liftIO . throwIO $ NonLinearSpine n

 | otherwise = (t:) <$> checkSpine scope xs

checkSpine _ (p Seq.:<| _) = liftIO . throwIO $ SpineProj p

checkSpine _ Seq.Empty = pure []


newtype NonLinearSpine = NonLinearSpine { getDupeName :: Name }

 deriving (Show, Typeable, Exception)


newtype SpineProjection = SpineProj { getSpineProjection :: Projection }

 deriving (Show, Typeable, Exception)