llmcc_rust/descriptor/
enumeration.rs

1use llmcc_core::context::CompileUnit;
2use llmcc_core::ir::{HirId, HirNode};
3use tree_sitter::Node;
4
5use super::function::{parse_type_expr, FnVisibility, TypeExpr};
6
7/// Classification of enum variants.
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum EnumVariantKind {
10    Unit,
11    Tuple,
12    Struct,
13}
14
15/// Captures metadata for a single enum variant.
16#[derive(Debug, Clone)]
17pub struct EnumVariant {
18    pub name: String,
19    pub kind: EnumVariantKind,
20    pub fields: Vec<EnumVariantField>,
21    pub discriminant: Option<String>,
22}
23
24/// Field metadata for tuple or struct variants.
25#[derive(Debug, Clone)]
26pub struct EnumVariantField {
27    pub name: Option<String>,
28    pub ty: Option<TypeExpr>,
29}
30
31/// Structured metadata for Rust enums.
32#[derive(Debug, Clone)]
33pub struct EnumDescriptor {
34    pub hir_id: HirId,
35    pub name: String,
36    pub fqn: String,
37    pub visibility: FnVisibility,
38    pub generics: Option<String>,
39    pub variants: Vec<EnumVariant>,
40}
41
42impl EnumDescriptor {
43    pub fn from_enum<'tcx>(
44        unit: CompileUnit<'tcx>,
45        node: &HirNode<'tcx>,
46        fqn: String,
47    ) -> Option<Self> {
48        let ts_node = match node.inner_ts_node() {
49            ts if ts.kind() == "enum_item" => ts,
50            _ => return None,
51        };
52
53        let name_node = ts_node.child_by_field_name("name")?;
54        let name = clean(&node_text(unit, name_node));
55        let header_text = unit
56            .file()
57            .get_text(ts_node.start_byte(), name_node.start_byte());
58        let visibility = FnVisibility::from_header(&header_text);
59
60        let generics = ts_node
61            .child_by_field_name("type_parameters")
62            .map(|n| clean(&node_text(unit, n)));
63
64        let variants = ts_node
65            .child_by_field_name("body")
66            .map(|body| parse_enum_variants(unit, body))
67            .unwrap_or_default();
68
69        Some(EnumDescriptor {
70            hir_id: node.hir_id(),
71            name,
72            fqn,
73            visibility,
74            generics,
75            variants,
76        })
77    }
78}
79
80fn parse_enum_variants<'tcx>(unit: CompileUnit<'tcx>, body: Node<'tcx>) -> Vec<EnumVariant> {
81    let mut variants = Vec::new();
82    let mut cursor = body.walk();
83    for child in body.named_children(&mut cursor) {
84        if child.kind() == "enum_variant" {
85            variants.push(parse_enum_variant(unit, child));
86        }
87    }
88    variants
89}
90
91fn parse_enum_variant<'tcx>(unit: CompileUnit<'tcx>, node: Node<'tcx>) -> EnumVariant {
92    let name_node = node
93        .child_by_field_name("name")
94        .unwrap_or_else(|| node.child(0).unwrap_or(node));
95    let name = clean(&node_text(unit, name_node));
96
97    let discriminant = node
98        .child_by_field_name("value")
99        .map(|n| clean(&node_text(unit, n)));
100
101    let (kind, fields) = match node.child_by_field_name("body") {
102        Some(body) => match body.kind() {
103            "field_declaration_list" => (
104                EnumVariantKind::Struct,
105                parse_named_variant_fields(unit, body),
106            ),
107            "ordered_field_declaration_list" | "tuple_field_declaration_list" => (
108                EnumVariantKind::Tuple,
109                parse_tuple_variant_fields(unit, body),
110            ),
111            other => parse_variant_body(unit, body, other),
112        },
113        None => (EnumVariantKind::Unit, Vec::new()),
114    };
115
116    EnumVariant {
117        name,
118        kind,
119        fields,
120        discriminant,
121    }
122}
123
124fn parse_variant_body<'tcx>(
125    unit: CompileUnit<'tcx>,
126    body: Node<'tcx>,
127    kind: &str,
128) -> (EnumVariantKind, Vec<EnumVariantField>) {
129    match kind {
130        "field_declaration_list" => (
131            EnumVariantKind::Struct,
132            parse_named_variant_fields(unit, body),
133        ),
134        "ordered_field_declaration_list" | "tuple_field_declaration_list" => (
135            EnumVariantKind::Tuple,
136            parse_tuple_variant_fields(unit, body),
137        ),
138        _ => (EnumVariantKind::Unit, Vec::new()),
139    }
140}
141
142fn parse_named_variant_fields<'tcx>(
143    unit: CompileUnit<'tcx>,
144    list: Node<'tcx>,
145) -> Vec<EnumVariantField> {
146    let mut fields = Vec::new();
147    let mut cursor = list.walk();
148    for child in list.named_children(&mut cursor) {
149        if child.kind() == "field_declaration" {
150            let name = child
151                .child_by_field_name("name")
152                .map(|n| clean(&node_text(unit, n)));
153            let ty = child
154                .child_by_field_name("type")
155                .map(|n| parse_type_expr(unit, n));
156            fields.push(EnumVariantField { name, ty });
157        }
158    }
159    fields
160}
161
162fn parse_tuple_variant_fields<'tcx>(
163    unit: CompileUnit<'tcx>,
164    list: Node<'tcx>,
165) -> Vec<EnumVariantField> {
166    let mut fields = Vec::new();
167    let mut cursor = list.walk();
168    for child in list.children(&mut cursor) {
169        match child.kind() {
170            "tuple_field_declaration" | "ordered_field_declaration" => {
171                let ty = child
172                    .child_by_field_name("type")
173                    .map(|n| parse_type_expr(unit, n))
174                    .or_else(|| {
175                        child
176                            .children(&mut child.walk())
177                            .find_map(|n| is_type_node(n.kind()).then(|| parse_type_expr(unit, n)))
178                    });
179                fields.push(EnumVariantField { name: None, ty });
180            }
181            kind if is_type_node(kind) => {
182                fields.push(EnumVariantField {
183                    name: None,
184                    ty: Some(parse_type_expr(unit, child)),
185                });
186            }
187            _ => {}
188        }
189    }
190    fields
191}
192
193fn is_type_node(kind: &str) -> bool {
194    matches!(
195        kind,
196        "type_identifier"
197            | "primitive_type"
198            | "scoped_type_identifier"
199            | "generic_type"
200            | "tuple_type"
201            | "reference_type"
202            | "impl_trait_type"
203    )
204}
205
206fn node_text<'tcx>(unit: CompileUnit<'tcx>, node: Node<'tcx>) -> String {
207    unit.file().get_text(node.start_byte(), node.end_byte())
208}
209
210fn clean(text: &str) -> String {
211    let mut out = String::new();
212    let mut last_was_ws = false;
213    for ch in text.chars() {
214        if ch.is_whitespace() {
215            if !last_was_ws && !out.is_empty() {
216                out.push(' ');
217            }
218            last_was_ws = true;
219        } else {
220            out.push(ch);
221            last_was_ws = false;
222        }
223    }
224    out.trim().to_string()
225}
llmcc_rust/descriptor/enumeration.rs

llmcc_rust/descriptor/
enumeration.rs
