elabProgram done (for now at least)

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
+elabSection _name contents = saveState $ concat <$> forM contents elabDef
-    modify (SectionContext [] [] <|)
-    concat <$> forM contents elabDef
 
-pushVariable :: Text -> IRExpr -> NonEmpty SectionContext -> NonEmpty SectionContext
+elabProgram :: IRProgram -> IRProgram
-pushVariable name ty (SectionContext defs vars :| rest) = SectionContext defs ((name, ty) : vars) :| rest
+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 (Def name ty body) = do
-elabDef (Axiom name ty) = undefined
+    tyVars <- fromMaybe [] <$> traverse usedVars ty
-elabDef (Section name contents) = elabSection name contents
+    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