Add support for custom types; implement struct and enum handling in AST

0xGeorgii · 0xGeorgii · commit 2fe647d9dc61 · 2025-01-22T12:55:57.000+09:00
diff --git a/sdk/src/ast/build.rs b/sdk/src/ast/build.rs
@@ -3,19 +3,74 @@
 use std::{cell::RefCell, rc::Rc};
 
 use quote::ToTokens;
+use syn::{ItemEnum, ItemStruct};
 use uuid::Uuid;
 
 use crate::{location, Codebase, OpenState};
 
 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,
     },
-    inner_type::Type,
     statement::{Block, Statement},
 };
 
+pub(crate) fn build_contract(struct_item: &ItemStruct) -> Rc<Contract> {
+    let mut fields = Vec::new();
+    for field in &struct_item.fields {
+        let field_name = match &field.ident {
+            Some(ident) => ident.to_string(),
+            None => "unnamed".to_string(),
+        };
+        let field_type = Type::T(field.ty.to_token_stream().to_string());
+        fields.push((field_name, field_type));
+    }
+    Rc::new(Contract {
+        id: Uuid::new_v4().as_u128(),
+        location: location!(struct_item),
+        name: Contract::contract_name_from_syn_item(struct_item),
+        fields,
+        children: RefCell::new(Vec::new()),
+    })
+}
+
+pub(crate) fn build_struct_custom_type(struct_item: &ItemStruct) -> CustomType {
+    let mut fields = Vec::new();
+    for field in &struct_item.fields {
+        let field_name = match &field.ident {
+            Some(ident) => ident.to_string(),
+            None => "unnamed".to_string(),
+        };
+        let field_type = Type::T(field.ty.to_token_stream().to_string());
+        fields.push((field_name, field_type));
+    }
+    CustomType::Struct(Rc::new(StructType {
+        id: Uuid::new_v4().as_u128(),
+        location: location!(struct_item),
+        name: struct_item.ident.to_string(),
+        fields,
+        children: RefCell::new(Vec::new()),
+    }))
+}
+
+pub(crate) fn build_enum_custom_type(enum_item: &ItemEnum) -> CustomType {
+    let mut variants = Vec::new();
+    for variant in &enum_item.variants {
+        let variant_name = variant.ident.to_string();
+        variants.push(variant_name);
+    }
+    CustomType::Enum(Rc::new(EnumType {
+        id: Uuid::new_v4().as_u128(),
+        location: location!(enum_item),
+        name: enum_item.ident.to_string(),
+        variants,
+        children: RefCell::new(Vec::new()),
+    }))
+}
+
 pub(crate) fn build_array_expression(
     codebase: &mut Codebase<OpenState>,
     array_expr: &syn::ExprArray,
diff --git a/sdk/src/ast/contract.rs b/sdk/src/ast/contract.rs
@@ -1,4 +1,5 @@
 #![warn(clippy::pedantic)]
+use super::custom_type::Type;
 use super::function::Function;
 use super::node::{Location, Node, TLocation};
 use super::node_type::ContractChildType;
@@ -13,6 +14,7 @@ pub struct Contract {
     pub id: u128,
     pub location: Location,
     pub name: String,
+    pub fields: Vec<(String, Type)>,
     pub children: RefCell<Vec<ContractChildType>>,
 }
 
diff --git a/sdk/src/ast/custom_type.rs b/sdk/src/ast/custom_type.rs
@@ -0,0 +1,57 @@
+use std::{cell::RefCell, rc::Rc};
+
+use soroban_security_rules_macro_lib::node_location;
+
+use super::{
+    node::{Location, TLocation},
+    node_type::CustomTypeChildType,
+};
+
+#[derive(Clone, serde::Serialize, serde::Deserialize)]
+pub enum Type {
+    T(String), //TODO: Implement this
+}
+
+#[derive(Clone, serde::Serialize, serde::Deserialize)]
+pub enum CustomType {
+    Struct(Rc<StructType>),
+    Enum(Rc<EnumType>),
+}
+
+impl CustomType {
+    #[must_use = "This function creates a new instance of CustomType"]
+    pub fn id(&self) -> u128 {
+        match self {
+            CustomType::Struct(struct_type) => struct_type.id,
+            CustomType::Enum(enum_type) => enum_type.id,
+        }
+    }
+
+    #[must_use = "This function creates a new instance of CustomType"]
+    pub fn location(&self) -> Location {
+        match self {
+            CustomType::Struct(struct_type) => struct_type.location.clone(),
+            CustomType::Enum(enum_type) => enum_type.location.clone(),
+        }
+    }
+}
+
+#[node_location]
+#[derive(serde::Serialize, serde::Deserialize)]
+pub struct StructType {
+    pub id: u128,
+    pub location: Location,
+    pub name: String,
+    pub fields: Vec<(String, Type)>,
+    pub children: RefCell<Vec<CustomTypeChildType>>,
+}
+
+#[node_location]
+#[derive(serde::Serialize, serde::Deserialize)]
+pub struct EnumType {
+    pub id: u128,
+    pub location: Location,
+    pub name: String,
+    pub variants: Vec<String>,
+    pub children: RefCell<Vec<CustomTypeChildType>>,
+}
diff --git a/sdk/src/ast/expression.rs b/sdk/src/ast/expression.rs
@@ -1,6 +1,6 @@
 #![warn(clippy::pedantic)]
+use super::custom_type::Type;
 use super::function::Function;
-use super::inner_type::Type;
 use super::node::{Location, Node, TLocation};
 use super::node_type::{FunctionCallChildType, MemberAccessChildType, MethodCallChildType};
 use super::statement::Block;
diff --git a/sdk/src/ast/inner_type.rs b/sdk/src/ast/inner_type.rs
diff --git a/sdk/src/ast/mod.rs b/sdk/src/ast/mod.rs
@@ -1,10 +1,10 @@
 //! Semantic types for the AST.
 pub mod build;
 pub mod contract;
+pub mod custom_type;
 pub mod expression;
 pub mod file;
 pub mod function;
-pub mod inner_type;
 pub mod node;
 pub mod node_type;
 pub mod statement;
diff --git a/sdk/src/ast/node_type.rs b/sdk/src/ast/node_type.rs
@@ -4,6 +4,7 @@ use serde::{Deserialize, Serialize};
 
 use super::{
     contract::Contract,
+    custom_type::CustomType,
     expression::{Expression, ExpressionParentType, FunctionCall, MethodCall},
     file::File,
     function::{FnParameter, Function},
@@ -36,6 +37,7 @@ impl TypeNode {
 pub enum NodeKind {
     File(Rc<File>),
     Contract(RcContract),
+    CustomType(CustomType),
     Function(RcFunction),
     FnParameter(RcFnParameter),
     Statement(Statement),
@@ -46,6 +48,7 @@ pub enum NodeKind {
 #[derive(Clone, serde::Serialize, serde::Deserialize)]
 pub enum FileChildType {
     Contract(RcContract),
+    CustomType(CustomType),
 }
 
 #[derive(Clone, serde::Serialize, serde::Deserialize)]
@@ -110,6 +113,7 @@ pub fn get_node_kind_node_id(node: &NodeKind) -> u128 {
     match node {
         NodeKind::File(f) => f.id,
         NodeKind::Contract(c) => c.id,
+        NodeKind::CustomType(c) => c.id(),
         NodeKind::Function(f) => f.id,
         NodeKind::FnParameter(p) => p.id,
         NodeKind::Statement(s) => s.id(),
@@ -137,6 +141,7 @@ pub fn get_node_location(node: &NodeKind) -> Location {
             end_col: f.source_code.len(),
         },
         NodeKind::Contract(c) => c.location(),
+        NodeKind::CustomType(c) => c.location(),
         NodeKind::Function(f) => f.location(),
         NodeKind::FnParameter(p) => p.location(),
         NodeKind::Statement(s) => s.location(),
diff --git a/sdk/src/codebase.rs b/sdk/src/codebase.rs
@@ -5,14 +5,15 @@ use crate::ast::node_type::NodeKind;
 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_eblock_expression, build_identifier,
-    build_member_access_expression, build_method_call_expression, build_reference_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,
 };
+use crate::custom_type::Type;
 use crate::errors::SDKErr;
 use crate::expression::{Closure, Continue, Expression, Identifier};
 use crate::file::File;
 use crate::function::{FnParameter, Function};
-use crate::inner_type::Type;
 use crate::node_type::{ContractChildType, FileChildType, FunctionChildType, TypeNode};
 use crate::statement::Statement;
 use crate::{location, source_code, NodesStorage};
@@ -114,26 +115,40 @@ impl Codebase<OpenState> {
                             .iter()
                             .any(|attr| attr.path().is_ident("contract"))
                         {
-                            let contract = Contract {
-                                id: Uuid::new_v4().as_u128(),
-                                location: location!(struct_item),
-                                name: Contract::contract_name_from_syn_item(struct_item),
-                                children: RefCell::new(Vec::new()),
-                            };
-                            let rc_contract = Rc::new(contract);
+                            let rc_contract = build_contract(struct_item);
                             rc_file
                                 .children
                                 .borrow_mut()
                                 .push(FileChildType::Contract(rc_contract.clone()));
                             codebase.add_node(NodeKind::Contract(rc_contract.clone()), rc_file.id);
+                        } else if struct_item
+                            .attrs
+                            .iter()
+                            .any(|attr| attr.path().is_ident("contracttype"))
+                        {
+                            let rc_custom_type = build_struct_custom_type(struct_item);
+                            rc_file
+                                .children
+                                .borrow_mut()
+                                .push(FileChildType::CustomType(rc_custom_type.clone()));
+                            codebase
+                                .add_node(NodeKind::CustomType(rc_custom_type.clone()), rc_file.id);
+                        }
+                    }
+                    syn::Item::Enum(enum_item) => {
+                        if enum_item
+                            .attrs
+                            .iter()
+                            .any(|attr| attr.path().is_ident("contracttype"))
+                        {
+                            let rc_custom_type = build_enum_custom_type(enum_item);
+                            rc_file
+                                .children
+                                .borrow_mut()
+                                .push(FileChildType::CustomType(rc_custom_type.clone()));
+                            codebase
+                                .add_node(NodeKind::CustomType(rc_custom_type.clone()), rc_file.id);
                         }
-                        // else if struct_item
-                        //     .attrs
-                        //     .iter()
-                        //     .any(|attr| attr.path().is_ident("contracttype"))
-                        // {
-                        //     items_to_revisit.push(syn::Item::Struct(struct_item.clone()));
-                        // }
                     }
                     syn::Item::Impl(impl_item) => {
                         if impl_item
diff --git a/sdk/src/utils/test.rs b/sdk/src/utils/test.rs
@@ -86,6 +86,7 @@ pub(crate) fn create_mock_contract(id: u128) -> Contract {
         id,
         name: "TestContract".to_string(),
         location: create_mock_location(),
+        fields: vec![],
         children: RefCell::new(vec![]),
     }
 }
@@ -100,6 +101,7 @@ pub(crate) fn create_mock_contract_with_inner_struct(
         id,
         name,
         location,
+        fields: vec![],
         children: RefCell::new(vec![]),
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -86,6 +86,7 @@ pub(crate) fn create_mock_contract(id: u128) -> Contract {`
`86`	`86`	`id,`
`87`	`87`	`name: "TestContract".to_string(),`
`88`	`88`	`location: create_mock_location(),`
	`89`	`+ fields: vec![],`
`89`	`90`	`children: RefCell::new(vec![]),`
`90`	`91`	`}`
`91`	`92`	`}`
`@@ -100,6 +101,7 @@ pub(crate) fn create_mock_contract_with_inner_struct(`
`100`	`101`	`id,`
`101`	`102`	`name,`
`102`	`103`	`location,`
	`104`	`+ fields: vec![],`
`103`	`105`	`children: RefCell::new(vec![]),`
`104`	`106`	`}`
`105`	`107`	`}`