Typecheck custom binops

examples/custom_binop.reid

@@ -0,0 +1,12 @@
+// Arithmetic, function calls and imports!
+
+impl binop (lhs: u16) + (rhs: u32) -> u32 {
+  return (lhs as u32) + rhs;
+}
+
+fn main() -> u32 {
+  let value = 6;
+  let other = 15;
+
+  return value * other + 7 * -value;
+}

reid/src/ast/mod.rs

@@ -207,6 +207,17 @@ pub enum TopLevelStatement {
     ExternFunction(FunctionSignature),
     FunctionDefinition(FunctionDefinition),
     TypeDefinition(TypeDefinition),
+    BinopDefinition(BinopDefinition),
+}
+
+#[derive(Debug)]
+pub struct BinopDefinition {
+    pub lhs: (String, Type),
+    pub op: BinaryOperator,
+    pub rhs: (String, Type),
+    pub return_ty: Type,
+    pub block: Block,
+    pub signature_range: TokenRange,
 }
 
 #[derive(Debug)]

reid/src/ast/parse.rs

@@ -554,7 +554,11 @@ impl Parse for Block {
             statements.push(statement);
         }
         stream.expect(Token::BraceClose)?;
-        Ok(Block(statements, return_stmt, stream.get_range().unwrap()))
+        Ok(Block(
+            statements,
+            return_stmt,
+            stream.get_range_prev().unwrap(),
+        ))
     }
 }
 
@@ -785,7 +789,46 @@ impl Parse for TopLevelStatement {
                     range,
                 })
             }
+            Some(Token::Impl) => Stmt::BinopDefinition(stream.parse()?),
             _ => Err(stream.expecting_err("import or fn")?)?,
         })
     }
 }
+
+impl Parse for BinopDefinition {
+    fn parse(mut stream: TokenStream) -> Result<Self, Error> {
+        stream.expect(Token::Impl)?;
+        stream.expect(Token::Binop)?;
+
+        stream.expect(Token::ParenOpen)?;
+        let Some(Token::Identifier(lhs_name)) = stream.next() else {
+            return Err(stream.expected_err("lhs name")?);
+        };
+        stream.expect(Token::Colon)?;
+        let lhs_type = stream.parse()?;
+        stream.expect(Token::ParenClose)?;
+
+        let operator = stream.parse()?;
+
+        stream.expect(Token::ParenOpen)?;
+        let Some(Token::Identifier(rhs_name)) = stream.next() else {
+            return Err(stream.expected_err("rhs name")?);
+        };
+        stream.expect(Token::Colon)?;
+        let rhs_type = stream.parse()?;
+        stream.expect(Token::ParenClose)?;
+
+        let signature_range = stream.get_range().unwrap();
+
+        stream.expect(Token::Arrow)?;
+
+        Ok(BinopDefinition {
+            lhs: (lhs_name, lhs_type),
+            op: operator,
+            rhs: (rhs_name, rhs_type),
+            return_ty: stream.parse()?,
+            block: stream.parse()?,
+            signature_range,
+        })
+    }
+}

reid/src/ast/process.rs

@@ -23,6 +23,7 @@ impl ast::Module {
         let mut imports = Vec::new();
         let mut functions = Vec::new();
         let mut typedefs = Vec::new();
+        let mut binops = Vec::new();
 
         use ast::TopLevelStatement::*;
         for stmt in &self.top_level_statements {
@@ -97,12 +98,30 @@ impl ast::Module {
                     };
                     typedefs.push(def);
                 }
+                BinopDefinition(ast::BinopDefinition {
+                    lhs,
+                    op,
+                    rhs,
+                    return_ty,
+                    block,
+                    signature_range,
+                }) => {
+                    binops.push(mir::BinopDefinition {
+                        lhs: (lhs.0.clone(), lhs.1 .0.into_mir(module_id)),
+                        op: op.mir(),
+                        rhs: (rhs.0.clone(), rhs.1 .0.into_mir(module_id)),
+                        return_ty: return_ty.0.into_mir(module_id),
+                        block: block.into_mir(module_id),
+                        meta: signature_range.as_meta(module_id),
+                    });
+                }
             }
         }
 
         mir::Module {
             name: self.name.clone(),
             module_id: module_id,
+            binop_defs: binops,
             imports,
             functions,
             path: self.path.clone(),

reid/src/lexer.rs

@@ -46,8 +46,12 @@ pub enum Token {
     While,
     /// `for`
     For,
-    /// `In`
+    /// `in`
     In,
+    /// `impl`
+    Impl,
+    /// `binop`
+    Binop,
 
     // Symbols
     /// `;`
@@ -145,6 +149,8 @@ impl ToString for Token {
             Token::For => String::from("for"),
             Token::In => String::from("in"),
             Token::While => String::from("while"),
+            Token::Impl => String::from("impl"),
+            Token::Binop => String::from("binop"),
             Token::Semi => String::from(';'),
             Token::Equals => String::from('='),
             Token::Colon => String::from(':'),
@@ -334,6 +340,8 @@ pub fn tokenize<T: Into<String>>(to_tokenize: T) -> Result<Vec<FullToken>, Error
                     "for" => Token::For,
                     "while" => Token::While,
                     "in" => Token::In,
+                    "impl" => Token::Impl,
+                    "binop" => Token::Binop,
                     _ => Token::Identifier(value),
                 };
                 variant

reid/src/mir/fmt.rs

@@ -40,6 +40,9 @@ impl Display for Module {
         for import in &self.imports {
             writeln!(inner_f, "{}", import)?;
         }
+        for binop in &self.binop_defs {
+            writeln!(inner_f, "{}", binop)?;
+        }
         for typedef in &self.typedefs {
             writeln!(inner_f, "{}", typedef)?;
         }
@@ -56,6 +59,17 @@ impl Display for Import {
     }
 }
 
+impl Display for BinopDefinition {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "impl binop ({}: {:#}) {} ({}: {:#}) -> {:#} ",
+            self.lhs.0, self.lhs.1, self.op, self.rhs.0, self.rhs.1, self.return_ty
+        )?;
+        Display::fmt(&self.block, f)
+    }
+}
+
 impl Display for TypeDefinition {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "type {} = ", self.name)?;

reid/src/mir/mod.rs

@@ -365,6 +365,26 @@ pub enum TypeDefinitionKind {
     Struct(StructType),
 }
 
+#[derive(Debug)]
+pub struct BinopDefinition {
+    pub lhs: (String, TypeKind),
+    pub op: BinaryOperator,
+    pub rhs: (String, TypeKind),
+    pub return_ty: TypeKind,
+    pub block: Block,
+    pub meta: Metadata,
+}
+
+impl BinopDefinition {
+    pub fn block_meta(&self) -> Metadata {
+        self.block.meta
+    }
+
+    pub fn signature(&self) -> Metadata {
+        self.meta
+    }
+}
+
 #[derive(Debug)]
 pub struct Module {
     pub name: String,
@@ -372,6 +392,7 @@ pub struct Module {
     pub imports: Vec<Import>,
     pub functions: Vec<FunctionDefinition>,
     pub typedefs: Vec<TypeDefinition>,
+    pub binop_defs: Vec<BinopDefinition>,
     pub path: Option<PathBuf>,
     pub tokens: Vec<FullToken>,
     pub is_main: bool,

reid/src/mir/typecheck.rs

@@ -132,6 +132,11 @@ impl<'t> Pass for TypeCheck<'t> {
             check_typedefs_for_recursion(&defmap, typedef, HashSet::new(), &mut state);
         }
 
+        for binop in &mut module.binop_defs {
+            let res = binop.typecheck(&self.refs, &mut state.inner());
+            state.ok(res, binop.block_meta());
+        }
+
         for function in &mut module.functions {
             let res = function.typecheck(&self.refs, &mut state.inner());
             state.ok(res, function.block_meta());
@@ -170,6 +175,48 @@ fn check_typedefs_for_recursion<'a, 'b>(
     }
 }
 
+impl BinopDefinition {
+    fn typecheck(
+        &mut self,
+        typerefs: &TypeRefs,
+        state: &mut TypecheckPassState,
+    ) -> Result<TypeKind, ErrorKind> {
+        for param in vec![&self.lhs, &self.rhs] {
+            let param_t = state.or_else(
+                param.1.assert_known(typerefs, state),
+                TypeKind::Vague(Vague::Unknown),
+                self.signature(),
+            );
+            let res = state
+                .scope
+                .variables
+                .set(
+                    param.0.clone(),
+                    ScopeVariable {
+                        ty: param_t.clone(),
+                        mutable: param_t.is_mutable(),
+                    },
+                )
+                .or(Err(ErrorKind::VariableAlreadyDefined(param.0.clone())));
+            state.ok(res, self.signature());
+        }
+
+        let return_type = self.return_ty.clone().assert_known(typerefs, state)?;
+
+        state.scope.return_type_hint = Some(self.return_ty.clone());
+        let inferred = self
+            .block
+            .typecheck(&mut state.inner(), &typerefs, Some(&return_type));
+
+        match inferred {
+            Ok(t) => return_type
+                .collapse_into(&t.1)
+                .or(Err(ErrorKind::ReturnTypeMismatch(return_type, t.1))),
+            Err(e) => Ok(state.or_else(Err(e), return_type, self.block_meta())),
+        }
+    }
+}
+
 impl FunctionDefinition {
     fn typecheck(
         &mut self,

reid/src/mir/typeinference.rs

@@ -12,8 +12,8 @@ use super::{
     pass::{Pass, PassResult, PassState},
     typecheck::{ErrorKind, ErrorTypedefKind},
     typerefs::{ScopeTypeRefs, TypeRef, TypeRefs},
-    Block, CustomTypeKey, ExprKind, Expression, FunctionDefinition, FunctionDefinitionKind,
-    IfExpression, Module, ReturnKind, StmtKind,
+    BinopDefinition, Block, CustomTypeKey, ExprKind, Expression, FunctionDefinition,
+    FunctionDefinitionKind, IfExpression, Module, ReturnKind, StmtKind,
     TypeKind::*,
     VagueType::*,
     WhileStatement,
@@ -55,6 +55,11 @@ impl<'t> Pass for TypeInference<'t> {
             }
         }
 
+        for binop in &mut module.binop_defs {
+            let res = binop.infer_types(&self.refs, &mut state.inner());
+            state.ok(res, binop.block_meta());
+        }
+
         for function in &mut module.functions {
             let res = function.infer_types(&self.refs, &mut state.inner());
             state.ok(res, function.block_meta());
@@ -63,6 +68,56 @@ impl<'t> Pass for TypeInference<'t> {
     }
 }
 
+impl BinopDefinition {
+    fn infer_types(
+        &mut self,
+        type_refs: &TypeRefs,
+        state: &mut TypeInferencePassState,
+    ) -> Result<(), ErrorKind> {
+        let scope_hints = ScopeTypeRefs::from(type_refs);
+
+        let lhs_ty = state.or_else(
+            self.lhs.1.assert_unvague(),
+            Vague(Unknown),
+            self.signature(),
+        );
+        state.ok(
+            scope_hints
+                .new_var(self.lhs.0.clone(), false, &lhs_ty)
+                .or(Err(ErrorKind::VariableAlreadyDefined(self.lhs.0.clone()))),
+            self.signature(),
+        );
+
+        let rhs_ty = state.or_else(
+            self.rhs.1.assert_unvague(),
+            Vague(Unknown),
+            self.signature(),
+        );
+
+        state.ok(
+            scope_hints
+                .new_var(self.rhs.0.clone(), false, &rhs_ty)
+                .or(Err(ErrorKind::VariableAlreadyDefined(self.rhs.0.clone()))),
+            self.signature(),
+        );
+
+        state.scope.return_type_hint = Some(self.return_ty.clone());
+        let ret_res = self.block.infer_types(state, &scope_hints);
+        let (_, mut ret_ty) = state.or_else(
+            ret_res,
+            (
+                ReturnKind::Soft,
+                scope_hints.from_type(&Vague(Unknown)).unwrap(),
+            ),
+            self.block_meta(),
+        );
+
+        ret_ty.narrow(&scope_hints.from_type(&self.return_ty).unwrap());
+
+        Ok(())
+    }
+}
+
 impl FunctionDefinition {
     fn infer_types(
         &mut self,