From 640354bb45306180d9ad4aac85c1966b00a2c2b1 Mon Sep 17 00:00:00 2001 From: William Ball Date: Fri, 6 Dec 2024 00:40:24 -0800 Subject: [PATCH] elabProgram done (for now at least) --- lib/Elaborator.hs | 137 +++++++++++++++++++++++++++++++++++++--------- 1 file changed, 111 insertions(+), 26 deletions(-) diff --git a/lib/Elaborator.hs b/lib/Elaborator.hs index c855207..cc56df4 100644 --- a/lib/Elaborator.hs +++ b/lib/Elaborator.hs @@ -1,12 +1,14 @@ +{-# LANGUAGE TupleSections #-} + module Elaborator where import Control.Monad.Except -import Data.List (elemIndex) +import Data.List (elemIndex, lookup) import Data.List.NonEmpty ((<|)) import Errors (Error (..), Result) import Expr (Expr) import qualified Expr as E -import IR (IRDef (..), IRExpr) +import IR (IRDef (..), IRExpr, IRProgram) import qualified IR as I import Relude.Extra.Lens @@ -17,7 +19,24 @@ data SectionContext = SectionContext , sectionVars :: [(Text, IRExpr)] -- variables and their types } -type ElabMonad = StateT (NonEmpty SectionContext) Result +type ElabMonad = State SectionContext + +lookupDefInCtxt :: Text -> SectionContext -> Maybe [(Text, IRExpr)] +lookupDefInCtxt def (SectionContext defs vars) = mapMaybe helper <$> lookup def defs + where + helper dep = (dep,) <$> dep `lookup` vars + +-- looks up a definition in the context and gives a list of the variables and +-- their types that it depends on +lookupDef :: Text -> ElabMonad (Maybe [(Text, IRExpr)]) +lookupDef def = lookupDefInCtxt def <$> get + +lookupVarInCtxt :: Text -> SectionContext -> Maybe IRExpr +lookupVarInCtxt var = lookup var . sectionVars + +-- looks up a variable in the context and returns its type +lookupVar :: Text -> ElabMonad (Maybe IRExpr) +lookupVar var = lookupVarInCtxt var <$> get sectionDefsL :: Lens' SectionContext [(Text, [Text])] sectionDefsL = lens sectionDefs setter @@ -33,28 +52,98 @@ saveState :: ElabMonad a -> ElabMonad a saveState action = get >>= (action <*) . put elabSection :: Text -> [IRDef] -> ElabMonad [IRDef] -elabSection _name contents = saveState $ do - modify (SectionContext [] [] <|) - concat <$> forM contents elabDef +elabSection _name contents = saveState $ concat <$> forM contents elabDef -pushVariable :: Text -> IRExpr -> NonEmpty SectionContext -> NonEmpty SectionContext -pushVariable name ty (SectionContext defs vars :| rest) = SectionContext defs ((name, ty) : vars) :| rest +elabProgram :: IRProgram -> IRProgram +elabProgram prog = evalState (elabSection "" prog) (SectionContext [] []) + +pushVariable :: Text -> IRExpr -> SectionContext -> SectionContext +pushVariable name ty (SectionContext defs vars) = SectionContext defs ((name, ty) : vars) + +pushDefinition :: Text -> [Text] -> SectionContext -> SectionContext +pushDefinition name defVars (SectionContext defs vars) = SectionContext ((name, defVars) : defs) vars + +removeName :: Text -> ElabMonad () +removeName name = do + modify $ over sectionDefsL (filter ((/= name) . fst)) + modify $ over sectionVarsL (filter ((/= name) . fst)) + +-- find all the section variables used in an expression +usedVars :: IRExpr -> ElabMonad [(Text, IRExpr)] +usedVars (I.Var name) = maybe [] (pure . (name,)) <$> lookupVar name +usedVars I.Star = pure [] +usedVars (I.Level _) = pure [] +usedVars (I.App m n) = (++) <$> usedVars m <*> usedVars n +usedVars (I.Abs name ty ascr body) = saveState $ do + ty' <- usedVars ty + ascr' <- traverse usedVars ascr + removeName name + ((ty' ++ (ascr' ?: [])) ++) <$> usedVars body +usedVars (I.Pi name ty ascr body) = saveState $ do + ty' <- usedVars ty + ascr' <- traverse usedVars ascr + removeName name + ((ty' ++ (ascr' ?: [])) ++) <$> usedVars body +usedVars (I.Let name ascr value body) = saveState $ do + ty' <- usedVars value + ascr' <- traverse usedVars ascr + removeName name + ((ty' ++ (ascr' ?: [])) ++) <$> usedVars body + +-- traverse the body of a definition, adding the necessary section arguments to +-- any definitions made within the section +traverseBody :: IRExpr -> ElabMonad IRExpr +traverseBody (I.Var name) = do + deps <- lookupDef name + case deps of + Nothing -> pure $ I.Var name + Just [] -> pure $ I.Var name + Just ((top, _) : rest) -> pure $ foldr (I.App . I.Var . fst) (I.Var top) rest +traverseBody I.Star = pure I.Star +traverseBody e@(I.Level _) = pure e +traverseBody (I.App m n) = I.App <$> traverseBody m <*> traverseBody n +traverseBody (I.Abs name ty ascr body) = saveState $ do + ty' <- traverseBody ty + ascr' <- traverse traverseBody ascr + removeName name + I.Abs name ty' ascr' <$> traverseBody body +traverseBody (I.Pi name ty ascr body) = saveState $ do + ty' <- traverseBody ty + ascr' <- traverse traverseBody ascr + removeName name + I.Pi name ty' ascr' <$> traverseBody body +traverseBody (I.Let name ascr value body) = saveState $ do + value' <- traverseBody value + ascr' <- traverse traverseBody ascr + removeName name + I.Let name ascr' value' <$> traverseBody body + +mkPi :: (Text, IRExpr) -> IRExpr -> IRExpr +mkPi (param, typ) = I.Pi param typ Nothing + +mkAbs :: (Text, IRExpr) -> IRExpr -> IRExpr +mkAbs (param, typ) = I.Abs param typ Nothing + +generalizeType :: IRExpr -> [(Text, IRExpr)] -> IRExpr +generalizeType = foldr mkPi + +generalizeVal :: IRExpr -> [(Text, IRExpr)] -> IRExpr +generalizeVal = foldr mkAbs elabDef :: IRDef -> ElabMonad [IRDef] -elabDef (Def name ty body) = undefined -elabDef (Axiom name ty) = undefined -elabDef (Section name contents) = elabSection name contents +elabDef (Def name ty body) = do + tyVars <- fromMaybe [] <$> traverse usedVars ty + bodyVars <- usedVars body + let totalVars = tyVars ++ bodyVars + modify $ pushDefinition name (map fst totalVars) + pure [Def name (flip generalizeType totalVars <$> ty) (generalizeVal body totalVars)] +elabDef (Axiom name ty) = do + vars <- usedVars ty + modify $ pushDefinition name (map fst vars) + pure [Axiom name (generalizeType ty vars)] +elabDef (Section name contents) = saveState $ elabSection name contents elabDef (Variable name ty) = [] <$ modify' (pushVariable name ty) -traverseDef :: [Text] -> IRExpr -> ElabMonad [Text] -traverseDef seen (I.Var _) = pure seen -traverseDef seen I.Star = pure seen -traverseDef seen (I.Level _) = pure seen -traverseDef seen (I.App m n) = (++) <$> traverseDef seen m <*> traverseDef seen n - --- game plan: --- 1. add arguments to section local definitions - -- saveBinders :: State Binders a -> State Binders a -- saveBinders action = do -- binders <- get @@ -65,10 +154,6 @@ traverseDef seen (I.App m n) = (++) <$> traverseDef seen m <*> traverseDef seen -- elaborate :: IRExpr -> Expr -- elaborate ir = evalState (elaborate' ir) [] -- where --- helper :: (Monad m) => Maybe a -> (a -> m b) -> m (Maybe b) --- helper Nothing _ = pure Nothing --- helper (Just x) f = Just <$> f x --- -- elaborate' :: IRExpr -> State Binders Expr -- elaborate' (I.Var n) = do -- binders <- get @@ -78,12 +163,12 @@ traverseDef seen (I.App m n) = (++) <$> traverseDef seen m <*> traverseDef seen -- elaborate' (I.App m n) = E.App <$> elaborate' m <*> elaborate' n -- elaborate' (I.Abs x t a b) = saveBinders $ do -- t' <- elaborate' t --- a' <- helper a elaborate' +-- a' <- traverse elaborate' a -- modify (x :) -- E.Abs x t' a' <$> elaborate' b -- elaborate' (I.Pi x t a b) = saveBinders $ do -- t' <- elaborate' t --- a' <- helper a elaborate' +-- a' <- traverse elaborate' a -- modify (x :) -- E.Pi x t' a' <$> elaborate' b -- elaborate' (I.Let name Nothing val body) = saveBinders $ do