llmcc_rust/descriptor/
function.rs

1use llmcc_core::context::CompileUnit;
2use llmcc_core::ir::{HirId, HirNode};
3use tree_sitter::Node;
4
5/// Visibility as written in source for a function item.
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub enum FnVisibility {
8    Private,
9    Public,
10    Crate,
11    Restricted(String),
12}
13
14/// Captured argument information (pattern plus optional type).
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct FunctionParameter {
17    pub pattern: String,
18    pub ty: Option<TypeExpr>,
19}
20
21/// Rich metadata extracted for each Rust function encountered in the HIR.
22#[derive(Debug, Clone)]
23pub struct FunctionDescriptor {
24    /// HIR identifier for the function node.
25    pub hir_id: HirId,
26    /// Short name as written in source (no qualification).
27    pub name: String,
28    /// Module/impl/trait context owning the function.
29    // pub owner: SymbolOwner,
30    /// Visibility as declared (`pub`, `pub(crate)`, etc.).
31    pub visibility: FnVisibility,
32    /// Whether the function is marked `async`.
33    pub is_async: bool,
34    /// Whether the function is marked `const`.
35    pub is_const: bool,
36    /// Whether the function is marked `unsafe`.
37    pub is_unsafe: bool,
38    /// Generic parameter clause (`<T: Foo>`), if present.
39    pub generics: Option<String>,
40    /// Optional `where` clause.
41    pub where_clause: Option<String>,
42    /// Ordered function parameters.
43    pub parameters: Vec<FunctionParameter>,
44    /// Parsed return type, if any.
45    pub return_type: Option<TypeExpr>,
46    /// Textual signature up to (but not including) the body.
47    pub signature: String,
48    /// Fully-qualified name derived from the owner chain.
49    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 owner = SymbolOwner::from_ts_node(unit, ts_node);
72        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            // owner,
107            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/// Representation of Rust types encountered in signatures/fields.
153#[derive(Debug, Clone, PartialEq, Eq)]
154pub enum TypeExpr {
155    /// A bare path with optional generic arguments (e.g. `Result<T, E>`).
156    Path {
157        segments: Vec<String>,
158        generics: Vec<TypeExpr>,
159    },
160    /// `&'a mut T` style references.
161    Reference {
162        is_mut: bool,
163        lifetime: Option<String>,
164        inner: Box<TypeExpr>,
165    },
166    /// Tuple types such as `(A, B)`.
167    Tuple(Vec<TypeExpr>),
168    /// `impl Trait` sugar captured verbatim.
169    ImplTrait { bounds: String },
170    /// Fallback for nodes we do not yet model.
171    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}