separate typed and untyped ASTs

src/ast/mod.rs

@@ -1,66 +1,67 @@
-pub mod expr;
-pub mod module;
-
 use std::path::Path;
 
-pub use crate::ast::expr::{BinaryOperator, Expr};
-use crate::ast::module::*;
-use crate::typing::Type;
+pub mod typed;
+pub mod untyped;
+
+#[derive(Debug, PartialEq, Clone)]
+pub enum BinaryOperator {
+    Add,
+    Sub,
+    Mul,
+    Div,
+    Modulo,
+    Equal,
+    NotEqual,
+}
+
+#[derive(Debug, PartialEq, Clone)]
+pub enum UnaryOperator {
+}
 
 pub type Identifier = String;
 
-// XXX: Is this enum actually useful? Is 3:30 AM btw
-#[derive(Debug, PartialEq)]
-pub enum Ast {
-    Module(Module),
+#[derive(Debug, PartialEq, Clone)]
+pub struct ModulePath {
+    components: Vec<String>,
 }
 
-#[derive(Debug, PartialEq)]
-pub enum Definition {
-    FunctionDefinition(FunctionDefinition),
-    //StructDefinition(StructDefinition),
+impl std::fmt::Display for ModulePath {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_fmt(format_args!("{}", self.components.join("::")))
+    }
 }
 
-#[derive(Debug, PartialEq)]
-pub struct Location {
-    pub file: Box<Path>,
+impl From<&Path> for ModulePath {
+    fn from(path: &Path) -> Self {
+        let meta = std::fs::metadata(path).unwrap();
+        ModulePath {
+            components: path
+                .components()
+                .map(|component| match component {
+                    std::path::Component::Normal(n) => {
+                        if meta.is_file() {
+                            n.to_str().unwrap().split(".").nth(0).unwrap().to_string()
+                        } else if meta.is_dir() {
+                            n.to_str().unwrap().to_string()
+                        } else {
+                            // XXX: symlinks?
+                            unreachable!()
+                        }
+                    }
+                    _ => unreachable!(),
+                })
+                .collect(),
+        }
+    }
 }
 
-#[derive(Debug, PartialEq)]
-pub struct FunctionDefinition {
-    pub name: Identifier,
-    pub parameters: Vec<Parameter>,
-    pub return_type: Option<Type>,
-    pub body: Box<Block>,
-    pub line_col: (usize, usize),
-}
-
-#[derive(Debug, PartialEq)]
-pub struct Block {
-    pub statements: Vec<Statement>,
-    pub value: Option<Expr>,
-}
-
-#[derive(Debug, PartialEq)]
-pub enum Statement {
-    DeclareStatement(Identifier, Expr),
-    AssignStatement(Identifier, Expr),
-    ReturnStatement(Option<Expr>),
-    CallStatement(Call),
-    UseStatement(ModulePath),
-    IfStatement(Expr, Block),
-    WhileStatement(Box<Expr>, Box<Block>),
-}
-
-#[derive(Debug, PartialEq)]
-pub struct Call {
-    pub callee: Expr,
-    pub args: Vec<Expr>,
-}
-
-#[derive(Debug, PartialEq)]
-pub struct Parameter {
-    pub name: Identifier,
-    pub typ: Type,
+impl From<&str> for ModulePath {
+    fn from(string: &str) -> Self {
+        ModulePath {
+            components: string.split("::").map(|c| c.to_string()).collect(),
+        }
+    }
 }
 
+#[derive(Eq, PartialEq, Debug)]
+pub struct Import(pub String);

src/ast/module.rs

@@ -1,66 +0,0 @@
-use std::path::Path;
-use super::Definition;
-
-#[derive(Debug, PartialEq, Clone)]
-pub struct ModulePath {
-    components: Vec<String>,
-}
-
-impl std::fmt::Display for ModulePath {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.write_fmt(format_args!("{}", self.components.join("::")))
-    }
-}
-
-impl From<&Path> for ModulePath {
-    fn from(path: &Path) -> Self {
-        let meta = std::fs::metadata(path).unwrap();
-        ModulePath {
-            components: path
-                .components()
-                .map(|component| match component {
-                    std::path::Component::Normal(n) => {
-                        if meta.is_file() {
-                            n.to_str().unwrap().split(".").nth(0).unwrap().to_string()
-                        } else if meta.is_dir() {
-                            n.to_str().unwrap().to_string()
-                        } else {
-                            // XXX: symlinks?
-                            unreachable!()
-                        }
-                    }
-                    _ => unreachable!(),
-                })
-                .collect(),
-        }
-    }
-}
-
-impl From<&str> for ModulePath {
-    fn from(string: &str) -> Self {
-        ModulePath {
-            components: string.split("::").map(|c| c.to_string()).collect(),
-        }
-    }
-}
-
-type ImportPath = ModulePath;
-
-#[derive(Debug, PartialEq)]
-pub struct Module {
-    pub file: Option<std::path::PathBuf>,
-    pub path: ModulePath,
-    pub definitions: Vec<Definition>,
-    pub imports: Vec<ImportPath>,
-}
-
-impl Module {
-    pub fn new(path: ModulePath) -> Self {
-        Module {
-            path,
-            file: None,
-            definitions: vec![],
-            imports: vec![],
-        }
-    }
-}

src/ast/typed/expr.rs

@@ -0,0 +1,67 @@
+use crate::ast::{
+    typed::{Block, Call},
+    BinaryOperator, UnaryOperator,
+};
+use crate::typing::Type;
+
+#[derive(Debug, PartialEq)]
+pub enum Expr {
+    BinaryExpression {
+        lhs: Box<Expr>,
+        op: BinaryOperator,
+        rhs: Box<Expr>,
+        typ: Type,
+    },
+    UnaryExpression {
+        op: UnaryOperator,
+        inner: Box<Expr>,
+    },
+    Variable {
+        name: String,
+        typ: Type,
+    },
+    Call {
+        call: Box<Call>,
+        typ: Type,
+    },
+    Block {
+        block: Box<Block>,
+        typ: Type,
+    },
+    /// Last field is either Expr::Block or Expr::IfExpr
+    IfExpr {
+        cond: Box<Expr>,
+        then_body: Box<Block>,
+        else_body: Box<Expr>,
+        typ: Type,
+    },
+    // Literals
+    UnitLiteral,
+    BooleanLiteral(bool),
+    IntegerLiteral(i64),
+    FloatLiteral(f64),
+    StringLiteral(String),
+}
+
+impl Expr {
+    pub fn typ(&self) -> &Type {
+        match self {
+            Expr::BinaryExpression { lhs, op, rhs, typ } => typ,
+            Expr::UnaryExpression { op, inner } => inner.typ(), // XXX: problems will arise here
+            Expr::Variable { name, typ } => typ,
+            Expr::Call { call, typ } => typ,
+            Expr::Block { block, typ } => typ,
+            Expr::IfExpr {
+                cond,
+                then_body,
+                else_body,
+                typ,
+            } => typ,
+            Expr::UnitLiteral => &Type::Unit,
+            Expr::BooleanLiteral(_) => &Type::Bool,
+            Expr::IntegerLiteral(_) => &Type::Int,
+            Expr::FloatLiteral(_) => &Type::Float,
+            Expr::StringLiteral(_) => &Type::Str,
+        }
+    }
+}

src/ast/typed/mod.rs

@@ -0,0 +1,47 @@
+pub mod expr;
+
+use crate::typing::Type;
+use super::{untyped::Parameter, Identifier, Import};
+use expr::Expr;
+
+#[derive(Debug, PartialEq)]
+pub enum Statement {
+    DeclareStatement(Identifier, Box<Expr>),
+    AssignStatement(Identifier, Box<Expr>),
+    ReturnStatement(Option<Expr>),
+    CallStatement(Box<Call>),
+    UseStatement(Box<Import>),
+    IfStatement(Box<Expr>, Box<Block>),
+    WhileStatement(Box<Expr>, Box<Block>),
+}
+
+#[derive(Debug, PartialEq)]
+pub struct Block {
+    pub statements: Vec<Statement>,
+    pub value: Option<Expr>,
+    typ: Type,
+}
+
+impl Block {
+    #[inline]
+    pub fn typ(&self) -> Type {
+        self.typ.clone()
+    }
+}
+
+#[derive(Debug, PartialEq)]
+pub struct FunctionDefinition {
+    pub name: Identifier,
+    pub parameters: Vec<Parameter>,
+    pub return_type: Option<Type>,
+    pub body: Box<Block>,
+    pub line_col: (usize, usize),
+}
+
+#[derive(Debug, PartialEq)]
+pub struct Call {
+    pub callee: Box<Expr>,
+    pub args: Vec<Expr>,
+    pub typ: Type,
+}
+

src/ast/untyped/expr.rs

@@ -1,7 +1,13 @@
+use crate::ast::{
+    untyped::{Block, Call},
+    Identifier,
+};
+
 use crate::ast::*;
 
 #[derive(Debug, PartialEq)]
 pub enum Expr {
+    UnitLiteral,
     BinaryExpression(Box<Expr>, BinaryOperator, Box<Expr>),
     Identifier(Identifier),
     Call(Box<Call>),
@@ -14,14 +20,3 @@ pub enum Expr {
     /// Last field is either Expr::Block or Expr::IfExpr
     IfExpr(Box<Expr>, Box<Block>, Box<Expr>),
 }
-
-#[derive(Debug, PartialEq, Clone)]
-pub enum BinaryOperator {
-    Add,
-    Sub,
-    Mul,
-    Div,
-    Modulo,
-    Equal,
-    NotEqual,
-}

src/ast/untyped/mod.rs

@@ -0,0 +1,61 @@
+pub mod expr;
+pub mod module;
+
+use std::path::Path;
+
+pub use crate::ast::untyped::expr::Expr;
+pub use crate::ast::*;
+// TODO: remove all usage of 'Type' in the untyped ast
+// (for now it is assumed that anything that parses
+// is a Type, but the checking should be done in the typing
+// phase)
+use crate::typing::Type;
+
+#[derive(Debug, PartialEq)]
+pub enum Definition {
+    FunctionDefinition(FunctionDefinition),
+    //StructDefinition(StructDefinition),
+}
+
+#[derive(Debug, PartialEq)]
+pub struct Location {
+    pub file: Box<Path>,
+}
+
+#[derive(Debug, PartialEq)]
+pub struct FunctionDefinition {
+    pub name: Identifier,
+    pub parameters: Vec<Parameter>,
+    pub return_type: Option<Type>,
+    pub body: Box<Block>,
+    pub line_col: (usize, usize),
+}
+
+#[derive(Debug, PartialEq)]
+pub struct Block {
+    pub statements: Vec<Statement>,
+    pub value: Option<Expr>,
+}
+
+#[derive(Debug, PartialEq)]
+pub enum Statement {
+    DeclareStatement(Identifier, Expr),
+    AssignStatement(Identifier, Expr),
+    ReturnStatement(Option<Expr>),
+    CallStatement(Call),
+    UseStatement(Import),
+    IfStatement(Expr, Block),
+    WhileStatement(Box<Expr>, Box<Block>),
+}
+
+#[derive(Debug, PartialEq)]
+pub struct Call {
+    pub callee: Box<Expr>,
+    pub args: Vec<Expr>,
+}
+
+#[derive(Debug, PartialEq)]
+pub struct Parameter {
+    pub name: Identifier,
+    pub typ: Type,
+}

src/ast/untyped/module.rs

@@ -0,0 +1,20 @@
+use super::{Definition, ModulePath, Import};
+
+#[derive(Debug, PartialEq)]
+pub struct Module {
+    pub file: Option<std::path::PathBuf>,
+    pub path: ModulePath,
+    pub definitions: Vec<Definition>,
+    pub imports: Vec<Import>,
+}
+
+impl Module {
+    pub fn new(path: ModulePath) -> Self {
+        Module {
+            path,
+            file: None,
+            definitions: vec![],
+            imports: vec![],
+        }
+    }
+}

src/lib.rs

@@ -1 +1,4 @@
-pub mod ast;
+mod ast;
+mod typing;
+mod jit;
+mod parsing;

src/main.rs

@@ -4,7 +4,7 @@ mod typing;
 
 use clap::{Parser, Subcommand};
 
-use crate::ast::module::Module;
+use crate::ast::untyped::module::Module;
 
 /// Experimental compiler for lila
 #[derive(Parser, Debug)]

src/parsing/backend/pest/mod.rs

@@ -6,13 +6,14 @@ use pest::iterators::Pair;
 use pest::pratt_parser::PrattParser;
 use pest::Parser;
 
-use crate::ast::module::{Module, ModulePath};
-use crate::ast::*;
+use crate::ast::untyped::module::Module;
+use crate::ast::untyped::*;
+use crate::ast::{Import, ModulePath};
 use crate::typing::Type;
 
 #[derive(pest_derive::Parser)]
 #[grammar = "parsing/backend/pest/grammar.pest"]
-struct KrParser;
+struct LilaParser;
 
 use lazy_static;
 lazy_static::lazy_static! {
@@ -38,7 +39,7 @@ pub fn parse_file(path: &Path) -> Result<Module, Error<Rule>> {
 }
 
 pub fn parse_as_module(source: &str, path: ModulePath) -> Result<Module, Error<Rule>> {
-    let mut pairs = KrParser::parse(Rule::source_file, &source)?;
+    let mut pairs = LilaParser::parse(Rule::source_file, &source)?;
 
     assert!(pairs.len() == 1);
     let module = parse_module(pairs.next().unwrap().into_inner().next().unwrap(), path);
@@ -59,7 +60,7 @@ pub fn parse_module(pair: Pair<Rule>, path: ModulePath) -> Module {
                 module.definitions.push(def);
             }
             Rule::use_statement => {
-                let path = parse_import_path(pair.into_inner().next().unwrap());
+                let path = parse_import(pair.into_inner().next().unwrap());
                 module.imports.push(path);
             }
             _ => panic!("unexpected rule in source_file: {:?}", pair.as_rule()),
@@ -112,8 +113,8 @@ fn parse_statement(pair: Pair<Rule>) -> Statement {
             Statement::CallStatement(call)
         }
         Rule::use_statement => {
-            let path = parse_import_path(pair.into_inner().next().unwrap());
-            Statement::UseStatement(path)
+            let import = parse_import(pair.into_inner().next().unwrap());
+            Statement::UseStatement(import)
         }
         Rule::if_statement => {
             let mut pairs = pair.into_inner();
@@ -133,8 +134,8 @@ fn parse_statement(pair: Pair<Rule>) -> Statement {
 
 type ImportPath = ModulePath;
 
-fn parse_import_path(pair: Pair<Rule>) -> ImportPath {
-    ModulePath::from(pair.as_str())
+fn parse_import(pair: Pair<Rule>) -> Import {
+    Import(pair.as_str().to_string())
 }
 
 fn parse_call(pair: Pair<Rule>) -> Call {
@@ -147,7 +148,10 @@ fn parse_call(pair: Pair<Rule>) -> Call {
         .into_inner()
         .map(parse_expression)
         .collect();
-    Call { callee, args }
+    Call {
+        callee: Box::new(callee),
+        args,
+    }
 }
 
 fn parse_expression(pair: Pair<Rule>) -> Expr {

src/parsing/mod.rs

@@ -1,5 +1,4 @@
 mod backend;
-
-pub use self::backend::pest::{parse_file, parse_module};
-
 mod tests;
+
+pub use self::backend::pest::{parse_file, parse_module, parse_as_module};

src/parsing/tests.rs

@@ -2,8 +2,9 @@
 fn test_addition_function() {
     use crate::parsing::backend::pest::parse_as_module;
     use crate::{
-        ast::module::{Module, ModulePath},
-        ast::*,
+        ast::untyped::module::Module,
+        ast::untyped::*,
+        ast::ModulePath,
         typing::Type,
     };
 

src/typing/error.rs

@@ -0,0 +1,92 @@
+use crate::typing::{BinaryOperator, Identifier, ModulePath, Type, TypeContext};
+
+#[derive(Debug)]
+pub struct TypeError {
+    file: Option<std::path::PathBuf>,
+    module: ModulePath,
+    function: Option<String>,
+    kind: TypeErrorKind,
+}
+
+impl std::fmt::Display for TypeError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("Error\n")?;
+        if let Some(path) = &self.file {
+            f.write_fmt(format_args!(" in file {}\n", path.display()))?;
+        }
+        f.write_fmt(format_args!(" in module {}\n", self.module))?;
+        if let Some(name) = &self.function {
+            f.write_fmt(format_args!(" in function {}\n", name))?;
+        }
+        f.write_fmt(format_args!("{:#?}", self.kind))?;
+        Ok(())
+    }
+}
+
+#[derive(Default)]
+pub struct TypeErrorBuilder {
+    file: Option<std::path::PathBuf>,
+    module: Option<ModulePath>,
+    function: Option<String>,
+    kind: Option<TypeErrorKind>,
+}
+
+impl TypeError {
+    pub fn builder() -> TypeErrorBuilder {
+        TypeErrorBuilder::default()
+    }
+}
+
+impl TypeErrorBuilder {
+    pub fn context(mut self, ctx: &TypeContext) -> Self {
+        self.file = ctx.file.clone();
+        self.module = Some(ctx.module.clone());
+        self.function = ctx.function.clone();
+        self
+    }
+
+    pub fn kind(mut self, kind: TypeErrorKind) -> Self {
+        self.kind = Some(kind);
+        self
+    }
+
+    pub fn build(self) -> TypeError {
+        TypeError {
+            file: self.file,
+            module: self.module.expect("TypeError builder is missing module"),
+            function: self.function,
+            kind: self.kind.expect("TypeError builder is missing kind"),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub enum TypeErrorKind {
+    InvalidBinaryOperator {
+        operator: BinaryOperator,
+        lht: Type,
+        rht: Type,
+    },
+    BlockTypeDoesNotMatchFunctionType {
+        block_type: Type,
+        function_type: Type,
+    },
+    ReturnTypeDoesNotMatchFunctionType {
+        function_type: Type,
+        return_type: Type,
+    },
+    UnknownIdentifier {
+        identifier: String,
+    },
+    AssignmentMismatch {
+        lht: Type,
+        rht: Type,
+    },
+    AssignUndeclared,
+    VariableRedeclaration,
+    ReturnStatementsMismatch,
+    UnknownFunctionCalled(Identifier),
+    WrongFunctionArguments,
+    ConditionIsNotBool,
+    IfElseMismatch,
+}

src/typing/mod.rs

@@ -1,9 +1,11 @@
 use std::collections::HashMap;
 
-use crate::ast::{
-    module::{Module, ModulePath},
-    *,
-};
+use crate::ast::untyped::*;
+use crate::ast::untyped::module::Module;
+use crate::ast::ModulePath;
+
+mod error;
+use crate::typing::error::{TypeError, TypeErrorKind};
 
 #[derive(Debug, PartialEq, Clone)]
 pub enum Type {
@@ -25,22 +27,22 @@ impl From<&str> for Type {
     }
 }
 
-impl FunctionDefinition {
+impl untyped::FunctionDefinition {
     fn signature(&self) -> (Vec<Type>, Type) {
-        let return_type = self.return_type.unwrap_or(Type::Unit);
-        let params_types = self.parameters.iter().map(|p| p.typ).collect();
+        let return_type = self.return_type.clone().unwrap_or(Type::Unit);
+        let params_types = self.parameters.iter().map(|p| p.typ.clone()).collect();
         (params_types, return_type)
     }
 }
 
 impl Module {
     pub fn type_check(&self) -> Result<(), TypeError> {
-        let mut ctx = TypeContext::new(self.path);
+        let mut ctx = TypeContext::new(self.path.clone());
         ctx.file = self.file.clone();
 
         // Register all function signatures
         for Definition::FunctionDefinition(func) in &self.definitions {
-            if let Some(previous) = ctx.functions.insert(func.name.clone(), func.signature()) {
+            if let Some(_previous) = ctx.functions.insert(func.name.clone(), func.signature()) {
                 todo!("handle redefinition of function or identical function names across different files");
             }
         }
@@ -50,103 +52,13 @@ impl Module {
         // Type-check the function bodies
         for Definition::FunctionDefinition(func) in &self.definitions {
             func.typ(&mut ctx)?;
+            ctx.variables.clear();
         }
 
         Ok(())
     }
 }
 
-#[derive(Debug)]
-pub struct TypeError {
-    file: Option<std::path::PathBuf>,
-    module: ModulePath,
-    function: Option<String>,
-    kind: TypeErrorKind,
-}
-
-impl std::fmt::Display for TypeError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.write_str("Error\n")?;
-        if let Some(path) = &self.file {
-            f.write_fmt(format_args!(" in file {}\n", path.display()))?;
-        }
-        f.write_fmt(format_args!(" in module {}\n", self.module))?;
-        if let Some(name) = &self.function {
-            f.write_fmt(format_args!(" in function {}\n", name))?;
-        }
-        f.write_fmt(format_args!("{:#?}", self.kind))?;
-        Ok(())
-    }
-}
-
-#[derive(Default)]
-struct TypeErrorBuilder {
-    file: Option<std::path::PathBuf>,
-    module: Option<ModulePath>,
-    function: Option<String>,
-    kind: Option<TypeErrorKind>,
-}
-
-impl TypeError {
-    fn builder() -> TypeErrorBuilder {
-        TypeErrorBuilder::default()
-    }
-}
-
-impl TypeErrorBuilder {
-    fn context(mut self, ctx: &TypeContext) -> Self {
-        self.file = ctx.file.clone();
-        self.module = Some(ctx.module.clone());
-        self.function = ctx.function.clone();
-        self
-    }
-
-    fn kind(mut self, kind: TypeErrorKind) -> Self {
-        self.kind = Some(kind);
-        self
-    }
-
-    fn build(self) -> TypeError {
-        TypeError {
-            file: self.file,
-            module: self.module.expect("TypeError builder is missing module"),
-            function: self.function,
-            kind: self.kind.expect("TypeError builder is missing kind"),
-        }
-    }
-}
-
-#[derive(Debug)]
-pub enum TypeErrorKind {
-    InvalidBinaryOperator {
-        operator: BinaryOperator,
-        lht: Type,
-        rht: Type,
-    },
-    BlockTypeDoesNotMatchFunctionType {
-        block_type: Type,
-        function_type: Type,
-    },
-    ReturnTypeDoesNotMatchFunctionType {
-        function_type: Type,
-        return_type: Type,
-    },
-    UnknownIdentifier {
-        identifier: String,
-    },
-    AssignmentMismatch {
-        lht: Type,
-        rht: Type,
-    },
-    AssignUndeclared,
-    VariableRedeclaration,
-    ReturnStatementsMismatch,
-    UnknownFunctionCalled(Identifier),
-    WrongFunctionArguments,
-    ConditionIsNotBool,
-    IfElseMismatch,
-}
-
 pub struct TypeContext {
     pub file: Option<std::path::PathBuf>,
     pub module: ModulePath,
@@ -157,36 +69,37 @@ pub struct TypeContext {
 
 impl TypeContext {
     pub fn new(path: ModulePath) -> Self {
-        TypeContext {
+        let builtin_functions =
+            HashMap::from([(String::from("println"), (vec![Type::Str], Type::Unit))]);
+
+        Self {
             file: None,
             module: path,
             function: None,
-            functions: Default::default(),
+            functions: builtin_functions,
             variables: Default::default(),
         }
     }
 }
 
 /// Trait for nodes which have a deducible type.
-pub trait Typ {
+pub trait TypeCheck {
     /// Try to resolve the type of the node.
     fn typ(&self, ctx: &mut TypeContext) -> Result<Type, TypeError>;
 }
 
-impl Typ for FunctionDefinition {
+impl TypeCheck for FunctionDefinition {
     fn typ(&self, ctx: &mut TypeContext) -> Result<Type, TypeError> {
-        let func = self;
+        ctx.function = Some(self.name.clone());
 
-        ctx.function = Some(func.name.clone());
-
-        for param in &func.parameters {
+        for param in &self.parameters {
             ctx.variables.insert(param.name.clone(), param.typ.clone());
         }
 
-        let body_type = &func.body.typ(ctx)?;
+        let body_type = &self.body.typ(ctx)?;
 
         // If the return type is not specified, it is unit.
-        let func_return_type = match &func.return_type {
+        let func_return_type = match &self.return_type {
             Some(typ) => typ,
             None => &Type::Unit,
         };
@@ -203,7 +116,7 @@ impl Typ for FunctionDefinition {
         }
 
         // Check coherence with return statements.
-        for statement in &func.body.statements {
+        for statement in &self.body.statements {
             if let Statement::ReturnStatement(value) = statement {
                 let ret_type = match value {
                     Some(expr) => expr.typ(ctx)?,
@@ -225,7 +138,7 @@ impl Typ for FunctionDefinition {
     }
 }
 
-impl Typ for Block {
+impl TypeCheck for Block {
     fn typ(&self, ctx: &mut TypeContext) -> Result<Type, TypeError> {
         let mut return_typ: Option<Type> = None;
 
@@ -320,9 +233,9 @@ impl Typ for Block {
     }
 }
 
-impl Typ for Call {
+impl TypeCheck for Call {
     fn typ(&self, ctx: &mut TypeContext) -> Result<Type, TypeError> {
-        match &self.callee {
+        match &*self.callee {
             Expr::Identifier(ident) => {
                 let signature = match ctx.functions.get(ident) {
                     Some(sgn) => sgn.clone(),
@@ -356,7 +269,7 @@ impl Typ for Call {
     }
 }
 
-impl Typ for Expr {
+impl TypeCheck for Expr {
     fn typ(&self, ctx: &mut TypeContext) -> Result<Type, TypeError> {
         match self {
             Expr::Identifier(identifier) => {
@@ -426,6 +339,7 @@ impl Typ for Expr {
                 }
             },
             Expr::StringLiteral(_) => Ok(Type::Str),
+            Expr::UnitLiteral => Ok(Type::Unit),
             Expr::Call(call) => call.typ(ctx),
             Expr::Block(block) => block.typ(ctx),
             Expr::IfExpr(cond, true_block, else_value) => {
@@ -450,3 +364,115 @@ impl Typ for Expr {
         }
     }
 }
+
+struct Typed<T: TypeCheck> {
+    inner: T,
+    typ: Type,
+}
+
+trait IntoTyped<T: TypeCheck> {
+    fn into_typed(self: Self, ctx: &mut TypeContext) -> Result<Typed<T>, TypeError>;
+}
+
+impl IntoTyped<Block> for Block {
+    fn into_typed(self: Block, ctx: &mut TypeContext) -> Result<Typed<Block>, TypeError> {
+        let mut return_typ: Option<Type> = None;
+
+        // Check declarations and assignments.
+        for statement in &self.statements {
+            match statement {
+                Statement::DeclareStatement(ident, expr) => {
+                    let typ = expr.typ(ctx)?;
+                    if let Some(_typ) = ctx.variables.insert(ident.clone(), typ) {
+                        // XXX: Shadowing? (illegal for now)
+                        return Err(TypeError::builder()
+                            .context(ctx)
+                            .kind(TypeErrorKind::VariableRedeclaration)
+                            .build());
+                    }
+                }
+
+                Statement::AssignStatement(ident, expr) => {
+                    let rhs_typ = expr.typ(ctx)?;
+                    let Some(lhs_typ) = ctx.variables.get(ident) else {
+                        return Err(TypeError::builder()
+                                   .context(ctx)
+                                   .kind(TypeErrorKind::AssignUndeclared)
+                                   .build());
+                    };
+
+                    // Ensure same type on both sides.
+                    if rhs_typ != *lhs_typ {
+                        return Err(TypeError::builder()
+                            .context(ctx)
+                            .kind(TypeErrorKind::AssignmentMismatch {
+                                lht: lhs_typ.clone(),
+                                rht: rhs_typ.clone(),
+                            })
+                            .build());
+                    }
+                }
+
+                Statement::ReturnStatement(maybe_expr) => {
+                    let expr_typ = if let Some(expr) = maybe_expr {
+                        expr.typ(ctx)?
+                    } else {
+                        Type::Unit
+                    };
+                    if let Some(typ) = &return_typ {
+                        if expr_typ != *typ {
+                            return Err(TypeError::builder()
+                                .context(ctx)
+                                .kind(TypeErrorKind::ReturnStatementsMismatch)
+                                .build());
+                        }
+                    } else {
+                        return_typ = Some(expr_typ);
+                    }
+                }
+
+                Statement::CallStatement(call) => {
+                    call.typ(ctx)?;
+                }
+
+                Statement::UseStatement(_path) => {
+                    // TODO: import the signatures (and types)
+                }
+
+                Statement::IfStatement(cond, block) => {
+                    if cond.typ(ctx)? != Type::Bool {
+                        return Err(TypeError::builder()
+                            .context(ctx)
+                            .kind(TypeErrorKind::ConditionIsNotBool)
+                            .build());
+                    }
+                    block.typ(ctx)?;
+                }
+
+                Statement::WhileStatement(cond, block) => {
+                    if cond.typ(ctx)? != Type::Bool {
+                        return Err(TypeError::builder()
+                            .context(ctx)
+                            .kind(TypeErrorKind::ConditionIsNotBool)
+                            .build());
+                    }
+                    block.typ(ctx)?;
+                }
+            }
+        }
+
+        // Check if there is an expression at the end of the block.
+        let typ = if let Some(expr) = &self.value {
+            expr.typ(ctx)?
+        } else {
+            Type::Unit
+        };
+
+        Ok(Typed { inner: self, typ })
+
+        // TODO/FIXME: find a way to return `return_typ` so that the
+        // top-level block (the function) can check if this return type
+        // (and eventually those from other block) matches the type of
+        // the function.
+    }
+}