Misc features; read description

- Refactor of how arrow functions are handled
- Inferring types of closure params/return based on hint
- Allow using `5i/5u` shorthand instead of `5i32/5u32`

compiler/ast/nodes.oc

@@ -201,6 +201,7 @@ struct Function {
     //! (Eventually) resolved return type of the function
     return_type: &Type
     body: &AST
+    is_arrow: bool
 
     scope: &Scope
     closure_scope: &Scope

compiler/lexer.oc

@@ -229,7 +229,15 @@ def Lexer::lex_numeric_literal(&this) {
             .inc()
         }
         let len = .i - start
-        let suffix = .source.substring(start, len)
+        let suffix = if {
+            len > 1 => .source.substring(start, len)
+            initial_char == 'i' => "i32"
+            initial_char == 'u' => "u32"
+            else => {
+                .errors.push(Error::new(Span(start_loc, .loc), "Invalid numeric literal suffix"))
+                yield .source.substring(start, len)
+            }
+        }
         token.suffix = Token::from_ident(suffix, Span(start_loc, .loc))
     }
 

compiler/parser.oc

@@ -1554,7 +1554,7 @@ def Parser::add_doc_comment(&this, sym: &Symbol, token: &Token) {
     }
 }
 
-def Parser::parse_function_args(&this, func: &Function, end: TokenType) {
+def Parser::parse_function_args(&this, func: &Function, end: TokenType, need_types: bool = true) {
     let seen_default = false
     while not .token_is_eof_or(end) {
 
@@ -1584,8 +1584,10 @@ def Parser::parse_function_args(&this, func: &Function, end: TokenType) {
             }
         }
         if not type? {
-            .consume(TokenType::Colon)
-            type = .parse_type()
+            if need_types or .token_is(Colon) {
+                .consume(TokenType::Colon)
+                type = .parse_type()
+            }
         }
 
         let default_value = null as &AST
@@ -1610,32 +1612,21 @@ def Parser::parse_function_args(&this, func: &Function, end: TokenType) {
     }
 }
 
-def Parser::parse_function_body(&this, returns: bool): &AST {
+def Parser::parse_function_body(&this, func: &Function): &AST {
     if .token().type != FatArrow {
+        func.is_arrow = false
         return .parse_block()
     }
 
+    func.is_arrow = true
     let arrow = .consume(TokenType::FatArrow)
     if .token_is(TokenType::OpenCurly) {
         .error(Error::new(.token().span, "Expected an expression for an arrow function"))
         return null
     }
 
     let stmt = .parse_expression(end_type: TokenType::Newline)
-    if returns {
-        let ret = AST::new(Return, stmt.span)
-        ret.u.ret.expr = stmt
-        ret.u.ret.return_span = arrow.span
-        stmt = ret
-    }
-
-    let body = AST::new(Block, stmt.span)
-
-    let statements = Vector<&AST>::new()
-    statements.push(stmt)
-    body.u.block.statements = statements
-
-    return body
+    return stmt
 }
 
 def Parser::parse_closure(&this): &Function {
@@ -1644,20 +1635,18 @@ def Parser::parse_closure(&this): &Function {
     func.name_ast = null
     func.kind = Closure
 
-    .parse_function_args(func, Line)
+    .parse_function_args(func, Line, need_types: false)
     let end_span = .consume(Line).span
 
-    let returns = false
     if .consume_if(TokenType::Colon) {
         func.return_type = .parse_type()
         end_span = func.return_type.span
-        returns = true
     } else {
-        func.return_type = Type::new_unresolved_base(BaseType::Void, start.span)
+        func.return_type = null  // Leave this as null so we can infer it later
     }
 
     func.parsed_return_type = func.return_type
-    func.body = .parse_function_body(returns)
+    func.body = .parse_function_body(func)
     func.span = start.span.join(func.body.span)
     return func
 }
@@ -1767,7 +1756,7 @@ def Parser::parse_function(&this): &Function {
 
     .curr_func = func
 
-    func.body = .parse_function_body(returns)
+    func.body = .parse_function_body(func)
 
     .curr_func = null
     func.span = start.span.join(func.body.span)

compiler/passes/code_generator.oc

@@ -1424,6 +1424,30 @@ def CodeGenerator::gen_type(&this, type: &Type) {
     .gen_type_and_name(type, name: "")
 }
 
+def CodeGenerator::gen_function_body(&this, func: &Function) {
+    let ret_type = func.return_type
+    if func.is_arrow {
+        .out += "{\n"
+        .indent += 1
+        .gen_indent()
+
+        if ret_type.base != Void {
+            .out += "return "
+            .gen_expression(func.body)
+            .out += ";"
+        } else {
+            .gen_statement(func.body)
+        }
+
+        .indent -= 1
+        .gen_indent()
+        .out += "}"
+
+    } else {
+        .gen_block(func.body)
+    }
+}
+
 def CodeGenerator::gen_function(&this, func: &Function) {
     if func.kind == Method and func.parent_type.base == Structure {
         let struc = func.parent_type.u.struc
@@ -1435,7 +1459,7 @@ def CodeGenerator::gen_function(&this, func: &Function) {
     .gen_debug_info(func.sym.span)
     .gen_function_decl(func)
     .out += " "
-    .gen_block(func.body)
+    .gen_function_body(func)
     .out += "\n\n"
 }
 
@@ -1644,7 +1668,7 @@ def CodeGenerator::gen_closure_func(&this, clos: &Function) {
     )
     .out += " {\n"
     .out <<= `  {ctx_type} *{ctx} = ({ctx_type} *)_{ctx};\n`
-    .gen_block(clos.body)
+    .gen_function_body(clos)
     .out += "\n}\n\n"
 }
 

compiler/passes/typechecker.oc

@@ -1674,6 +1674,42 @@ def TypeChecker::check_expression_helper(&this, node: &AST, hint: &Type): &Type
             clos.closure_scope = .o.scope()
             clos.closed_vars = Map<str, &Symbol>::new()
 
+            let hint_matches = if {
+                not hint? => false
+                hint.base != Closure => false
+                hint.u.func.params.size != clos.params.size => false
+                else => true
+            }
+
+            if hint_matches {
+                let hfunc = hint.u.func
+                for let i = 0; i < clos.params.size; i += 1 {
+                    let param = clos.params[i]
+                    let hint_type = hfunc.params[i].type
+                    let param_type = if {
+                        param.type? => param.type
+                        else => hint_type
+                    }
+                    param.type = param_type
+                }
+
+                let hint_ret_type = hfunc.return_type
+                let clos_ret_type = if {
+                    clos.return_type? => clos.return_type
+                    else => hint_ret_type
+                }
+                clos.return_type = clos_ret_type
+            } else {
+                for param : clos.params.iter() {
+                    if not param.type? {
+                        .error(Error::new(param.sym.span, "Cannot infer type, specify it explicitly"))
+                    }
+                }
+            }
+            if not clos.return_type? {
+                clos.return_type = .get_base_type(BaseType::Void, node.span)
+            }
+
             .check_function_declaration(clos)
             .check_function(clos)
             return clos.type
@@ -2333,7 +2369,24 @@ def TypeChecker::check_function(&this, func: &Function) {
     if func.sym? and func.sym.is_extern then return
 
     .o.push_scope(new_scope)
-    .check_statement(func.body)
+    let ret_type = func.return_type
+    if func.is_arrow and ret_type? and ret_type.base != BaseType::Void {
+        let stmt_type = .check_expression(func.body, hint: ret_type)
+        if not stmt_type? {
+            .error(Error::new(func.body.span, "Arrow function must yield a value"))
+        } else if not stmt_type.eq(ret_type) {
+            .error(Error::new_hint(
+                func.body.span, `Expected return type {ret_type.str()}, but got {stmt_type.str()}`,
+                ret_type.span, `Arrow function has return type {ret_type.str()}`
+            ))
+        } else {
+            func.body.returns = true
+        }
+
+    } else {
+        .check_statement(func.body)
+    }
+
     if not func.body.returns and func.return_type.base != BaseType::Void and not func.sym.full_name.eq("main") {
         .error(Error::new(func.sym.span, "Function does not always return"))
     }
@@ -2528,6 +2581,7 @@ def TypeChecker::check_const_expression(&this, node: &AST, hint: &Type = null):
             }
             yield sym.u.var.type
         }
+        // FIXME: What about inferring integer types based on hint?
         IntLiteral  => .get_base_type(BaseType::I32, node.span)
         FloatLiteral => .get_base_type(BaseType::F32, node.span)
         BoolLiteral => .get_base_type(BaseType::Bool, node.span)

compiler/types.oc

@@ -267,7 +267,7 @@ def Type::str(&this): str => match .base {
             if param.type? {
                 buf += param.type.str()
             } else {
-                buf += "null"
+                buf += "<unknown>"
             }
             if i < .u.func.params.size - 1 {
                 buf += ", "

tests/bad/closures/infer_bad_param_type.oc

@@ -0,0 +1,8 @@
+/// fail: Expected `x` with type @fn(i32, i32): i32, but got @fn(i32, u32): i32
+
+def foo(x: @fn(i32,i32): i32): i32 => x(1, 2)
+
+def main() {
+    let t = foo(|a, b: u32| => a + b as i32 + 3)
+    println(`t={t}`)
+}

tests/bad/closures/infer_bad_return_type.oc

@@ -0,0 +1,8 @@
+/// fail: Expected `x` with type @fn(i32, i32): i32, but got @fn(i32, i32): str
+
+def foo(x: @fn(i32,i32): i32): i32 => x(1, 2)
+
+def main() {
+    let t = foo(|a, b|: str => "hello")
+    println(`t={t}`)
+}

tests/bad/closures/infer_closure_expr_wrong.oc

@@ -0,0 +1,9 @@
+/// fail: Expected return type i32, but got str
+
+def foo(x: @fn(i32,i32): i32): i32 => x(1, 2)
+
+def main() {
+    // Infered type is @fn(i32, i32): i32, but we are returning str
+    let t = foo(|a, b| => "hello")
+    println(`t={t}`)
+}

tests/bad/closures/infer_no_hint.oc

@@ -0,0 +1,9 @@
+/// fail: Cannot infer type, specify it explicitly
+
+def foo(x: @fn(i32,i32): i32): i32 => x(1, 2)
+
+def main() {
+    let x = |a, b| => a + b + 3
+    let t = foo(x)
+    println(`t={t}`)
+}

tests/closure_infer_types.oc

@@ -0,0 +1,12 @@
+/// out: "t=6"
+
+def foo(x: @fn(i32,i32): i32): i32 => x(1, 2)
+
+def main() {
+    let t = foo(|a, b| => a + b + 3)
+    t = foo(|a, b|: i32 => a + b + 3)
+    t = foo(|a: i32, b| => a + b + 3)
+    t = foo(|a: i32, b: i32| => a + b + 3)
+    t = foo(|a, b: i32| => a + b + 3)
+    println(`t={t}`)
+}

tests/typed_int_lits.oc

@@ -11,4 +11,5 @@ def main() {
     let d: i64 = -40i64
     let e: f32 = 50.0f32
     let f: f64 = 60.0f64
+    let g: i32 = 60i
 }