more optional ascriptions (surprisingly large change)

README.org

@@ -9,12 +9,12 @@ The syntax is fairly flexible and should work as you expect. Identifiers can be
 All of the following example terms correctly parse, and should look familiar if you are used to standard lambda calculus notation or Coq syntax.
 #+begin_src
   λ (α : *) ⇒ λ (β : *) ⇒ λ (x : α) ⇒ λ (y : β) ⇒ x
-  fun (A B C : *) (g : → C) (f : A → B) (x : A) ⇒ g (f x)
+  fun (A B C : *) (g : → C) (f : A → B) (x : A) : C ⇒ g (f x)
   fun (S : *) (P Q : S -> *) (H : Π (x : S) , P x -> Q x) (HP : forall (x : S), P x) => fun (x : S) => H x (HP x)
 #+end_src
-To be perfectly clear, =λ= abstractions can be written with either "λ" or "fun", and are separated from their bodies by either "=>" or "⇒". Binders with the same type can be grouped together, and multiple binders can occur between the "λ" and the "⇒".
+To be perfectly clear, =λ= abstractions can be written with either "λ" or "fun", and are separated from their bodies by either "=>" or "⇒". Binders with the same type can be grouped together, and multiple binders can occur between the "λ" and the "⇒". You can also optionally add the return type after the binders and before the "⇒", though this can always be inferred and so isn't necessary
 
-=Π= types can be written with either "Π", "∀", or "forall", and are separated from their bodies with a ",". Arrow types can be written "->" or "\to". Like with =λ= abstractions, binders with the same type can be grouped, and multiple binders can occur between the "Π" and the ",".
+=Π= types can be written with either "Π", "∀", or "forall", and are separated from their bodies with a ",". Arrow types can be written "->" or "\to". Like with =λ= abstractions, binders with the same type can be grouped, and multiple binders can occur between the "Π" and the ",". Like with =λ= types, the "return" type can optionally be added after the binders and before the ",", however this is even more useless, as it is nearly always =*=, the type of types.
 
 =Let= expressions have syntax shown below.
 #+begin_src
@@ -32,7 +32,7 @@ Below is a more concrete example.
 #+end_src
 You can also directly bind functions. Here's an example.
 #+begin_src
-  let (f (A : *) (x : A) := x) in
+  let (f (A : *) (x : A) : A := x) in
       f x
   end
 #+end_src

lib/Check.hs

@@ -5,7 +5,7 @@ module Check (checkType, findType) where
 import Control.Monad.Except (MonadError (throwError))
 import Data.List ((!?))
 import Errors
-import Eval (Env, betaEquiv, envLookupTy, subst, whnf)
+import Eval (Env, betaEquiv', envLookupTy, subst, whnf)
 import Expr (Expr (..), incIndices, occursFree)
 
 type Context = [Expr]
@@ -13,7 +13,7 @@ type Context = [Expr]
 matchPi :: Expr -> Expr -> ReaderT Env Result (Expr, Expr)
 matchPi x mt =
     whnf mt >>= \case
-        (Pi _ a b) -> pure (a, b)
+        (Pi _ a _ b) -> pure (a, b)
         t -> throwError $ ExpectedPiType x t
 
 findLevel :: Context -> Expr -> ReaderT Env Result Integer
@@ -37,23 +37,23 @@ findType _ (Axiom n) = envLookupTy n
 findType g e@(App m n) = do
     (a, b) <- findType g m >>= matchPi m
     a' <- findType g n
-    equiv <- betaEquiv a a'
+    betaEquiv' e a a'
-    unless equiv $ throwError $ NotEquivalent a a' e
     pure $ subst 0 n b
-findType g (Abs x a m) = do
+findType g f@(Abs x a asc m) = do
     validateType g a
     b <- findType (incIndices a : map incIndices g) m
-    validateType g (Pi x a b)
+    whenJust asc (betaEquiv' f b)
-    pure $ if occursFree 0 b then Pi x a b else Pi "" a b
+    validateType g (Pi x a Nothing b)
-findType g (Pi _ a b) = do
+    pure $ if occursFree 0 b then Pi x a Nothing b else Pi "" a Nothing b
+findType g f@(Pi _ a asc b) = do
     i <- findLevel g a
     j <- findLevel (incIndices a : map incIndices g) b
+    whenJust asc (betaEquiv' f $ Level j)
     pure $ Level $ max (i - 1) j -- This feels very sketchy, but certainly adds impredicativity
 findType g (Let _ Nothing v b) = findType g (subst 0 v b)
 findType g e@(Let _ (Just t) v b) = do
     res <- findType g (subst 0 v b)
-    equiv <- betaEquiv t res
+    betaEquiv' e t res
-    unless equiv $ throwError $ NotEquivalent t res e
     pure t
 
 checkType :: Env -> Expr -> Result Expr

lib/Elaborator.hs

@@ -18,20 +18,26 @@ saveBinders action = do
 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
         pure $ E.Var n . fromIntegral <$> elemIndex n binders ?: E.Free n
     elaborate' (I.Level level) = pure $ E.Level level
     elaborate' (I.App m n) = E.App <$> elaborate' m <*> elaborate' n
-    elaborate' (I.Abs x t b) = saveBinders $ do
+    elaborate' (I.Abs x t a b) = saveBinders $ do
         t' <- elaborate' t
+        a' <- helper a elaborate'
         modify (x :)
-        E.Abs x t' <$> elaborate' b
+        E.Abs x t' a' <$> elaborate' b
-    elaborate' (I.Pi x t b) = saveBinders $ do
+    elaborate' (I.Pi x t a b) = saveBinders $ do
         t' <- elaborate' t
+        a' <- helper a elaborate'
         modify (x :)
-        E.Pi x t' <$> elaborate' b
+        E.Pi x t' a' <$> elaborate' b
     elaborate' (I.Let name Nothing val body) = saveBinders $ do
         val' <- elaborate' val
         modify (name :)

lib/Eval.hs

@@ -41,8 +41,8 @@ subst _ _ (Free s) = Free s
 subst _ _ (Axiom s) = Axiom s
 subst _ _ (Level i) = Level i
 subst k s (App m n) = App (subst k s m) (subst k s n)
-subst k s (Abs x m n) = Abs x (subst k s m) (subst (k + 1) (incIndices s) n)
+subst k s (Abs x m a n) = Abs x (subst k s m) a (subst (k + 1) (incIndices s) n)
-subst k s (Pi x m n) = Pi x (subst k s m) (subst (k + 1) (incIndices s) n)
+subst k s (Pi x m a n) = Pi x (subst k s m) a (subst (k + 1) (incIndices s) n)
 subst k s (Let x t v b) = Let x t (subst k s v) (subst (k + 1) (incIndices s) b)
 
 envLookupVal :: Text -> ReaderT Env Result Expr
@@ -53,7 +53,7 @@ envLookupTy n = asks ((_ty <$>) . M.lookup n) >>= maybe (throwError $ UnboundVar
 
 -- reduce until β reducts or let simplifications are impossible in head position
 whnf :: Expr -> ReaderT Env Result Expr
-whnf (App (Abs _ _ v) n) = whnf $ subst 0 n v
+whnf (App (Abs _ _ _ v) n) = whnf $ subst 0 n v
 whnf (App m n) = do
     m' <- whnf m
     if m == m'
@@ -64,10 +64,10 @@ whnf (Let _ _ v b) = whnf $ subst 0 v b
 whnf e = pure e
 
 reduce :: Expr -> ReaderT Env Result Expr
-reduce (App (Abs _ _ v) n) = pure $ subst 0 n v
+reduce (App (Abs _ _ _ v) n) = pure $ subst 0 n v
 reduce (App m n) = App <$> reduce m <*> reduce n
-reduce (Abs x t v) = Abs x <$> reduce t <*> reduce v
+reduce (Abs x t a v) = Abs x <$> reduce t <*> pure a <*> reduce v
-reduce (Pi x t v) = Pi x <$> reduce t <*> reduce v
+reduce (Pi x t a v) = Pi x <$> reduce t <*> pure a <*> reduce v
 reduce (Free n) = envLookupVal n
 reduce (Let _ _ v b) = pure $ subst 0 v b
 reduce e = pure e
@@ -92,13 +92,16 @@ betaEquiv e1 e2
             (e, Free n) -> envLookupVal n >>= betaEquiv e
             (Axiom n, Axiom m) -> pure $ n == m
             (Level i, Level j) -> pure $ i == j
-            (Abs _ t1 v1, Abs _ t2 v2) -> (&&) <$> betaEquiv t1 t2 <*> betaEquiv v1 v2 -- i want idiom brackets
+            (Abs _ t1 _ v1, Abs _ t2 _ v2) -> (&&) <$> betaEquiv t1 t2 <*> betaEquiv v1 v2 -- i want idiom brackets
-            (Pi _ t1 v1, Pi _ t2 v2) -> (&&) <$> betaEquiv t1 t2 <*> betaEquiv v1 v2
+            (Pi _ t1 _ v1, Pi _ t2 _ v2) -> (&&) <$> betaEquiv t1 t2 <*> betaEquiv v1 v2
             (App m1 n1, App m2 n2) -> (&&) <$> betaEquiv m1 m2 <*> betaEquiv n1 n2
             (Let _ _ v b, e) -> betaEquiv (subst 0 v b) e
             (e, Let _ _ v b) -> betaEquiv (subst 0 v b) e
             _ -> pure False -- remaining cases impossible or false
 
+betaEquiv' :: Expr -> Expr -> Expr -> ReaderT Env Result ()
+betaEquiv' ctxt e1 e2 = unlessM (betaEquiv e1 e2) $ throwError $ NotEquivalent e1 e2 ctxt
+
 checkBeta :: Env -> Expr -> Expr -> Expr -> Result ()
 checkBeta env e1 e2 ctxt = case runReaderT (betaEquiv e1 e2) env of
     Left err -> Left err

lib/Expr.hs

@@ -6,8 +6,8 @@ data Expr where
     Axiom :: Text -> Expr
     Level :: Integer -> Expr
     App :: Expr -> Expr -> Expr
-    Abs :: Text -> Expr -> Expr -> Expr
+    Abs :: Text -> Expr -> Maybe Expr -> Expr -> Expr
-    Pi :: Text -> Expr -> Expr -> Expr
+    Pi :: Text -> Expr -> Maybe Expr -> Expr -> Expr
     Let :: Text -> Maybe Expr -> Expr -> Expr -> Expr
     deriving (Show, Ord)
 
@@ -17,8 +17,8 @@ instance Eq Expr where
     (Axiom s) == (Axiom t) = s == t
     (Level i) == (Level j) = i == j
     (App e1 e2) == (App f1 f2) = e1 == f1 && e2 == f2
-    (Abs _ t1 b1) == (Abs _ t2 b2) = t1 == t2 && b1 == b2
+    (Abs _ _ t1 b1) == (Abs _ _ t2 b2) = t1 == t2 && b1 == b2
-    (Pi _ t1 b1) == (Pi _ t2 b2) = t1 == t2 && b1 == b2
+    (Pi _ _ t1 b1) == (Pi _ _ t2 b2) = t1 == t2 && b1 == b2
     (Let _ _ v1 b1) == (Let _ _ v2 b2) = v1 == v2 && b1 == b2
     _ == _ = False
 
@@ -28,8 +28,8 @@ occursFree _ (Free _) = False
 occursFree _ (Axiom _) = False
 occursFree _ (Level _) = False
 occursFree n (App a b) = on (||) (occursFree n) a b
-occursFree n (Abs _ a b) = occursFree n a || occursFree (n + 1) b
+occursFree n (Abs _  a _ b) = occursFree n a || occursFree (n + 1) b
-occursFree n (Pi _ a b) = occursFree n a || occursFree (n + 1) b
+occursFree n (Pi _  a _ b) = occursFree n a || occursFree (n + 1) b
 occursFree n (Let _ _ v b) = occursFree n v || occursFree (n + 1) b
 
 shiftIndices :: Integer -> Integer -> Expr -> Expr
@@ -40,8 +40,8 @@ shiftIndices _ _ (Free s) = Free s
 shiftIndices _ _ (Axiom s) = Axiom s
 shiftIndices _ _ (Level i) = Level i
 shiftIndices d c (App m n) = App (shiftIndices d c m) (shiftIndices d c n)
-shiftIndices d c (Abs x m n) = Abs x (shiftIndices d c m) (shiftIndices d (c + 1) n)
+shiftIndices d c (Abs x m a n) = Abs x (shiftIndices d c m) a (shiftIndices d (c + 1) n)
-shiftIndices d c (Pi x m n) = Pi x (shiftIndices d c m) (shiftIndices d (c + 1) n)
+shiftIndices d c (Pi x m a n) = Pi x (shiftIndices d c m) a (shiftIndices d (c + 1) n)
 shiftIndices d c (Let x t v b) = Let x t (shiftIndices d c v) (shiftIndices d (c + 1) b)
 
 incIndices :: Expr -> Expr
@@ -57,7 +57,7 @@ type ParamGroup = ([Text], Expr)
 type Binding = (Text, [ParamGroup], Expr)
 
 collectLambdas :: Expr -> ([Param], Expr)
-collectLambdas (Abs x ty body) = ((x, ty) : params, final)
+collectLambdas (Abs x ty _ body) = ((x, ty) : params, final)
   where
     (params, final) = collectLambdas body
 collectLambdas e = ([], e)
@@ -71,18 +71,18 @@ collectLets (Let x _ val body) = ((x, params', val') : bindings, final)
 collectLets e = ([], e)
 
 collectPis :: Expr -> ([Param], Expr)
-collectPis p@(Pi "" _ _) = ([], p)
+collectPis p@(Pi "" _ _ _) = ([], p)
-collectPis (Pi x ty body) = ((x, ty) : params, final)
+collectPis (Pi x ty _ body) = ((x, ty) : params, final)
   where
     (params, final) = collectPis body
 collectPis e = ([], e)
 
 cleanNames :: Expr -> Expr
 cleanNames (App m n) = App (cleanNames m) (cleanNames n)
-cleanNames (Abs x ty body) = Abs x (cleanNames ty) (cleanNames body)
+cleanNames (Abs x ty a body) = Abs x (cleanNames ty) a (cleanNames body)
-cleanNames (Pi x ty body)
+cleanNames (Pi x ty a body)
-    | occursFree 0 body = Pi x (cleanNames ty) (cleanNames body)
+    | occursFree 0 body = Pi x (cleanNames ty) a (cleanNames body)
-    | otherwise = Pi "" ty (cleanNames body)
+    | otherwise = Pi "" ty a (cleanNames body)
 cleanNames e = e
 
 groupParams :: [Param] -> [ParamGroup]
@@ -116,7 +116,7 @@ helper _ (Level i)
 helper k (App e1 e2) = if k > 3 then parenthesize res else res
   where
     res = helper 3 e1 <> " " <> helper 4 e2
-helper k (Pi "" t1 t2) = if k > 2 then parenthesize res else res
+helper k (Pi "" t1 _ t2) = if k > 2 then parenthesize res else res
   where
     res = helper 3 t1 <> " -> " <> helper 2 t2
 helper k e@(Pi{}) = if k > 2 then parenthesize res else res

lib/IR.hs

@@ -12,11 +12,13 @@ data IRExpr
     | Abs
         { absParamName :: Text
         , absParamType :: IRExpr
+        , absAscription :: Maybe IRExpr
         , absBody :: IRExpr
         }
     | Pi
         { piParamName :: Text
         , piParamType :: IRExpr
+        , piAscription :: Maybe IRExpr
         , piBody :: IRExpr
         }
     | Let

lib/Parser.hs

@@ -68,21 +68,29 @@ pSomeParams = lexeme $ concat <$> some pParamGroup
 pManyParams :: Parser [Param]
 pManyParams = lexeme $ concat <$> many pParamGroup
 
+mkAbs :: Maybe IRExpr -> (Text, IRExpr) -> IRExpr -> IRExpr
+mkAbs ascription (param, typ) = Abs param typ ascription
+
+mkPi :: Maybe IRExpr -> (Text, IRExpr) -> IRExpr -> IRExpr
+mkPi ascription (param, typ) = Pi param typ ascription
+
 pLAbs :: Parser IRExpr
 pLAbs = lexeme $ label "λ-abstraction" $ do
     _ <- defChoice $ "λ" :| ["fun"]
     params <- pSomeParams
+    ascription <- optional pAscription
     _ <- defChoice $ "=>" :| ["⇒"]
     body <- pIRExpr
-    pure $ foldr (uncurry Abs) body params
+    pure $ foldr (mkAbs ascription) body params
 
 pPAbs :: Parser IRExpr
 pPAbs = lexeme $ label "Π-abstraction" $ do
     _ <- defChoice $ "∏" :| ["forall", "∀"]
     params <- pSomeParams
+    ascription <- optional pAscription
     eat ","
     body <- pIRExpr
-    pure $ foldr (uncurry Pi) body params
+    pure $ foldr (mkPi ascription) body params
 
 pBinding :: Parser (Text, Maybe IRExpr, IRExpr)
 pBinding = lexeme $ label "binding" $ do
@@ -95,8 +103,8 @@ pBinding = lexeme $ label "binding" $ do
     eat ")"
     pure
         ( ident
-        , flip (foldr (uncurry Pi)) params <$> ascription
+        , flip (foldr (mkPi Nothing)) params <$> ascription
-        , foldr (uncurry Abs) value params
+        , foldr (mkAbs Nothing) value params
         )
 
 pLet :: Parser IRExpr
@@ -115,7 +123,7 @@ pArrow :: Parser IRExpr
 pArrow = lexeme $ label "->" $ do
     a <- pAppTerm
     _ <- defChoice $ "->" :| ["→"]
-    Pi "" a <$> pIRExpr
+    Pi "" a Nothing <$> pIRExpr
 
 pApp :: Parser IRExpr
 pApp = lexeme $ foldl1 App <$> some pTerm
@@ -134,7 +142,7 @@ pAxiom = lexeme $ label "axiom" $ do
     eat "axiom"
     ident <- pIdentifier
     params <- pManyParams
-    ascription <- fmap (flip (foldr (uncurry Pi)) params) pAscription
+    ascription <- fmap (flip (foldr (mkPi Nothing)) params) pAscription
     eat ";"
     pure $ Axiom ident ascription
 
@@ -143,11 +151,11 @@ pDef = lexeme $ label "definition" $ do
     eat "def"
     ident <- pIdentifier
     params <- pManyParams
-    ascription <- fmap (flip (foldr (uncurry Pi)) params) <$> optional pAscription
+    ascription <- fmap (flip (foldr (mkPi Nothing)) params) <$> optional pAscription
     eat ":="
     body <- pIRExpr
     eat ";"
-    pure $ Def ident ascription $ foldr (uncurry Abs) body params
+    pure $ Def ident ascription $ foldr (mkAbs Nothing) body params
 
 pIRDef :: Parser IRDef
 pIRDef = pAxiom <|> pDef