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.Set as Set

import qualified Data.Text as T

import Data.Traversable

import Data.Set (Set)

import Data.Typeable

import Data.Foldable

import Data.IORef

import Data.Maybe


import {-# SOURCE #-} Elab.WiredIn

import Elab.Monad


import Presyntax.Presyntax (Plicity(..))


import Syntax


import System.IO.Unsafe

import Syntax.Pretty


eval :: Term -> ElabM Value

eval t = asks (flip evalWithEnv 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 (reverse args)

 Nothing -> pure vl

forceIO x = pure x


applProj :: Value -> Projection -> Value

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

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 (reverse sp')

 Nothing -> pure $ VNe hd sp'

 hd -> pure $ VNe hd sp'

 where

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

 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


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


evalWithEnv :: ElabEnv -> Term -> Value

evalWithEnv env (Ref x) =

 case Map.lookup x (getEnv env) of

 Just (_, vl) -> vl

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

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


evalWithEnv env (Lam p s t) =

 VLam p $ Closure s $ \a ->

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


evalWithEnv env (Pi p s d t) =

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

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


evalWithEnv _ (Meta m) = VNe (HMeta m) []


evalWithEnv env (Sigma s d t) =

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

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


evalWithEnv e (Pair a b) = VPair (evalWithEnv e a) (evalWithEnv e b)


evalWithEnv e (Proj1 a) = vProj1 (evalWithEnv e a)

evalWithEnv e (Proj2 a) = vProj2 (evalWithEnv e a)


evalWithEnv _ Type = VType

evalWithEnv _ Typeω = VTypeω

evalWithEnv _ I = VI

evalWithEnv _ I0 = VI0

evalWithEnv _ I1 = VI1


evalWithEnv e (IAnd x y) = iand (evalWithEnv e x) (evalWithEnv e y)

evalWithEnv e (IOr x y) = ior (evalWithEnv e x) (evalWithEnv e y)

evalWithEnv e (INot x) = inot (evalWithEnv e x)


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

vApp p (VLam p' k) arg = assert (p == p') $ clCont k arg

vApp p (VNe h sp) arg = VNe h (PApp p arg:sp)

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


vProj1 :: Value -> Value

vProj1 (VPair a _) = a

vProj1 (VNe h sp) = VNe h (PProj1:sp)

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


vProj2 :: Value -> Value

vProj2 (VPair _ b) = b

vProj2 (VNe h sp) = VNe h (PProj2:sp)

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) (zip a a')

 | otherwise = 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 _ _ = 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 _ _ = 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 Ex f (quote meta)))

VType `isConvertibleTo` VTypeω = pure id

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) (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 -> [Projection] -> Value -> ElabM ()

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

 env <- ask

 liftIO $ putStrLn (showValue (VNe (HMeta m) sp) ++ " ≡? " ++ showValue rhs)

 names <- checkSpine Set.empty sp

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

 let tm = quote rhs

 lam = evalWithEnv 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 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


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

checkSpine scope (PApp Ex (VVar n@(Bound t)):xs)

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

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

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

checkSpine _ [] = pure []


newtype NonLinearSpine = NonLinearSpine { getDupeName :: Name }

 deriving (Show, Typeable, Exception)


newtype SpineProjection = SpineProj { getSpineProjection :: Projection }

 deriving (Show, Typeable, Exception)