Add support for ForLoop and If expressions; enhance expression handling in AST

Author: 0xGeorgii

sdk/src/ast/build.rs

@@ -12,8 +12,8 @@ use super::{
     contract::Contract,
     custom_type::{CustomType, EnumType, StructType, Type},
     expression::{
-        Array, Assign, BinEx, Binary, Break, Cast, ConstBlock, EBlock, Expression, FunctionCall,
-        Identifier, MemberAccess, MethodCall, Reference,
+        Array, Assign, BinEx, Binary, Break, Cast, ConstBlock, EBlock, Expression, ForLoop,
+        FunctionCall, Identifier, If, MemberAccess, MethodCall, Reference,
     },
     statement::{Block, Statement},
 };
@@ -254,20 +254,54 @@ pub(crate) fn build_block_statement(
     }))
 }
 
-pub(crate) fn build_eblock_expression(block: &Rc<Block>) -> Expression {
-    let id = Uuid::new_v4().as_u128();
+pub(crate) fn build_eblock_expression(block: &Rc<Block>, id: u128) -> Expression {
     Expression::EBlock(Rc::new(EBlock {
         id,
         location: block.location.clone(),
         block: block.clone(),
     }))
 }
 
-pub(crate) fn build_const_block_expression(block: &Rc<Block>) -> Expression {
-    let id = Uuid::new_v4().as_u128();
+pub(crate) fn build_const_block_expression(block: &Rc<Block>, id: u128) -> Expression {
     Expression::Const(Rc::new(ConstBlock {
         id,
         location: block.location.clone(),
         block: block.clone(),
     }))
 }
+
+pub(crate) fn build_for_loop_expression(
+    codebase: &mut Codebase<OpenState>,
+    for_loop: &syn::ExprForLoop,
+    block: &Rc<Block>,
+    id: u128,
+) -> Expression {
+    let expression = codebase.build_expression(&for_loop.expr, id, false);
+    Expression::ForLoop(Rc::new(ForLoop {
+        id,
+        location: location!(for_loop),
+        expression,
+        block: block.clone(),
+    }))
+}
+
+pub(crate) fn build_if_expression(
+    codebase: &mut Codebase<OpenState>,
+    if_expr: &syn::ExprIf,
+    then_branch: Rc<Block>,
+    id: u128,
+) -> Expression {
+    let condition = codebase.build_expression(&if_expr.cond, id, false);
+    let else_branch = if let Some((_, else_expr)) = &if_expr.else_branch {
+        Some(codebase.build_expression(else_expr, id, false))
+    } else {
+        None
+    };
+    Expression::If(Rc::new(If {
+        id,
+        location: location!(if_expr),
+        condition,
+        then_branch,
+        else_branch,
+    }))
+}

sdk/src/ast/expression.rs

@@ -20,7 +20,9 @@ pub enum Expression {
     Closure(Rc<Closure>),
     Const(Rc<ConstBlock>),
     Continue(Rc<Continue>),
+    ForLoop(Rc<ForLoop>),
     FunctionCall(Rc<FunctionCall>),
+    If(Rc<If>),
     MethodCall(Rc<MethodCall>),
     MemberAccess(Rc<MemberAccess>),
     Reference(Rc<Reference>),
@@ -41,7 +43,9 @@ impl Expression {
             Expression::Closure(c) => c.id,
             Expression::Const(c) => c.id,
             Expression::Continue(c) => c.id,
+            Expression::ForLoop(f) => f.id,
             Expression::FunctionCall(f) => f.id,
+            Expression::If(i) => i.id,
             Expression::MethodCall(m) => m.id,
             Expression::MemberAccess(m) => m.id,
             Expression::Reference(r) => r.id,
@@ -62,7 +66,9 @@ impl Expression {
             Expression::Closure(c) => c.location.clone(),
             Expression::Const(c) => c.location.clone(),
             Expression::Continue(c) => c.location.clone(),
+            Expression::ForLoop(f) => f.location.clone(),
             Expression::FunctionCall(f) => f.location.clone(),
+            Expression::If(i) => i.location.clone(),
             Expression::MethodCall(m) => m.location.clone(),
             Expression::MemberAccess(m) => m.location.clone(),
             Expression::Reference(r) => r.location.clone(),
@@ -407,6 +413,25 @@ pub struct Continue {
     pub location: Location,
 }
 
+#[node_location]
+#[derive(serde::Serialize, serde::Deserialize)]
+pub struct ForLoop {
+    pub id: u128,
+    pub location: Location,
+    pub expression: Expression,
+    pub block: Rc<Block>,
+}
+
+#[node_location]
+#[derive(serde::Serialize, serde::Deserialize)]
+pub struct If {
+    pub id: u128,
+    pub location: Location,
+    pub condition: Expression,
+    pub then_branch: Rc<Block>,
+    pub else_branch: Option<Expression>,
+}
+
 #[cfg(test)]
 mod function_call_tests {
     use super::*;

sdk/src/codebase.rs

@@ -6,8 +6,9 @@ use crate::build::{
     build_array_expression, build_assign_expresison, build_binary_expression,
     build_block_statement, build_break_expression, build_cast_expression,
     build_const_block_expression, build_contract, build_eblock_expression, build_enum_custom_type,
-    build_identifier, build_member_access_expression, build_method_call_expression,
-    build_reference_expression, build_struct_custom_type,
+    build_for_loop_expression, build_identifier, build_if_expression,
+    build_member_access_expression, build_method_call_expression, build_reference_expression,
+    build_struct_custom_type,
 };
 use crate::custom_type::Type;
 use crate::errors::SDKErr;
@@ -322,10 +323,11 @@ impl Codebase<OpenState> {
                 break_expr
             }
             syn::Expr::Block(block_expr) => {
+                let id = Uuid::new_v4().as_u128();
                 let block_statement = build_block_statement(self, &block_expr.block);
-                self.add_node(NodeKind::Statement(block_statement.clone()), parent_id);
+                self.add_node(NodeKind::Statement(block_statement.clone()), id);
                 let block = match block_statement {
-                    Statement::Block(block) => build_eblock_expression(&block),
+                    Statement::Block(block) => build_eblock_expression(&block, id),
                     _ => Expression::Empty,
                 };
                 self.add_node(
@@ -396,7 +398,7 @@ impl Codebase<OpenState> {
                 let block_statement = build_block_statement(self, &expr_const.block);
                 self.add_node(NodeKind::Statement(block_statement.clone()), id);
                 let const_block = match block_statement {
-                    Statement::Block(block) => build_const_block_expression(&block),
+                    Statement::Block(block) => build_const_block_expression(&block, id),
                     _ => Expression::Empty,
                 };
                 self.add_node(
@@ -416,6 +418,22 @@ impl Codebase<OpenState> {
                 );
                 cont_expr
             }
+            syn::Expr::ForLoop(expr_forloop) => {
+                let id = Uuid::new_v4().as_u128();
+                let block_statement = build_block_statement(self, &expr_forloop.body);
+                self.add_node(NodeKind::Statement(block_statement.clone()), id);
+                let for_loop = match block_statement {
+                    Statement::Block(block) => {
+                        build_for_loop_expression(self, expr_forloop, &block, id)
+                    }
+                    _ => Expression::Empty,
+                };
+                self.add_node(
+                    NodeKind::Statement(Statement::Expression(for_loop.clone())),
+                    parent_id,
+                );
+                for_loop
+            }
             syn::Expr::Field(field_expr) => {
                 let reference = build_member_access_expression(self, field_expr, is_tried);
                 self.add_node(
@@ -424,6 +442,22 @@ impl Codebase<OpenState> {
                 );
                 reference
             }
+            syn::Expr::If(expr_if) => {
+                let id = Uuid::new_v4().as_u128();
+                let then_block = build_block_statement(self, &expr_if.then_branch);
+                self.add_node(NodeKind::Statement(then_block.clone()), id);
+                match then_block {
+                    Statement::Block(block) => {
+                        let if_expr = build_if_expression(self, expr_if, block.clone(), id);
+                        self.add_node(
+                            NodeKind::Statement(Statement::Expression(if_expr.clone())),
+                            parent_id,
+                        );
+                        if_expr
+                    }
+                    _ => Expression::Empty,
+                }
+            }
             syn::Expr::Try(expr_try) => self.build_expression(&expr_try.expr, parent_id, true),
             syn::Expr::MethodCall(method_call) => {
                 let method_call = build_method_call_expression(self, method_call, is_tried);