Refactor AST and Type Inference: Rename AssignExpression to AssignStatement; update related handling in Builder and TypeChecker; adjust TypeInfo and WatEmitter for consistency

0xGeorgii · 0xGeorgii · commit 68760e143d4c · 2025-05-04T23:02:44.000+09:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -33,5 +33,5 @@ inf-wasmparser = "0.0.8"
 wat-fmt = "0.0.8"
 wasm-fmt = { path = "./wasm-fmt", version = "0.0.1" }
 tree-sitter = "0.25.3"
-tree-sitter-inference = "0.0.30"
+tree-sitter-inference = "0.0.31"
 anyhow = "1.0.98"
diff --git a/ast/src/builder.rs b/ast/src/builder.rs
@@ -7,7 +7,7 @@ use crate::{
     symbols::SymbolType,
     t_ast::TypedAst,
     types::{
-        ArrayIndexAccessExpression, ArrayLiteral, AssertStatement, AssignExpression, AstNode,
+        ArrayIndexAccessExpression, ArrayLiteral, AssertStatement, AssignStatement, AstNode,
         BinaryExpression, Block, BlockType, BoolLiteral, BreakStatement, ConstantDefinition,
         Definition, EnumDefinition, Expression, ExternalFunctionDefinition, FunctionCallExpression,
         FunctionDefinition, FunctionType, GenericType, Identifier, IfStatement, Literal, Location,
@@ -540,6 +540,9 @@ impl<'a> Builder<'a, InitState> {
 
     fn build_statement(&mut self, parent_id: u32, node: &Node, code: &[u8]) -> Statement {
         match node.kind() {
+            "assign_statement" => {
+                Statement::Assign(self.build_assign_statement(parent_id, node, code))
+            }
             "block" | "forall_block" | "assume_block" | "exists_block" | "unique_block" => {
                 Statement::Block(self.build_block(parent_id, node, code))
             }
@@ -675,9 +678,6 @@ impl<'a> Builder<'a, InitState> {
     fn build_expression(&mut self, parent_id: u32, node: &Node, code: &[u8]) -> Expression {
         let node_kind = node.kind();
         match node_kind {
-            "assign_expression" => {
-                Expression::Assign(self.build_assign_expression(parent_id, node, code), None)
-            }
             "array_index_access_expression" => Expression::ArrayIndexAccess(
                 self.build_array_index_access_expression(parent_id, node, code),
                 None,
@@ -716,20 +716,20 @@ impl<'a> Builder<'a, InitState> {
         }
     }
 
-    fn build_assign_expression(
+    fn build_assign_statement(
         &mut self,
         parent_id: u32,
         node: &Node,
         code: &[u8],
-    ) -> Rc<AssignExpression> {
+    ) -> Rc<AssignStatement> {
         let id = Self::get_node_id();
         let location = Self::get_location(node, code);
         let left = self.build_expression(id, &node.child_by_field_name("left").unwrap(), code);
         let right = self.build_expression(id, &node.child_by_field_name("right").unwrap(), code);
 
-        let node = Rc::new(AssignExpression::new(id, location, left, right));
+        let node = Rc::new(AssignStatement::new(id, location, left, right));
         self.arena.add_node(
-            AstNode::Expression(Expression::Assign(node.clone(), None)),
+            AstNode::Statement(Statement::Assign(node.clone())),
             parent_id,
         );
         node
@@ -1134,6 +1134,7 @@ impl<'a> Builder<'a, InitState> {
         let location = Self::get_location(node, code);
         let mut arguments = None;
         let mut cursor = node.walk();
+        let mut returns = None;
 
         let founded_arguments = node
             .children_by_field_name("argument", &mut cursor)
@@ -1142,7 +1143,9 @@ impl<'a> Builder<'a, InitState> {
         if !founded_arguments.is_empty() {
             arguments = Some(founded_arguments);
         }
-        let returns = self.build_type(id, &node.child_by_field_name("returns").unwrap(), code);
+        if let Some(returns_type_node) = node.child_by_field_name("returns") {
+            returns = Some(self.build_type(id, &returns_type_node, code));
+        }
         let node = Rc::new(FunctionType::new(id, location, arguments, returns));
         self.arena.add_node(
             AstNode::Expression(Expression::Type(Type::Function(node.clone()), None)),
diff --git a/ast/src/symbols.rs b/ast/src/symbols.rs
@@ -62,7 +62,7 @@ impl SymbolTable {
                             0,
                             f.location.clone(),
                             param_types,
-                            return_ty,
+                            Some(return_ty),
                         )));
                         table.map.insert(name.clone(), func_ty);
                     }
diff --git a/ast/src/type_infer.rs b/ast/src/type_infer.rs
@@ -1,6 +1,6 @@
 use anyhow::bail;
 
-use crate::types::{Definition, FunctionDefinition, Identifier, TypeInfo};
+use crate::types::{Definition, FunctionDefinition, Identifier, Statement, TypeInfo};
 #[allow(clippy::all, unused_imports, dead_code)]
 use crate::types::{
     Expression, Literal, Location, OperatorKind, SimpleType, SourceFile, Type, TypeArray,
@@ -340,12 +340,6 @@ impl TypeChecker {
     #[allow(clippy::needless_pass_by_value)]
     fn infer_variables(&mut self, function_definition: Rc<FunctionDefinition>) {
         self.symbol_table.push_scope();
-        // let mut generic_type_param_placeholders: HashMap<String, Option<String>> = HashMap::new();
-        // if let Some(type_parameters) = &function_definition.type_parameters {
-        //     for tp in type_parameters {
-        //         generic_type_param_placeholders.insert(tp.name(), None);
-        //     }
-        // }
         if let Some(arguments) = &function_definition.arguments {
             for argument in arguments {
                 if let Err(err) = self
@@ -356,8 +350,151 @@ impl TypeChecker {
                 }
             }
         }
+        for stmt in &mut function_definition.body.statements() {
+            self.infer_statement(
+                stmt,
+                function_definition.returns.clone(),
+                &function_definition
+                    .type_parameters
+                    .as_ref()
+                    .unwrap_or(&vec![])
+                    .iter()
+                    .map(|p| p.name())
+                    .collect(),
+            );
+        }
         self.symbol_table.pop_scope();
     }
+
+    fn infer_statement(
+        &mut self,
+        statement: &mut Statement,
+        return_type: Option<Type>,
+        type_parameters: &Vec<String>,
+    ) {
+        match statement {
+            Statement::Block(block_type) => {
+                self.symbol_table.push_scope();
+                for stmt in &mut block_type.statements() {
+                    self.infer_statement(stmt, return_type.clone(), type_parameters);
+                }
+                self.symbol_table.pop_scope();
+            }
+            Statement::Expression(expression) => match expression {
+                Expression::Assign(assign_expression, type_info) => {
+                    let target_type = self.infer_expression(&mut assign_expression.left);
+                    let value_type = self.infer_expression(&mut assign_expression.right);
+                    if let (Some(target_type), Some(value_type)) = (target_type, value_type) {
+                        if !Self::types_equal(&target_type, &value_type) {
+                            self.errors.push(format!(
+                                "Cannot assign value of type {:?} to variable of type {:?}",
+                                value_type, target_type
+                            ));
+                        }
+                    }
+                }
+            },
+            Statement::Return(return_statement) => todo!(),
+            Statement::Loop(loop_statement) => todo!(),
+            Statement::Break(break_statement) => todo!(),
+            Statement::If(if_statement) => todo!(),
+            Statement::VariableDefinition(variable_definition_statement) => {
+                if let Some(mut initial_value) = variable_definition_statement.value.clone() {
+                    if let Some(init_type) = self.infer_expression(&mut initial_value) {
+                        if !Self::types_equal(&init_type, &variable_definition_statement.ty) {
+                            self.errors.push(format!(
+                                "Type mismatch in variable definition: expected {:?}, found {:?}",
+                                variable_definition_statement.ty, init_type
+                            ));
+                        }
+                    }
+                }
+                if let Err(err) = self.symbol_table.push_variable_to_scope(
+                    variable_definition_statement.name(),
+                    variable_definition_statement.ty.clone(),
+                ) {
+                    self.errors.push(err.to_string());
+                }
+                //TODO handle the case when the variable is not initialized
+            }
+            Statement::TypeDefinition(type_definition_statement) => todo!(),
+            Statement::Assert(assert_statement) => todo!(),
+            Statement::ConstantDefinition(constant_definition) => todo!(),
+        }
+    }
+
+    fn infer_expression(&mut self, expression: &mut Expression) -> Option<Type> {
+        match expression {
+            Expression::ArrayIndexAccess(array_index_access_expression, type_info) => todo!(),
+            Expression::MemberAccess(member_access_expression, type_info) => todo!(),
+            Expression::FunctionCall(function_call_expression, type_info) => todo!(),
+            Expression::PrefixUnary(prefix_unary_expression, type_info) => todo!(),
+            Expression::Parenthesized(parenthesized_expression, type_info) => {
+                self.infer_expression(expression)
+            }
+            Expression::Binary(binary_expression, type_info) => todo!(),
+            Expression::Literal(literal, type_info) => todo!(),
+            Expression::Identifier(identifier, type_info) => todo!(),
+            Expression::Type(_, type_info) => todo!(),
+            Expression::Uzumaki(uzumaki_expression, type_info) => todo!(),
+        }
+    }
+
+    fn types_equal(left: &Type, right: &Type) -> bool {
+        match (left, right) {
+            (Type::Array(left), Type::Array(right)) => {
+                Self::types_equal(&left.element_type, &right.element_type)
+            }
+            (Type::Simple(left), Type::Simple(right)) => left.name == right.name,
+            (Type::Generic(left), Type::Generic(right)) => {
+                left.base.name() == right.base.name() && left.parameters == right.parameters
+            }
+            (Type::Qualified(left), Type::Qualified(right)) => left.name() == right.name(),
+            (Type::QualifiedName(left), Type::QualifiedName(right)) => {
+                left.qualifier() == right.qualifier() && left.name() == right.name()
+            }
+            (Type::Custom(left), Type::Custom(right)) => left.name() == right.name(),
+            (Type::Function(left), Type::Function(right)) => {
+                let left_has_return_type = left.returns.is_some();
+                let right_has_return_type = right.returns.is_some();
+                if left_has_return_type != right_has_return_type {
+                    return false;
+                }
+                if left_has_return_type {
+                    if let (Some(left_return_type), Some(right_return_type)) =
+                        (&left.returns, &right.returns)
+                    {
+                        if !Self::types_equal(left_return_type, right_return_type) {
+                            return false;
+                        }
+                    }
+                }
+                let left_has_parameters = left.parameters.is_some();
+                let right_has_parameters = right.parameters.is_some();
+                if left_has_parameters != right_has_parameters {
+                    return false;
+                }
+                if left_has_parameters {
+                    if let (Some(left_parameters), Some(right_parameters)) =
+                        (&left.parameters, &right.parameters)
+                    {
+                        if left_parameters.len() != right_parameters.len() {
+                            return false;
+                        }
+                        for (left_param, right_param) in
+                            left_parameters.iter().zip(right_parameters.iter())
+                        {
+                            if !Self::types_equal(left_param, right_param) {
+                                return false;
+                            }
+                        }
+                    }
+                }
+                return true;
+            }
+            _ => false,
+        }
+    }
 }
 
 // pub struct TypeContext<'a> {
@@ -482,44 +619,6 @@ impl TypeChecker {
 //     }
 // }
 
-// pub fn traverse_source_files(
-//     source_files: &[crate::types::SourceFile],
-//     symbols: &SymbolTable,
-// ) -> Result<(), TypeError> {
-//     let ctx = TypeContext { symbols };
-//     for sf in source_files {
-//         for def in &sf.definitions {
-//             if let crate::types::Definition::Function(func_rc) = def {
-//                 traverse_function(func_rc, &ctx)?;
-//             }
-//         }
-//     }
-//     Ok(())
-// }
-
-// fn traverse_function(
-//     func_rc: &std::rc::Rc<crate::types::FunctionDefinition>,
-//     ctx: &TypeContext,
-// ) -> Result<(), TypeError> {
-//     let func = func_rc.as_ref();
-//     // TODO: insert parameter types into ctx.symbols if needed
-//     traverse_block(&func.body, ctx)
-// }
-
-// fn traverse_block(
-//     block_type: &crate::types::BlockType,
-//     ctx: &TypeContext,
-// ) -> Result<(), TypeError> {
-//     use crate::types::BlockType;
-//     if let BlockType::Block(b_rc) = block_type {
-//         let block = b_rc.as_ref();
-//         for stmt in &block.statements {
-//             traverse_statement(stmt, ctx)?;
-//         }
-//     }
-//     Ok(())
-// }
-
 // fn traverse_statement(stmt: &crate::types::Statement, ctx: &TypeContext) -> Result<(), TypeError> {
 //     use crate::types::Statement;
 //     match stmt {
diff --git a/ast/src/types.rs b/ast/src/types.rs
@@ -47,8 +47,8 @@ pub struct TypeInfo {
 
 impl TypeInfo {
     #[must_use]
-    pub fn new(ty: Type) -> Self {
-        match &ty {
+    pub fn new(ty: &Type) -> Self {
+        match ty {
             Type::Simple(simple) => Self {
                 name: simple.name.clone(),
                 type_params: vec![],
@@ -66,22 +66,17 @@ impl TypeInfo {
                 type_params: vec![],
             },
             Type::Array(array) => Self {
-                name: format!("Array<{}>", TypeInfo::new(array.element_type.clone()).name),
+                name: format!("Array<{}>", TypeInfo::new(&array.element_type).name),
                 type_params: vec![],
             },
             Type::Function(func) => {
                 //REVISIT
                 let param_types = func
                     .parameters
                     .as_ref()
-                    .map(|params| {
-                        params
-                            .iter()
-                            .map(|p| TypeInfo::new(p.clone()))
-                            .collect::<Vec<_>>()
-                    })
+                    .map(|params| params.iter().map(TypeInfo::new).collect::<Vec<_>>())
                     .unwrap_or_default();
-                let return_type = TypeInfo::new(func.returns.clone());
+                let return_type = TypeInfo::new(&func.returns);
                 Self {
                     name: format!("Function<{}, {}>", param_types.len(), return_type.name),
                     type_params: vec![],
@@ -273,6 +268,7 @@ ast_enums! {
     pub enum Statement {
         @inner_enum Block(BlockType),
         @inner_enum Expression(Expression),
+        Assign(Rc<AssignStatement>),
         Return(Rc<ReturnStatement>),
         Loop(Rc<LoopStatement>),
         Break(Rc<BreakStatement>),
@@ -284,7 +280,6 @@ ast_enums! {
     }
 
     pub enum Expression {
-        Assign(Rc<AssignExpression>, Option<TypeInfo>),
         ArrayIndexAccess(Rc<ArrayIndexAccessExpression>, Option<TypeInfo>),
         MemberAccess(Rc<MemberAccessExpression>, Option<TypeInfo>),
         FunctionCall(Rc<FunctionCallExpression>, Option<TypeInfo>),
@@ -444,7 +439,7 @@ ast_nodes! {
 
     pub struct VariableDefinitionStatement {
         pub name: Rc<Identifier>,
-        pub type_: Type,
+        pub ty: Type,
         pub value: Option<Expression>,
         pub is_uzumaki: bool,
     }
@@ -454,7 +449,7 @@ ast_nodes! {
         pub ty: Type,
     }
 
-    pub struct AssignExpression {
+    pub struct AssignStatement {
         pub left: Expression,
         pub right: Expression,
     }
@@ -524,7 +519,7 @@ ast_nodes! {
 
     pub struct FunctionType {
         pub parameters: Option<Vec<Type>>,
-        pub returns: Type,
+        pub returns: Option<Type>,
     }
 
     pub struct QualifiedName {
diff --git a/ast/src/types_impl.rs b/ast/src/types_impl.rs
diff --git a/wat-codegen/src/wat_emitter.rs b/wat-codegen/src/wat_emitter.rs

Original file line number	Diff line number	Diff line change
`@@ -62,7 +62,7 @@ impl SymbolTable {`
`62`	`62`	`0,`
`63`	`63`	`f.location.clone(),`
`64`	`64`	`param_types,`
`65`		`- return_ty,`
	`65`	`+ Some(return_ty),`
`66`	`66`	`)));`
`67`	`67`	`table.map.insert(name.clone(), func_ty);`
`68`	`68`	`}`