1use llmcc_core::context::CompileUnit;
2use llmcc_core::ir::{HirId, HirNode};
3use tree_sitter::Node;
4
5#[derive(Debug, Clone, PartialEq, Eq)]
7pub enum FnVisibility {
8 Private,
9 Public,
10 Crate,
11 Restricted(String),
12}
13
14#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct FunctionParameter {
17 pub pattern: String,
18 pub ty: Option<TypeExpr>,
19}
20
21#[derive(Debug, Clone)]
23pub struct FunctionDescriptor {
24 pub hir_id: HirId,
26 pub name: String,
28 pub visibility: FnVisibility,
32 pub is_async: bool,
34 pub is_const: bool,
36 pub is_unsafe: bool,
38 pub generics: Option<String>,
40 pub where_clause: Option<String>,
42 pub parameters: Vec<FunctionParameter>,
44 pub return_type: Option<TypeExpr>,
46 pub signature: String,
48 pub fqn: String,
50}
51
52impl FunctionDescriptor {
53 pub fn from_hir<'tcx>(
54 unit: CompileUnit<'tcx>,
55 node: &HirNode<'tcx>,
56 fqn: String,
57 ) -> Option<Self> {
58 let ts_node = match node.inner_ts_node() {
59 ts if ts.kind() == "function_item" => ts,
60 _ => return None,
61 };
62
63 let name_node = ts_node.child_by_field_name("name")?;
64 let name = node_text(unit, name_node);
65 let header_text = unit
66 .file()
67 .get_text(ts_node.start_byte(), name_node.start_byte());
68 let fn_index = header_text.rfind("fn").unwrap_or(header_text.len());
69 let header_clean = clean(&header_text[..fn_index]);
70
71 let visibility = FnVisibility::from_header(&header_clean);
73 let is_async = header_clean
74 .split_whitespace()
75 .any(|token| token == "async");
76 let is_const = header_clean
77 .split_whitespace()
78 .any(|token| token == "const");
79 let is_unsafe = header_clean
80 .split_whitespace()
81 .any(|token| token == "unsafe");
82 let body_start = ts_node
83 .child_by_field_name("body")
84 .map(|body| body.start_byte())
85 .unwrap_or_else(|| ts_node.end_byte());
86 let signature = clean(&unit.file().get_text(ts_node.start_byte(), body_start));
87
88 let generics = ts_node
89 .child_by_field_name("type_parameters")
90 .map(|n| node_text(unit, n));
91 let where_clause = ts_node
92 .child_by_field_name("where_clause")
93 .map(|n| node_text(unit, n))
94 .or_else(|| extract_where_clause(&signature));
95 let parameters = ts_node
96 .child_by_field_name("parameters")
97 .map(|n| parse_parameters(unit, n))
98 .unwrap_or_default();
99 let return_type = ts_node
100 .child_by_field_name("return_type")
101 .map(|n| parse_return_type_node(unit, n));
102
103 Some(FunctionDescriptor {
104 hir_id: node.hir_id(),
105 name,
106 visibility,
108 is_async,
109 is_const,
110 is_unsafe,
111 generics,
112 where_clause,
113 parameters,
114 return_type,
115 signature,
116 fqn,
117 })
118 }
119}
120
121fn extract_where_clause(signature: &str) -> Option<String> {
122 let signature = signature.trim();
123 let idx = signature.find("where ")?;
124 let clause = signature[idx..].trim();
125 if clause.is_empty() {
126 return None;
127 }
128 Some(clause.trim_end_matches(',').to_string())
129}
130
131impl FnVisibility {
132 pub(crate) fn from_header(header: &str) -> Self {
133 if let Some(index) = header.find("pub") {
134 let rest = &header[index..];
135 let compressed: String = rest.chars().filter(|c| !c.is_whitespace()).collect();
136 if compressed.starts_with("pub(") && compressed.ends_with(')') {
137 let inner = &compressed[4..compressed.len() - 1];
138 if inner == "crate" {
139 FnVisibility::Crate
140 } else {
141 FnVisibility::Restricted(inner.to_string())
142 }
143 } else {
144 FnVisibility::Public
145 }
146 } else {
147 FnVisibility::Private
148 }
149 }
150}
151
152#[derive(Debug, Clone, PartialEq, Eq)]
154pub enum TypeExpr {
155 Path {
157 segments: Vec<String>,
158 generics: Vec<TypeExpr>,
159 },
160 Reference {
162 is_mut: bool,
163 lifetime: Option<String>,
164 inner: Box<TypeExpr>,
165 },
166 Tuple(Vec<TypeExpr>),
168 ImplTrait { bounds: String },
170 Unknown(String),
172}
173
174impl TypeExpr {
175 pub fn path_segments(&self) -> Option<&[String]> {
176 match self {
177 TypeExpr::Path { segments, .. } => Some(segments),
178 _ => None,
179 }
180 }
181
182 pub fn generics(&self) -> Option<&[TypeExpr]> {
183 match self {
184 TypeExpr::Path { generics, .. } => Some(generics),
185 _ => None,
186 }
187 }
188}
189
190fn parse_parameters<'tcx>(
191 unit: CompileUnit<'tcx>,
192 params_node: Node<'tcx>,
193) -> Vec<FunctionParameter> {
194 let mut params = Vec::new();
195 let mut cursor = params_node.walk();
196 for child in params_node.named_children(&mut cursor) {
197 match child.kind() {
198 "parameter" => {
199 let pattern = child
200 .child_by_field_name("pattern")
201 .map(|n| node_text(unit, n))
202 .unwrap_or_else(|| node_text(unit, child));
203 let ty = child
204 .child_by_field_name("type")
205 .map(|n| parse_type_expr(unit, n));
206 params.push(FunctionParameter { pattern, ty });
207 }
208 "self_parameter" => {
209 params.push(FunctionParameter {
210 pattern: node_text(unit, child),
211 ty: None,
212 });
213 }
214 _ => {}
215 }
216 }
217 params
218}
219
220pub(crate) fn parse_type_expr<'tcx>(unit: CompileUnit<'tcx>, node: Node<'tcx>) -> TypeExpr {
221 match node.kind() {
222 "type_identifier" | "primitive_type" => TypeExpr::Path {
223 segments: node_text(unit, node)
224 .split("::")
225 .map(|s| s.to_string())
226 .collect(),
227 generics: Vec::new(),
228 },
229 "scoped_type_identifier" => TypeExpr::Path {
230 segments: node_text(unit, node)
231 .split("::")
232 .map(|s| s.to_string())
233 .collect(),
234 generics: Vec::new(),
235 },
236 "generic_type" => parse_generic_type(unit, node),
237 "reference_type" => parse_reference_type(unit, node),
238 "tuple_type" => {
239 let mut types = Vec::new();
240 let mut cursor = node.walk();
241 for child in node.named_children(&mut cursor) {
242 if is_type_node(child.kind()) {
243 types.push(parse_type_expr(unit, child));
244 }
245 }
246 TypeExpr::Tuple(types)
247 }
248 "impl_trait_type" => TypeExpr::ImplTrait {
249 bounds: node_text(unit, node),
250 },
251 _ => TypeExpr::Unknown(node_text(unit, node)),
252 }
253}
254
255fn parse_generic_type<'tcx>(unit: CompileUnit<'tcx>, node: Node<'tcx>) -> TypeExpr {
256 let mut base_segments: Vec<String> = Vec::new();
257 let mut generics = Vec::new();
258 let mut cursor = node.walk();
259 for child in node.named_children(&mut cursor) {
260 match child.kind() {
261 "type_identifier" | "scoped_type_identifier" => {
262 base_segments = node_text(unit, child)
263 .split("::")
264 .map(|s| s.to_string())
265 .collect();
266 }
267 "type_arguments" => {
268 generics = parse_type_arguments(unit, child);
269 }
270 _ => {}
271 }
272 }
273 if base_segments.is_empty() {
274 base_segments = node_text(unit, node)
275 .split("::")
276 .map(|s| s.to_string())
277 .collect();
278 }
279 TypeExpr::Path {
280 segments: base_segments,
281 generics,
282 }
283}
284
285fn parse_type_arguments<'tcx>(unit: CompileUnit<'tcx>, node: Node<'tcx>) -> Vec<TypeExpr> {
286 let mut args = Vec::new();
287 let mut cursor = node.walk();
288 for child in node.named_children(&mut cursor) {
289 match child.kind() {
290 "type_argument" => {
291 if let Some(inner) = child.child_by_field_name("type") {
292 args.push(parse_type_expr(unit, inner));
293 }
294 }
295 kind if is_type_node(kind) => {
296 args.push(parse_type_expr(unit, child));
297 }
298 _ => {}
299 }
300 }
301 args
302}
303
304fn parse_reference_type<'tcx>(unit: CompileUnit<'tcx>, node: Node<'tcx>) -> TypeExpr {
305 let mut lifetime = None;
306 let mut is_mut = false;
307 let mut inner = None;
308 let mut cursor = node.walk();
309 for child in node.named_children(&mut cursor) {
310 match child.kind() {
311 "lifetime" => lifetime = Some(node_text(unit, child)),
312 "mutable_specifier" => is_mut = true,
313 kind if is_type_node(kind) => inner = Some(parse_type_expr(unit, child)),
314 _ => {}
315 }
316 }
317 let inner = inner.unwrap_or_else(|| TypeExpr::Unknown(node_text(unit, node)));
318 TypeExpr::Reference {
319 is_mut,
320 lifetime,
321 inner: Box::new(inner),
322 }
323}
324
325fn is_type_node(kind: &str) -> bool {
326 matches!(
327 kind,
328 "type_identifier"
329 | "scoped_type_identifier"
330 | "generic_type"
331 | "reference_type"
332 | "tuple_type"
333 | "primitive_type"
334 | "impl_trait_type"
335 )
336}
337
338fn parse_return_type_node<'tcx>(unit: CompileUnit<'tcx>, node: Node<'tcx>) -> TypeExpr {
339 let mut expr_opt = None;
340 let mut cursor = node.walk();
341 for child in node.named_children(&mut cursor) {
342 if expr_opt.is_none() && is_type_node(child.kind()) {
343 expr_opt = Some(parse_type_expr(unit, child));
344 }
345 }
346
347 let mut expr = if let Some(expr) = expr_opt {
348 expr
349 } else if let Some(inner) = node.child_by_field_name("type") {
350 parse_type_expr(unit, inner)
351 } else {
352 parse_type_expr(unit, node)
353 };
354
355 let mut cursor = node.walk();
356 for child in node.named_children(&mut cursor) {
357 if child.kind() == "type_arguments" {
358 let extra = parse_type_arguments(unit, child);
359 if !extra.is_empty() {
360 if let TypeExpr::Path { generics, .. } = &mut expr {
361 generics.extend(extra);
362 }
363 }
364 }
365 }
366
367 if let TypeExpr::Unknown(text) = &expr {
368 if text.starts_with("impl ") {
369 expr = TypeExpr::ImplTrait {
370 bounds: text.clone(),
371 };
372 }
373 }
374
375 expr
376}
377
378fn clean(text: &str) -> String {
379 let mut out = String::new();
380 let mut last_was_ws = false;
381 for ch in text.chars() {
382 if ch.is_whitespace() {
383 if !last_was_ws && !out.is_empty() {
384 out.push(' ');
385 }
386 last_was_ws = true;
387 } else {
388 out.push(ch);
389 last_was_ws = false;
390 }
391 }
392 out.trim().to_string()
393}
394
395fn node_text<'tcx>(unit: CompileUnit<'tcx>, node: Node<'tcx>) -> String {
396 clean(&unit.file().get_text(node.start_byte(), node.end_byte()))
397}