amelia
/
cubical


{-# LANGUAGE DeriveAnyClass #-}

{-# LANGUAGE GeneralizedNewtypeDeriving #-}

{-# LANGUAGE DerivingVia #-}

{-# LANGUAGE DeriveAnyClass #-}

module Elab.Monad where


import Control.Monad.Reader

import Control.Exception


import Data.Text.Prettyprint.Doc.Render.Terminal (AnsiStyle)

import qualified Data.Map.Strict as Map

import Data.Text.Prettyprint.Doc

import Data.Map.Strict (Map)

import Data.Sequence (Seq)

import Data.Text (Text)

import Data.Set (Set)

import Data.Typeable

import Data.IORef


import qualified Presyntax.Presyntax as P


import Syntax


data ElabEnv =

 ElabEnv { getEnv :: Map Name (NFType, Value)


 , nameMap :: Map Text Name

 , pingPong :: {-# UNPACK #-} !Int

 , commHook :: Doc AnsiStyle -> IO ()


 , currentSpan :: Maybe (P.Posn, P.Posn)

 , currentFile :: Maybe Text


 , whereBound :: Map Name (P.Posn, P.Posn)

 , definedNames :: Set Name


 , boundaries :: Map Name Boundary


 , unsolvedMetas :: {-# UNPACK #-} !(IORef (Map MV [(Seq Projection, Value)]))


 , loadedFiles :: [String]

 }


newtype ElabM a = ElabM { runElab :: ElabEnv -> IO a }

 deriving (Functor, Applicative, Monad, MonadReader ElabEnv, MonadIO)

 via ReaderT ElabEnv IO


emptyEnv :: IO ElabEnv

emptyEnv = do

 u <- newIORef mempty

 pure $ ElabEnv { getEnv = mempty

 , nameMap = mempty

 , pingPong = 0

 , commHook = const (pure ())


 , currentSpan = Nothing

 , currentFile = Nothing


 , whereBound = mempty

 , definedNames = mempty


 , boundaries = mempty


 , unsolvedMetas = u

 , loadedFiles = []

 }


addBoundary :: Name -> Boundary -> ElabM a -> ElabM a

addBoundary nm boundary = local (\e -> e { boundaries = Map.insert nm boundary (boundaries e)} )


assume :: Name -> Value -> (Name -> ElabM a) -> ElabM a

assume nm ty k = defineInternal nm ty VVar k


define :: Name -> Value -> Value -> (Name -> ElabM a) -> ElabM a

define nm vty val = defineInternal nm vty (const val)


makeLetDef :: Name -> Value -> Value -> (Name -> ElabM a) -> ElabM a

makeLetDef nm vty val = defineInternal nm vty (\nm -> GluedVl (HVar nm) mempty val)


replaceLetDef :: Name -> Value -> Value -> ElabM a -> ElabM a

replaceLetDef nm vty val = redefine nm vty (GluedVl (HVar nm) mempty val)


assumes :: [Name] -> Value -> ([Name] -> ElabM a) -> ElabM a

assumes nms ty k = do

 let

 go acc [] k = k acc

 go acc (x:xs) k = assume x ty $ \n -> go (n:acc) xs k

 in go [] nms k


defineInternal :: Name -> Value -> (Name -> Value) -> (Name -> ElabM a) -> ElabM a

defineInternal nm vty val k =

 do

 env <- ask

 let (env', name') = go env

 local (const env') (k name')

 where

 go x =

 let

 nm' = case Map.lookup (getNameText nm) (nameMap x) of

 Just name -> incName nm name

 Nothing -> nm

 in ( x { getEnv = Map.insert nm' (vty, val nm') (getEnv x)

 , nameMap = Map.insert (getNameText nm) nm' (nameMap x)

 , whereBound = maybe (whereBound x) (flip (Map.insert nm') (whereBound x)) (currentSpan x)

 }

 , nm')


redefine :: Name -> Value -> Value -> ElabM a -> ElabM a

redefine nm vty val = local go where

 go x = x { getEnv = Map.insert nm (vty, val) (getEnv x)

 , nameMap = Map.insert (getNameText nm) nm (nameMap x)

 , whereBound = maybe (whereBound x) (flip (Map.insert nm) (whereBound x)) (currentSpan x)

 }


getValue :: Name -> ElabM Value

getValue nm = do

 vl <- asks (Map.lookup nm . getEnv)

 case vl of

 Just v -> pure (snd v)

 Nothing -> throwElab $ NotInScope nm


getNfType :: Name -> ElabM NFType

getNfType nm = do

 vl <- asks (Map.lookup nm . getEnv)

 case vl of

 Just v -> pure (fst v)

 Nothing -> throwElab $ NotInScope nm


getNameFor :: Text -> ElabM Name

getNameFor x = do

 vl <- asks (Map.lookup x . nameMap)

 case vl of

 Just v -> pure v

 Nothing -> liftIO . throwIO $ NotInScope (Bound x 0)


switch :: ElabM a -> ElabM a

switch k =

 do

 depth <- asks pingPong

 when (depth >= 128) $ throwElab StackOverflow

 local go k

 where go e = e { pingPong = pingPong e + 1 }


newtype NotInScope = NotInScope { nameNotInScope :: Name }

 deriving (Show, Typeable)

 deriving anyclass (Exception)


data AttachedNote = AttachedNote { getNote :: Doc AnsiStyle, getExc :: SomeException }

 deriving (Show, Typeable)

 deriving anyclass (Exception)


withNote :: ElabM a -> Doc AnsiStyle -> ElabM a

withNote k note = do

 env <- ask

 liftIO $

 runElab k env

 `catch` \e -> throwIO (AttachedNote note e)


data WhileChecking = WhileChecking { startPos :: P.Posn, endPos :: P.Posn, exc :: SomeException }

 deriving (Show, Typeable, Exception)


withSpan :: P.Posn -> P.Posn -> ElabM a -> ElabM a

withSpan a b k = do

 env <- ask

 liftIO $

 runElab k env{ currentSpan = Just (a, b) }

 `catches` [ Handler $ \e@WhileChecking{} -> throwIO e

 , Handler $ \e -> throwIO (WhileChecking a b e)

 ]


data SeeAlso = SeeAlso { saStartPos :: P.Posn, saEndPos :: P.Posn, saExc :: SomeException }

 deriving (Show, Typeable, Exception)


seeAlso :: ElabM a -> Name -> ElabM a

seeAlso k nm = do

 env <- ask

 case Map.lookup nm (whereBound env) of

 Just l ->

 liftIO $ runElab k env

 `catch` \e -> throwIO (SeeAlso (fst l) (snd l) e)

 Nothing -> k


catchElab :: Exception e => ElabM a -> (e -> ElabM a) -> ElabM a

catchElab k h = do

 env <- ask

 liftIO $ runElab k env `catch` \e -> runElab (h e) env


tryElab :: Exception e => ElabM a -> ElabM (Either e a)

tryElab k = do

 env <- ask

 liftIO $ (Right <$> runElab k env) `catch` \e -> pure (Left e)


throwElab :: Exception e => e -> ElabM a

throwElab = liftIO . throwIO


incName :: Name -> Name -> Name

incName (Bound x _) n = Bound x (getNameNum n + 1)

incName (Defined x _) n = Defined x (getNameNum n + 1)

incName (ConName x _ s a) n = ConName x (getNameNum n + 1) s a