run tests so far, fix statements so far

compiler/compiler/src/lib.rs

@@ -130,6 +130,8 @@ impl<'a> Compiler<'a> {
     fn compiler_stages(&self) -> Result<SymbolTable<'_>> {
         let symbol_table = CreateSymbolTable::do_pass((&self.ast, self.handler))?;
 
+        TypeChecker::do_pass((&self.ast, &mut symbol_table.clone(), self.handler))?;
+
         Ok(symbol_table)
     }
 

compiler/passes/src/symbol_table/table.rs

@@ -65,7 +65,7 @@ impl<'a> SymbolTable<'a> {
     }
 
     pub fn lookup_variable(&self, symbol: &Symbol) -> Option<&VariableSymbol<'a>> {
-        self.variables.variables.get(symbol)
+        self.variables.lookup_variable(symbol)
     }
 
     pub fn push_variable_scope(&mut self) {

compiler/passes/src/symbol_table/variable_scope.rs

@@ -49,6 +49,16 @@ impl<'a> VariableScope<'a> {
         self.parent = None;
         self.variables.clear();
     }
+
+    pub fn lookup_variable(&self, symbol: &Symbol) -> Option<&VariableSymbol<'a>> {
+        if let Some(var) = self.variables.get(symbol) {
+            Some(var)
+        } else if let Some(parent) = &self.parent {
+            parent.lookup_variable(symbol)
+        } else {
+            None
+        }
+    }
 }
 
 impl<'a> Display for VariableScope<'a> {

compiler/passes/src/type_checker/check.rs

@@ -48,18 +48,49 @@ impl<'a> TypeChecker<'a> {
             _ => {}
         }
     }
+
+    fn compare_expr_type(&self, compare: Result<Option<Type>>, expr: &Expression, span: &Span) {
+        match expr {
+            Expression::Identifier(ident) => {
+                if let Some(var) = self.symbol_table.lookup_variable(&ident.name) {
+                    self.assert_type(compare, var.type_.clone(), span);
+                } else {
+                    self.handler
+                        .emit_err(TypeCheckerError::unknown_returnee(ident.name, span).into());
+                }
+            }
+            expr => self.compare_types(compare, expr.get_type(), span),
+        }
+    }
+
+    fn assert_expr_type(&self, expr: &Expression, expected: Type, span: &Span) {
+        match expr {
+            Expression::Identifier(ident) => {
+                if let Some(var) = self.symbol_table.lookup_variable(&ident.name) {
+                    if var.type_ != &expected {
+                        self.handler
+                            .emit_err(TypeCheckerError::type_should_be(var.type_, expected, span).into());
+                    }
+                } else {
+                    self.handler
+                        .emit_err(TypeCheckerError::unknown_returnee(ident.name, span).into());
+                }
+            }
+            expr => self.assert_type(expr.get_type(), expected, span),
+        }
+    }
 }
 
 impl<'a> ExpressionVisitor<'a> for TypeChecker<'a> {}
 
-// we can safely unwrap all self.parent instances because
-// statements should always have some parent block
 impl<'a> StatementVisitor<'a> for TypeChecker<'a> {
     fn visit_return(&mut self, input: &'a ReturnStatement) -> VisitResult {
+        // we can safely unwrap all self.parent instances because
+        // statements should always have some parent block
         let parent = self.parent.unwrap();
 
         if let Some(func) = self.symbol_table.lookup_fn(&parent) {
-            self.compare_types(func.get_type(), input.get_type(), input.span());
+            self.compare_expr_type(func.get_type(), &input.expression, input.span());
         }
 
         VisitResult::VisitChildren
@@ -103,7 +134,7 @@ impl<'a> StatementVisitor<'a> for TypeChecker<'a> {
                 _ => {}
             }
 
-            self.assert_type(input.value.get_type(), var.type_.clone(), input.span())
+            self.assert_expr_type(&input.value, var.type_.clone(), input.span())
         } else {
             self.handler
                 .emit_err(TypeCheckerError::unknown_assignee(&input.assignee.identifier.name, input.span()).into())
@@ -113,19 +144,27 @@ impl<'a> StatementVisitor<'a> for TypeChecker<'a> {
     }
 
     fn visit_conditional(&mut self, input: &'a ConditionalStatement) -> VisitResult {
-        self.assert_type(input.condition.get_type(), Type::Boolean, input.span());
+        self.assert_expr_type(&input.condition, Type::Boolean, input.span());
 
         VisitResult::VisitChildren
     }
 
     fn visit_iteration(&mut self, input: &'a IterationStatement) -> VisitResult {
-        // TODO: need to change symbol table to some other repr for variables.
-        // self.symbol_table.insert_variable(input.variable.name, input);
+        if let Err(err) = self.symbol_table.insert_variable(
+            input.variable.name,
+            VariableSymbol {
+                type_: &input.type_,
+                span: input.span(),
+                declaration: Declaration::Const,
+            },
+        ) {
+            self.handler.emit_err(err);
+        }
 
         let iter_var_type = input.get_type();
 
-        self.compare_types(iter_var_type.clone(), input.start.get_type(), input.span());
-        self.compare_types(iter_var_type, input.stop.get_type(), input.span());
+        self.compare_expr_type(iter_var_type.clone(), &input.start, input.span());
+        self.compare_expr_type(iter_var_type, &input.stop, input.span());
 
         VisitResult::VisitChildren
     }
@@ -133,7 +172,7 @@ impl<'a> StatementVisitor<'a> for TypeChecker<'a> {
     fn visit_console(&mut self, input: &'a ConsoleStatement) -> VisitResult {
         match &input.function {
             ConsoleFunction::Assert(expr) => {
-                self.assert_type(expr.get_type(), Type::Boolean, expr.span());
+                self.assert_expr_type(expr, Type::Boolean, expr.span());
             }
             ConsoleFunction::Error(_) | ConsoleFunction::Log(_) => {
                 todo!("need to discuss this");

compiler/passes/src/type_checker/checker.rs

@@ -20,13 +20,13 @@ use leo_span::Symbol;
 use crate::SymbolTable;
 
 pub struct TypeChecker<'a> {
-    pub(crate) symbol_table: SymbolTable<'a>,
+    pub(crate) symbol_table: &'a mut SymbolTable<'a>,
     pub(crate) handler: &'a Handler,
     pub(crate) parent: Option<Symbol>,
 }
 
 impl<'a> TypeChecker<'a> {
-    pub fn new(symbol_table: SymbolTable<'a>, handler: &'a Handler) -> Self {
+    pub fn new(symbol_table: &'a mut SymbolTable<'a>, handler: &'a Handler) -> Self {
         Self {
             symbol_table,
             handler,

compiler/passes/src/type_checker/mod.rs

@@ -26,7 +26,7 @@ use leo_ast::{Ast, VisitorDirector};
 use leo_errors::{emitter::Handler, Result};
 
 impl<'a> Pass<'a> for TypeChecker<'a> {
-    type Input = (&'a Ast, SymbolTable<'a>, &'a Handler);
+    type Input = (&'a Ast, &'a mut SymbolTable<'a>, &'a Handler);
     type Output = Result<()>;
 
     fn do_pass((ast, symbol_table, handler): Self::Input) -> Self::Output {

leo/errors/src/errors/type_checker/type_checker_error.rs

@@ -82,4 +82,14 @@ create_messages!(
         ),
         help: None,
     }
+
+    /// For when the user tries to return a unknown variable.
+    @formatted
+    unknown_returnee {
+        args: (var: impl Display),
+        msg: format!(
+            "Unknown returnee `{var}`",
+        ),
+        help: None,
+    }
 );