From b93268085aab14c80a400c299da5d04d2781098e Mon Sep 17 00:00:00 2001
From: Griffin Smith <root@gws.fyi>
Date: Sun, 14 Mar 2021 17:01:25 -0400
Subject: Implement top-level ascription of declarations

---
 ach/Makefile      |  2 +-
 src/ast/mod.rs    |  1 +
 src/parser/mod.rs | 30 +++++++++++++++++++++++++++++-
 src/tc/mod.rs     | 42 +++++++++++++++++++++++++++++-------------
 tests/compile.rs  | 13 +++++++++++++
 5 files changed, 73 insertions(+), 15 deletions(-)

diff --git a/ach/Makefile b/ach/Makefile
index 869a0d0f8a3e..3a8cd2865e87 100644
--- a/ach/Makefile
+++ b/ach/Makefile
@@ -12,4 +12,4 @@ default: simple
 .PHONY: clean
 
 clean:
-	@rm -f *.ll *.o simple
+	@rm -f *.ll *.o simple functions
diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 1884ba69f43c..3a2261aeda23 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -266,6 +266,7 @@ impl<'a> Fun<'a> {
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub enum Decl<'a> {
     Fun { name: Ident<'a>, body: Fun<'a> },
+    Ascription { name: Ident<'a>, type_: Type<'a> },
 }
 
 ////
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 8599ccabfc23..dd7874aff853 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -98,7 +98,20 @@ named!(fun_decl(&str) -> Decl, do_parse!(
         })
 ));
 
+named!(ascription_decl(&str) -> Decl, do_parse!(
+    name: ident
+        >> multispace0
+        >> complete!(char!(':'))
+        >> multispace0
+        >> type_: type_
+        >> (Decl::Ascription {
+            name,
+            type_
+        })
+));
+
 named!(pub decl(&str) -> Decl, alt!(
+    ascription_decl |
     fun_decl
 ));
 
@@ -108,7 +121,7 @@ named!(pub toplevel(&str) -> Vec<Decl>, terminated!(many0!(decl), multispace0));
 mod tests {
     use std::convert::TryInto;
 
-    use crate::ast::{BinaryOperator, Expr, Literal, Type};
+    use crate::ast::{BinaryOperator, Expr, FunctionType, Literal, Type};
 
     use super::*;
     use expr::tests::ident_expr;
@@ -166,4 +179,19 @@ mod tests {
         );
         assert_eq!(res.len(), 3);
     }
+
+    #[test]
+    fn top_level_ascription() {
+        let res = test_parse!(toplevel, "id : fn a -> a");
+        assert_eq!(
+            res,
+            vec![Decl::Ascription {
+                name: "id".try_into().unwrap(),
+                type_: Type::Function(FunctionType {
+                    args: vec![Type::Var("a".try_into().unwrap())],
+                    ret: Box::new(Type::Var("a".try_into().unwrap()))
+                })
+            }]
+        )
+    }
 }
diff --git a/src/tc/mod.rs b/src/tc/mod.rs
index 4c088c885749..559ac993cc9b 100644
--- a/src/tc/mod.rs
+++ b/src/tc/mod.rs
@@ -305,22 +305,38 @@ impl<'ast> Typechecker<'ast> {
         }
     }
 
-    pub(crate) fn tc_decl(&mut self, decl: ast::Decl<'ast>) -> Result<hir::Decl<'ast, Type>> {
+    pub(crate) fn tc_decl(
+        &mut self,
+        decl: ast::Decl<'ast>,
+    ) -> Result<Option<hir::Decl<'ast, Type>>> {
         match decl {
             ast::Decl::Fun { name, body } => {
-                let body = self.tc_expr(ast::Expr::Fun(Box::new(body)))?;
+                let mut expr = ast::Expr::Fun(Box::new(body));
+                if let Some(type_) = self.env.resolve(&name) {
+                    expr = ast::Expr::Ascription {
+                        expr: Box::new(expr),
+                        type_: self.finalize_type(type_.clone())?,
+                    };
+                }
+
+                let body = self.tc_expr(expr)?;
                 let type_ = body.type_().clone();
                 self.env.set(name.clone(), type_);
                 match body {
-                    hir::Expr::Fun { args, body, type_ } => Ok(hir::Decl::Fun {
+                    hir::Expr::Fun { args, body, type_ } => Ok(Some(hir::Decl::Fun {
                         name,
                         args,
                         body,
                         type_,
-                    }),
+                    })),
                     _ => unreachable!(),
                 }
             }
+            ast::Decl::Ascription { name, type_ } => {
+                let type_ = self.type_from_ast_type(type_);
+                self.env.set(name.clone(), type_);
+                Ok(None)
+            }
         }
     }
 
@@ -561,15 +577,15 @@ pub fn typecheck_expr(expr: ast::Expr) -> Result<hir::Expr<ast::Type>> {
 
 pub fn typecheck_toplevel(decls: Vec<ast::Decl>) -> Result<Vec<hir::Decl<ast::Type>>> {
     let mut typechecker = Typechecker::new();
-    decls
-        .into_iter()
-        .map(|decl| {
-            let hir_decl = typechecker.tc_decl(decl)?;
-            let res = typechecker.finalize_decl(hir_decl)?;
-            typechecker.ctx.clear();
-            Ok(res)
-        })
-        .try_collect()
+    let mut res = Vec::with_capacity(decls.len());
+    for decl in decls {
+        if let Some(hir_decl) = typechecker.tc_decl(decl)? {
+            let hir_decl = typechecker.finalize_decl(hir_decl)?;
+            res.push(hir_decl);
+        }
+        typechecker.ctx.clear();
+    }
+    Ok(res)
 }
 
 #[cfg(test)]
diff --git a/tests/compile.rs b/tests/compile.rs
index 177391423c7d..7fa15ad9653e 100644
--- a/tests/compile.rs
+++ b/tests/compile.rs
@@ -8,6 +8,19 @@ const FIXTURES: &[(&str, i32)] = &[("simple", 5), ("functions", 9)];
 fn compile_and_run_files() {
     let ach = root().unwrap().join("ach");
 
+    println!("Running: `make clean`");
+    assert!(
+        Command::new("make")
+            .arg("clean")
+            .current_dir(&ach)
+            .spawn()
+            .unwrap()
+            .wait()
+            .unwrap()
+            .success(),
+        "make clean failed"
+    );
+
     for (fixture, exit_code) in FIXTURES {
         println!(">>> Testing: {}", fixture);
 
-- 
cgit 1.4.1