11from enum import Enum
2- from typing import Dict , List , Optional , Set
2+ from typing import Dict , List , Optional , Set , ForwardRef
33from pydantic import BaseModel , Field
44
55
66class RustVisibility (Enum ):
77 """Represents Rust visibility modifiers."""
8+
89 PUBLIC = "pub"
9- PRIVATE = ""
10+ PRIVATE = ""
1011 CRATE = "pub(crate)"
1112 SUPER = "pub(super)"
1213 IN_PATH = "pub(in path)"
1314
1415
1516class SafetyClassification (Enum ):
1617 """Classifies the safety level of Rust code."""
18+
1719 SAFE = "safe"
1820 UNSAFE = "unsafe"
19- UNSAFE_CONTAINER = "unsafe_container"
21+ UNSAFE_CONTAINER = "unsafe_container"
2022 FFI = "ffi"
2123
2224
2325class UnsafeReason (Enum ):
2426 """Reasons why code might be unsafe."""
27+
2528 RAW_POINTER_DEREF = "raw_pointer_deref"
2629 MUTABLE_STATIC = "mutable_static"
2730 FFI_CALL = "ffi_call"
2831 UNION_FIELD_ACCESS = "union_field_access"
2932 INLINE_ASSEMBLY = "inline_assembly"
3033 UNSAFE_TRAIT_IMPL = "unsafe_trait_impl"
31- CUSTOM = "custom"
34+ CUSTOM = "custom"
3235
3336
3437class UnsafeBlock (BaseModel ):
3538 """Represents an unsafe block within Rust code."""
39+
3640 start_line : int
3741 end_line : int
3842 reasons : List [UnsafeReason ] = Field (default_factory = list )
@@ -42,6 +46,7 @@ class UnsafeBlock(BaseModel):
4246
4347class RustType (BaseModel ):
4448 """Represents a Rust type."""
49+
4550 name : str
4651 is_reference : bool = False
4752 is_mutable : bool = False
@@ -55,13 +60,15 @@ class RustType(BaseModel):
5560
5661class RustAttribute (BaseModel ):
5762 """Represents a Rust attribute."""
63+
5864 name : str
5965 arguments : List [str ] = Field (default_factory = list )
6066 is_inner : bool = False # #![foo] vs #[foo]
6167
6268
6369class SafetyAnalysis (BaseModel ):
6470 """Analyzes and tracks safety-related information."""
71+
6572 classification : SafetyClassification
6673 unsafe_blocks : List [UnsafeBlock ] = Field (default_factory = list )
6774 unsafe_fn_calls : List [str ] = Field (default_factory = list )
@@ -74,6 +81,7 @@ class SafetyAnalysis(BaseModel):
7481
7582class RustCallable (BaseModel ):
7683 """Represents a Rust function or method."""
84+
7785 name : str
7886 visibility : RustVisibility = RustVisibility .PRIVATE
7987 doc_comment : Optional [str ] = None
@@ -97,23 +105,24 @@ class RustCallable(BaseModel):
97105 variable_declarations : List ["RustVariableDeclaration" ] = Field (default_factory = list )
98106 cyclomatic_complexity : Optional [int ] = None
99107 safety_analysis : SafetyAnalysis
100-
108+
101109 def is_fully_safe (self ) -> bool :
102110 """Check if the function is completely safe (no unsafe blocks or calls)."""
103111 return (
104112 self .safety_analysis .classification == SafetyClassification .SAFE
105113 and not self .safety_analysis .unsafe_blocks
106114 and not self .safety_analysis .unsafe_fn_calls
107115 )
108-
116+
109117 def contains_unsafe (self ) -> bool :
110118 """Check if the function contains any unsafe code."""
111119 return (
112- self .safety_analysis .classification in [SafetyClassification .UNSAFE , SafetyClassification .UNSAFE_CONTAINER ]
120+ self .safety_analysis .classification
121+ in [SafetyClassification .UNSAFE , SafetyClassification .UNSAFE_CONTAINER ]
113122 or bool (self .safety_analysis .unsafe_blocks )
114123 or bool (self .safety_analysis .unsafe_fn_calls )
115124 )
116-
125+
117126 def get_unsafe_reasons (self ) -> Set [UnsafeReason ]:
118127 """Get all reasons for unsafe usage in this function."""
119128 reasons = set ()
@@ -124,6 +133,7 @@ def get_unsafe_reasons(self) -> Set[UnsafeReason]:
124133
125134class RustModule (BaseModel ):
126135 """Represents a Rust module."""
136+
127137 name : str
128138 doc_comment : Optional [str ] = None
129139 attributes : List [RustAttribute ] = Field (default_factory = list )
@@ -132,36 +142,47 @@ class RustModule(BaseModel):
132142 functions : Dict [str , RustCallable ] = Field (default_factory = dict )
133143 safe_functions : Dict [str , RustCallable ] = Field (default_factory = dict )
134144 unsafe_functions : Dict [str , RustCallable ] = Field (default_factory = dict )
135- submodules : Dict [str , ' RustModule' ] = Field (default_factory = dict )
145+ submodules : Dict [str , " RustModule" ] = Field (default_factory = dict )
136146 constants : List ["RustVariableDeclaration" ] = Field (default_factory = list )
137147 macros : List [str ] = Field (default_factory = list )
138148 use_declarations : List [str ] = Field (default_factory = list )
139149 extern_crates : List [str ] = Field (default_factory = list )
140150 is_unsafe : bool = False
141151 file_path : Optional [str ] = None
142152 is_mod_rs : bool = False
153+ is_root_module : bool = False
143154
144155 def categorize_functions (self ):
145156 """Categorize functions based on their safety analysis."""
146157 self .safe_functions .clear ()
147158 self .unsafe_functions .clear ()
148-
159+
149160 for name , func in self .functions .items ():
150161 if func .is_fully_safe ():
151162 self .safe_functions [name ] = func
152163 else :
153164 self .unsafe_functions [name ] = func
154165
155166
167+ class RustDependency (BaseModel ):
168+ """Represents a Rust Dependency, Maybe internal or external"""
169+
170+ name : str
171+ is_external : bool = True
172+ crate : Optional [ForwardRef ("RustCrate" )]
173+
174+
156175class RustCrate (BaseModel ):
157176 """Represents a complete Rust crate."""
177+
158178 name : str
159179 version : str
160- root_module : RustModule
161- dependencies : List ["RustDependencyEdge" ] = Field ( default_factory = list )
180+ is_lib : bool = False
181+ modules : List [RustModule ]
162182 edition : str = "2021"
163183 features : List [str ] = Field (default_factory = list )
164-
184+ dependencies : List [RustDependency ] = Field (default_factory = list )
185+
165186 def analyze_safety (self ) -> Dict [str , int ]:
166187 """Analyze safety statistics across the crate."""
167188 stats = {
@@ -171,7 +192,7 @@ def analyze_safety(self) -> Dict[str, int]:
171192 "unsafe_blocks" : 0 ,
172193 "ffi_functions" : 0 ,
173194 }
174-
195+
175196 def analyze_module (module : RustModule ):
176197 for func in module .functions .values ():
177198 stats ["total_functions" ] += 1
@@ -182,26 +203,26 @@ def analyze_module(module: RustModule):
182203 if func .safety_analysis .classification == SafetyClassification .FFI :
183204 stats ["ffi_functions" ] += 1
184205 stats ["unsafe_blocks" ] += len (func .safety_analysis .unsafe_blocks )
185-
206+
186207 for submodule in module .submodules .values ():
187208 analyze_module (submodule )
188-
209+
189210 analyze_module (self .root_module )
190211 return stats
191-
212+
192213 def get_unsafe_functions (self ) -> List [tuple [str , RustCallable ]]:
193214 """Get all unsafe functions in the crate with their module paths."""
194215 unsafe_fns = []
195-
216+
196217 def collect_unsafe (module : RustModule , path : str ):
197218 for name , func in module .functions .items ():
198219 if not func .is_fully_safe ():
199220 full_path = f"{ path } ::{ name } " if path else name
200221 unsafe_fns .append ((full_path , func ))
201-
222+
202223 for submod_name , submod in module .submodules .items ():
203224 new_path = f"{ path } ::{ submod_name } " if path else submod_name
204225 collect_unsafe (submod , new_path )
205-
226+
206227 collect_unsafe (self .root_module , "" )
207- return unsafe_fns
228+ return unsafe_fns
0 commit comments