Skip to main content

assura_codegen/
hir.rs

1//! Typed Rust HIR (High-level Intermediate Representation) for codegen.
2//!
3//! Instead of building Rust source code via string concatenation,
4//! codegen builds structured `RustItem` / `RustStmt` / `RustExpr` trees.
5//! These are converted to `syn` AST nodes, then formatted by `prettyplease`.
6//!
7//! Benefits:
8//! - Structural validation: malformed HIR caught at construction, not formatting.
9//! - Testability: assert on HIR structure (e.g., "2 Assert nodes"), not strings.
10//! - No string-escaping bugs: identifiers and types are typed, not interpolated.
11
12use std::fmt::Write as _;
13
14// ---------------------------------------------------------------------------
15// Core HIR types
16// ---------------------------------------------------------------------------
17
18/// A top-level Rust item (function, struct, enum, trait, const, mod, etc.).
19#[derive(Debug, Clone, PartialEq)]
20pub enum RustItem {
21    /// `pub fn name(params) -> ret { body }`
22    Fn(RustFn),
23    /// `pub struct Name { fields }` or `pub struct Name;`
24    Struct(RustStruct),
25    /// `pub enum Name { variants }`
26    Enum(RustEnum),
27    /// `pub trait Name { methods }`
28    Trait(RustTrait),
29    /// `impl Trait for Type { methods }`
30    Impl(RustImpl),
31    /// `pub mod name { items }`
32    Mod(RustMod),
33    /// `pub const NAME: Type = value;`
34    Const(RustConst),
35    /// `use path;`
36    Use(String),
37    /// `// comment` or `/// doc comment`
38    Comment(String),
39    /// An attribute like `#[allow(...)]` or `#![allow(...)]` applied to the crate/file.
40    InnerAttr(String),
41    /// Raw Rust code string (escape hatch for migration).
42    ///
43    /// Used during incremental migration: parts of codegen not yet converted
44    /// to HIR can emit raw strings that are spliced into the output.
45    Raw(String),
46}
47
48/// A Rust function definition.
49#[derive(Debug, Clone, PartialEq)]
50pub struct RustFn {
51    /// Function name.
52    pub name: String,
53    /// Generic type parameters (e.g., `["T", "U"]`).
54    pub type_params: Vec<String>,
55    /// Function parameters.
56    pub params: Vec<RustParam>,
57    /// Return type. `None` means no explicit return type (unit).
58    pub ret: Option<RustType>,
59    /// Function body statements.
60    pub body: Vec<RustStmt>,
61    /// Whether the function is `pub`.
62    pub is_pub: bool,
63    /// Whether the function is `unsafe`.
64    pub is_unsafe: bool,
65    /// Doc comments (each string is one `///` line).
66    pub doc: Vec<String>,
67    /// Attributes (e.g., `#[allow(dead_code)]`).
68    pub attrs: Vec<String>,
69    /// Whether this is an abstract method declaration (no body, e.g., trait methods).
70    pub is_abstract: bool,
71}
72
73impl Default for RustFn {
74    fn default() -> Self {
75        Self {
76            name: String::new(),
77            type_params: Vec::new(),
78            params: Vec::new(),
79            ret: None,
80            body: Vec::new(),
81            is_pub: true,
82            is_unsafe: false,
83            doc: Vec::new(),
84            attrs: Vec::new(),
85            is_abstract: false,
86        }
87    }
88}
89
90/// A function parameter.
91#[derive(Debug, Clone, PartialEq)]
92pub struct RustParam {
93    pub name: String,
94    pub ty: RustType,
95}
96
97/// A Rust type.
98#[derive(Debug, Clone, PartialEq)]
99pub enum RustType {
100    /// A simple named type: `i64`, `String`, `MyStruct`.
101    Named(String),
102    /// A generic type: `Vec<u8>`, `Result<T, E>`.
103    Generic(String, Vec<RustType>),
104    /// A reference: `&T` or `&[u8]`.
105    Ref(Box<RustType>),
106    /// A mutable reference: `&mut T`.
107    RefMut(Box<RustType>),
108    /// A tuple type: `(A, B, C)`.
109    Tuple(Vec<RustType>),
110    /// Unit type: `()`.
111    Unit,
112    /// Never type: `!`.
113    Never,
114    /// A raw type string (escape hatch for complex types not yet modeled).
115    Raw(String),
116}
117
118impl RustType {
119    /// Shorthand for common types.
120    pub fn i64() -> Self {
121        Self::Named("i64".into())
122    }
123    pub fn u64() -> Self {
124        Self::Named("u64".into())
125    }
126    pub fn f64() -> Self {
127        Self::Named("f64".into())
128    }
129    pub fn bool() -> Self {
130        Self::Named("bool".into())
131    }
132    pub fn string() -> Self {
133        Self::Named("String".into())
134    }
135    pub fn bytes() -> Self {
136        Self::Generic("Vec".into(), vec![Self::Named("u8".into())])
137    }
138    pub fn result(ok: RustType, err: RustType) -> Self {
139        Self::Generic("Result".into(), vec![ok, err])
140    }
141}
142
143/// A Rust statement inside a function body.
144#[derive(Debug, Clone, PartialEq)]
145pub enum RustStmt {
146    /// `let name: ty = init;` or `let name = init;`
147    Let {
148        name: String,
149        ty: Option<RustType>,
150        init: RustExpr,
151    },
152    /// `debug_assert!(cond, "label: msg");`
153    Assert { cond: String, label: String },
154    /// A `// comment` line.
155    Comment(String),
156    /// An expression used as a statement (e.g., `x;` or `x` as trailing expr).
157    Expr(RustExpr),
158    /// Raw Rust code (escape hatch during migration).
159    Raw(String),
160}
161
162/// A Rust expression.
163#[derive(Debug, Clone, PartialEq)]
164pub enum RustExpr {
165    /// An identifier: `x`, `RESULT_VAR`.
166    Ident(String),
167    /// A literal: `42`, `"hello"`, `true`.
168    Literal(String),
169    /// A function/method call: `func(args)`.
170    Call(String, Vec<RustExpr>),
171    /// A method call: `receiver.method(args)`.
172    MethodCall(Box<RustExpr>, String, Vec<RustExpr>),
173    /// `todo!("msg")`.
174    Todo(String),
175    /// `Ok(expr)`.
176    Ok(Box<RustExpr>),
177    /// `.clone()` on an expression.
178    Clone(Box<RustExpr>),
179    /// Raw Rust expression string (escape hatch).
180    Raw(String),
181}
182
183/// A Rust struct definition.
184#[derive(Debug, Clone, PartialEq)]
185pub struct RustStruct {
186    pub name: String,
187    pub type_params: Vec<String>,
188    pub fields: Vec<RustField>,
189    /// Derive attributes (e.g., `["Debug", "Clone", "PartialEq"]`).
190    pub derives: Vec<String>,
191    pub is_pub: bool,
192    /// Doc comments.
193    pub doc: Vec<String>,
194}
195
196impl Default for RustStruct {
197    fn default() -> Self {
198        Self {
199            name: String::new(),
200            type_params: Vec::new(),
201            fields: Vec::new(),
202            derives: vec!["Debug".into(), "Clone".into(), "PartialEq".into()],
203            is_pub: true,
204            doc: Vec::new(),
205        }
206    }
207}
208
209/// A struct field.
210#[derive(Debug, Clone, PartialEq)]
211pub struct RustField {
212    pub name: String,
213    pub ty: RustType,
214    pub is_pub: bool,
215}
216
217/// A Rust enum definition.
218#[derive(Debug, Clone, PartialEq)]
219pub struct RustEnum {
220    pub name: String,
221    pub type_params: Vec<String>,
222    pub variants: Vec<RustVariant>,
223    pub derives: Vec<String>,
224    pub is_pub: bool,
225    pub doc: Vec<String>,
226    /// Extra attributes beyond derives (e.g., `#[non_exhaustive]`).
227    pub attrs: Vec<String>,
228}
229
230impl Default for RustEnum {
231    fn default() -> Self {
232        Self {
233            name: String::new(),
234            type_params: Vec::new(),
235            variants: Vec::new(),
236            derives: vec!["Debug".into(), "Clone".into(), "PartialEq".into()],
237            is_pub: true,
238            doc: Vec::new(),
239            attrs: Vec::new(),
240        }
241    }
242}
243
244/// An enum variant.
245#[derive(Debug, Clone, PartialEq)]
246pub struct RustVariant {
247    pub name: String,
248    /// Tuple fields (empty for unit variants).
249    pub fields: Vec<RustType>,
250    /// Extra attributes (e.g., `#[error("...")]` for thiserror).
251    pub attrs: Vec<String>,
252}
253
254/// A Rust trait definition.
255#[derive(Debug, Clone, PartialEq)]
256pub struct RustTrait {
257    pub name: String,
258    pub type_params: Vec<String>,
259    pub supertraits: Vec<String>,
260    pub methods: Vec<RustFn>,
261    pub is_pub: bool,
262    pub doc: Vec<String>,
263}
264
265/// An `impl` block.
266#[derive(Debug, Clone, PartialEq)]
267pub struct RustImpl {
268    /// The trait being implemented (e.g., `"Display"`). `None` for inherent impls.
269    pub trait_name: Option<String>,
270    /// The type the impl is for (e.g., `"MyStruct"`).
271    pub target: String,
272    pub type_params: Vec<String>,
273    pub methods: Vec<RustFn>,
274}
275
276/// A module definition.
277#[derive(Debug, Clone, PartialEq)]
278pub struct RustMod {
279    pub name: String,
280    pub items: Vec<RustItem>,
281    pub is_pub: bool,
282    pub doc: Vec<String>,
283}
284
285/// A constant definition.
286#[derive(Debug, Clone, PartialEq)]
287pub struct RustConst {
288    pub name: String,
289    pub ty: RustType,
290    pub value: String,
291    pub is_pub: bool,
292    pub doc: Vec<String>,
293}
294
295// ---------------------------------------------------------------------------
296// HIR -> String rendering (via syn + prettyplease)
297// ---------------------------------------------------------------------------
298
299/// Render a list of `RustItem`s to formatted Rust source code.
300///
301/// Converts each item to its string representation, joins them,
302/// then formats the result via `syn::parse_file` + `prettyplease::unparse`.
303pub fn render_items(items: &[RustItem]) -> String {
304    let mut code = String::new();
305    for item in items {
306        render_item(item, &mut code);
307    }
308    // Format via prettyplease (same path as the existing format_rust)
309    match syn::parse_file(&code) {
310        Ok(syntax_tree) => prettyplease::unparse(&syntax_tree),
311        Err(e) => {
312            eprintln!("warning: HIR-generated Rust has syntax errors, skipping formatting: {e}");
313            format!("// WARNING: prettyplease formatting skipped (parse error: {e})\n\n{code}")
314        }
315    }
316}
317
318/// Render a single `RustItem` to a string, appending to `out`.
319fn render_item(item: &RustItem, out: &mut String) {
320    match item {
321        RustItem::Fn(f) => render_fn(f, out),
322        RustItem::Struct(s) => render_struct(s, out),
323        RustItem::Enum(e) => render_enum(e, out),
324        RustItem::Trait(t) => render_trait(t, out),
325        RustItem::Impl(i) => render_impl(i, out),
326        RustItem::Mod(m) => render_mod(m, out),
327        RustItem::Const(c) => render_const(c, out),
328        RustItem::Use(path) => {
329            let _ = writeln!(out, "use {path};");
330        }
331        RustItem::Comment(text) => {
332            let _ = writeln!(out, "// {text}");
333        }
334        RustItem::InnerAttr(attr) => {
335            let _ = writeln!(out, "#![{attr}]");
336        }
337        RustItem::Raw(code) => {
338            out.push_str(code);
339            if !code.ends_with('\n') {
340                out.push('\n');
341            }
342        }
343    }
344}
345
346fn render_fn(f: &RustFn, out: &mut String) {
347    render_fn_opts(f, out, &RenderOpts::default());
348}
349
350fn render_fn_opts(f: &RustFn, out: &mut String, opts: &RenderOpts) {
351    for doc in &f.doc {
352        let _ = writeln!(out, "/// {doc}");
353    }
354    for attr in &f.attrs {
355        let _ = writeln!(out, "#[{attr}]");
356    }
357
358    let vis = if f.is_pub { "pub " } else { "" };
359    let unsafe_kw = if f.is_unsafe { "unsafe " } else { "" };
360    let tps = if f.type_params.is_empty() {
361        String::new()
362    } else {
363        format!("<{}>", f.type_params.join(", "))
364    };
365
366    let params: Vec<String> = f
367        .params
368        .iter()
369        .map(|p| {
370            // Self receiver params are rendered without an explicit type
371            if p.name == "&self" || p.name == "self" || p.name == "&mut self" {
372                p.name.clone()
373            } else {
374                format!("{}: {}", p.name, render_type(&p.ty))
375            }
376        })
377        .collect();
378    let params_s = params.join(", ");
379
380    let ret = match &f.ret {
381        Some(ty) => format!(" -> {}", render_type(ty)),
382        None => String::new(),
383    };
384
385    if f.is_abstract {
386        let _ = writeln!(
387            out,
388            "{vis}{unsafe_kw}fn {}{tps}({params_s}){ret};\n",
389            f.name
390        );
391    } else {
392        let _ = writeln!(
393            out,
394            "{vis}{unsafe_kw}fn {}{tps}({params_s}){ret} {{",
395            f.name
396        );
397
398        for stmt in &f.body {
399            render_stmt_opts(stmt, out, 1, opts);
400        }
401
402        out.push_str("}\n\n");
403    }
404}
405
406fn render_struct(s: &RustStruct, out: &mut String) {
407    for doc in &s.doc {
408        let _ = writeln!(out, "/// {doc}");
409    }
410    if !s.derives.is_empty() {
411        let _ = writeln!(out, "#[derive({})]", s.derives.join(", "));
412    }
413
414    let vis = if s.is_pub { "pub " } else { "" };
415    let tps = if s.type_params.is_empty() {
416        String::new()
417    } else {
418        format!("<{}>", s.type_params.join(", "))
419    };
420
421    if s.fields.is_empty() {
422        let _ = writeln!(out, "{vis}struct {}{tps};\n", s.name);
423    } else {
424        let _ = writeln!(out, "{vis}struct {}{tps} {{", s.name);
425        for field in &s.fields {
426            let fvis = if field.is_pub { "pub " } else { "" };
427            let _ = writeln!(out, "    {fvis}{}: {},", field.name, render_type(&field.ty));
428        }
429        out.push_str("}\n\n");
430    }
431}
432
433fn render_enum(e: &RustEnum, out: &mut String) {
434    for doc in &e.doc {
435        let _ = writeln!(out, "/// {doc}");
436    }
437    if !e.derives.is_empty() {
438        let _ = writeln!(out, "#[derive({})]", e.derives.join(", "));
439    }
440    for attr in &e.attrs {
441        let _ = writeln!(out, "#[{attr}]");
442    }
443
444    let vis = if e.is_pub { "pub " } else { "" };
445    let tps = if e.type_params.is_empty() {
446        String::new()
447    } else {
448        format!("<{}>", e.type_params.join(", "))
449    };
450
451    let _ = writeln!(out, "{vis}enum {}{tps} {{", e.name);
452    for variant in &e.variants {
453        for attr in &variant.attrs {
454            let _ = writeln!(out, "    #[{attr}]");
455        }
456        if variant.fields.is_empty() {
457            let _ = writeln!(out, "    {},", variant.name);
458        } else {
459            let fields: Vec<String> = variant.fields.iter().map(render_type).collect();
460            let _ = writeln!(out, "    {}({}),", variant.name, fields.join(", "));
461        }
462    }
463    out.push_str("}\n\n");
464}
465
466fn render_trait(t: &RustTrait, out: &mut String) {
467    for doc in &t.doc {
468        let _ = writeln!(out, "/// {doc}");
469    }
470
471    let vis = if t.is_pub { "pub " } else { "" };
472    let tps = if t.type_params.is_empty() {
473        String::new()
474    } else {
475        format!("<{}>", t.type_params.join(", "))
476    };
477    let bounds = if t.supertraits.is_empty() {
478        String::new()
479    } else {
480        format!(": {}", t.supertraits.join(" + "))
481    };
482
483    let _ = writeln!(out, "{vis}trait {}{tps}{bounds} {{", t.name);
484    for method in &t.methods {
485        // Render as a trait method (with or without default body)
486        render_fn(method, out);
487    }
488    out.push_str("}\n\n");
489}
490
491fn render_impl(i: &RustImpl, out: &mut String) {
492    let tps = if i.type_params.is_empty() {
493        String::new()
494    } else {
495        format!("<{}>", i.type_params.join(", "))
496    };
497
498    match &i.trait_name {
499        Some(trait_name) => {
500            let _ = writeln!(out, "impl{tps} {trait_name} for {}{tps} {{", i.target);
501        }
502        None => {
503            let _ = writeln!(out, "impl{tps} {}{tps} {{", i.target);
504        }
505    }
506    for method in &i.methods {
507        render_fn(method, out);
508    }
509    out.push_str("}\n\n");
510}
511
512fn render_mod(m: &RustMod, out: &mut String) {
513    for doc in &m.doc {
514        let _ = writeln!(out, "/// {doc}");
515    }
516    let vis = if m.is_pub { "pub " } else { "" };
517    let _ = writeln!(out, "{vis}mod {} {{", m.name);
518    for item in &m.items {
519        render_item(item, out);
520    }
521    out.push_str("}\n\n");
522}
523
524fn render_const(c: &RustConst, out: &mut String) {
525    for doc in &c.doc {
526        let _ = writeln!(out, "/// {doc}");
527    }
528    let vis = if c.is_pub { "pub " } else { "" };
529    let _ = writeln!(
530        out,
531        "{vis}const {}: {} = {};",
532        c.name,
533        render_type(&c.ty),
534        c.value
535    );
536}
537
538fn render_type(ty: &RustType) -> String {
539    match ty {
540        RustType::Named(name) => name.clone(),
541        RustType::Generic(name, args) => {
542            let args_s: Vec<String> = args.iter().map(render_type).collect();
543            format!("{name}<{}>", args_s.join(", "))
544        }
545        RustType::Ref(inner) => format!("&{}", render_type(inner)),
546        RustType::RefMut(inner) => format!("&mut {}", render_type(inner)),
547        RustType::Tuple(items) => {
548            let items_s: Vec<String> = items.iter().map(render_type).collect();
549            format!("({})", items_s.join(", "))
550        }
551        RustType::Unit => "()".into(),
552        RustType::Never => "!".into(),
553        RustType::Raw(s) => s.clone(),
554    }
555}
556
557#[cfg(test)]
558pub(crate) fn render_stmt(stmt: &RustStmt, out: &mut String, indent: usize) {
559    render_stmt_opts(stmt, out, indent, &RenderOpts::default());
560}
561
562/// Options that control how statements are rendered.
563#[derive(Debug, Clone, Default)]
564pub(crate) struct RenderOpts {
565    /// When true, assertions emit `assura_runtime::contract_violation` calls
566    /// that persist in release builds instead of `debug_assert!`.
567    pub runtime_checks: bool,
568    /// The contract name for runtime violation reports.
569    pub contract_name: String,
570}
571
572pub(crate) fn render_stmt_opts(
573    stmt: &RustStmt,
574    out: &mut String,
575    indent: usize,
576    opts: &RenderOpts,
577) {
578    let pad = "    ".repeat(indent);
579    match stmt {
580        RustStmt::Let { name, ty, init } => {
581            let ty_s = match ty {
582                Some(t) => format!(": {}", render_type(t)),
583                None => String::new(),
584            };
585            let _ = writeln!(out, "{pad}let {name}{ty_s} = {};", render_expr(init));
586        }
587        RustStmt::Assert { cond, label } => {
588            // If expression references deep field chains (e.g., state.head.extra),
589            // emit as a comment since stub types don't have these fields.
590            if crate::has_deep_field_access(cond) {
591                let _ = writeln!(out, "{pad}// {label}: {}", cond.replace('"', "\\\""));
592            } else if opts.runtime_checks {
593                // Runtime checks: persist in release builds via assura_runtime
594                let escaped_cond = cond.replace('"', "\\\"");
595                let contract = opts.contract_name.replace('"', "\\\"");
596                if cond.contains('\n') {
597                    let flat_cond = cond.replace('\n', " ");
598                    let _ = writeln!(
599                        out,
600                        "{pad}if !({{ {cond} }}) {{ assura_runtime::contract_violation(\"{contract}\", \"{label}\", \"{escaped_cond}\", file!(), line!()); }}"
601                    );
602                    let _ = flat_cond.len(); // suppress unused warning
603                } else {
604                    let _ = writeln!(
605                        out,
606                        "{pad}if !({cond}) {{ assura_runtime::contract_violation(\"{contract}\", \"{label}\", \"{escaped_cond}\", file!(), line!()); }}"
607                    );
608                }
609            } else if cond.contains('\n') {
610                let msg = cond
611                    .replace('\n', " ")
612                    .replace('"', "\\\"")
613                    .replace('{', "{{")
614                    .replace('}', "}}");
615                let _ = writeln!(out, "{pad}debug_assert!({{ {cond} }}, \"{label}: {msg}\");");
616            } else {
617                // Escape braces so `if { } else { }` ensures do not break the
618                // format string used to build the debug_assert! source.
619                let escaped_cond = cond
620                    .replace('"', "\\\"")
621                    .replace('{', "{{")
622                    .replace('}', "}}");
623                let _ = writeln!(
624                    out,
625                    "{pad}debug_assert!({cond}, \"{label}: {escaped_cond}\");"
626                );
627            }
628        }
629        RustStmt::Comment(text) => {
630            let _ = writeln!(out, "{pad}// {text}");
631        }
632        RustStmt::Expr(expr) => {
633            let _ = writeln!(out, "{pad}{}", render_expr(expr));
634        }
635        RustStmt::Raw(code) => {
636            for line in code.lines() {
637                let _ = writeln!(out, "{pad}{line}");
638            }
639        }
640    }
641}
642
643fn render_expr(expr: &RustExpr) -> String {
644    match expr {
645        RustExpr::Ident(name) => name.clone(),
646        RustExpr::Literal(lit) => lit.clone(),
647        RustExpr::Call(func, args) => {
648            let args_s: Vec<String> = args.iter().map(render_expr).collect();
649            format!("{func}({})", args_s.join(", "))
650        }
651        RustExpr::MethodCall(receiver, method, args) => {
652            let args_s: Vec<String> = args.iter().map(render_expr).collect();
653            format!("{}.{method}({})", render_expr(receiver), args_s.join(", "))
654        }
655        RustExpr::Todo(msg) => format!("todo!(\"{msg}\")"),
656        RustExpr::Ok(inner) => format!("Ok({})", render_expr(inner)),
657        RustExpr::Clone(inner) => format!("{}.clone()", render_expr(inner)),
658        RustExpr::Raw(code) => code.clone(),
659    }
660}
661
662// ---------------------------------------------------------------------------
663// Builder helpers for common codegen patterns
664// ---------------------------------------------------------------------------
665
666/// Build a `#[derive(Debug, thiserror::Error)]` enum for contract error types.
667///
668/// Each variant gets a `#[error("VariantName")]` attribute for Display impl.
669pub fn build_error_enum(contract_name: &str, variants: &[String]) -> RustItem {
670    let enum_name = format!("{contract_name}Error");
671    RustItem::Enum(RustEnum {
672        name: enum_name,
673        variants: variants
674            .iter()
675            .map(|v| RustVariant {
676                name: v.clone(),
677                fields: vec![],
678                attrs: vec![format!("error(\"{v}\")")],
679            })
680            .collect(),
681        derives: vec!["Debug".into(), "thiserror::Error".into()],
682        ..RustEnum::default()
683    })
684}
685
686/// Build a `#[derive(Debug, Clone, PartialEq)]` enum with optional Display impl
687/// and exhaustiveness check function.
688///
689/// Returns a list of items: the enum definition, optionally a Display impl,
690/// and optionally an exhaustiveness check function.
691pub fn build_enum_def(e: &assura_ast::EnumDef) -> Vec<RustItem> {
692    let mut items = Vec::new();
693
694    let tps: Vec<String> = e.type_params.clone();
695
696    // Build enum
697    let rust_enum = RustEnum {
698        name: e.name.clone(),
699        type_params: tps.clone(),
700        variants: e
701            .variants
702            .iter()
703            .map(|v| RustVariant {
704                name: v.name.clone(),
705                fields: v
706                    .fields
707                    .iter()
708                    .map(|f| {
709                        let toks: Vec<String> = f.split_whitespace().map(String::from).collect();
710                        RustType::Raw(super::map_type_tokens(&toks))
711                    })
712                    .collect(),
713                attrs: vec![],
714            })
715            .collect(),
716        ..RustEnum::default()
717    };
718    items.push(RustItem::Enum(rust_enum));
719
720    // Display impl (only for non-generic enums)
721    if tps.is_empty() {
722        let arms: Vec<String> = e
723            .variants
724            .iter()
725            .map(|v| {
726                if v.fields.is_empty() {
727                    format!("{}::{} => write!(f, \"{}\")", e.name, v.name, v.name)
728                } else {
729                    let underscores: Vec<&str> = (0..v.fields.len()).map(|_| "_").collect();
730                    format!(
731                        "{}::{}({}) => write!(f, \"{}(...)\")",
732                        e.name,
733                        v.name,
734                        underscores.join(", "),
735                        v.name
736                    )
737                }
738            })
739            .collect();
740
741        let match_body = format!("match self {{ {} }}", arms.join(", "));
742        items.push(RustItem::Impl(RustImpl {
743            trait_name: Some("std::fmt::Display".into()),
744            target: e.name.clone(),
745            type_params: vec![],
746            methods: vec![RustFn {
747                name: "fmt".into(),
748                params: vec![
749                    RustParam {
750                        name: "&self".into(),
751                        ty: RustType::Raw("&Self".into()),
752                    },
753                    RustParam {
754                        name: "f".into(),
755                        ty: RustType::Raw("&mut std::fmt::Formatter<'_>".into()),
756                    },
757                ],
758                ret: Some(RustType::Raw("std::fmt::Result".into())),
759                body: vec![RustStmt::Raw(match_body)],
760                is_pub: false,
761                ..RustFn::default()
762            }],
763        }));
764    }
765
766    // Exhaustiveness check (only for non-generic, non-empty enums)
767    if !e.variants.is_empty() && tps.is_empty() {
768        let arms: Vec<String> = e
769            .variants
770            .iter()
771            .map(|v| {
772                if v.fields.is_empty() {
773                    format!("{}::{} => \"{}\"", e.name, v.name, v.name)
774                } else {
775                    let underscores: Vec<&str> = (0..v.fields.len()).map(|_| "_").collect();
776                    format!(
777                        "{}::{}({}) => \"{}\"",
778                        e.name,
779                        v.name,
780                        underscores.join(", "),
781                        v.name
782                    )
783                }
784            })
785            .collect();
786
787        let match_body = format!("match v {{ {} }}", arms.join(", "));
788        items.push(RustItem::Fn(RustFn {
789            name: format!("__exhaustive_check_{}", e.name.to_lowercase()),
790            params: vec![RustParam {
791                name: "v".into(),
792                ty: RustType::Ref(Box::new(RustType::Named(e.name.clone()))),
793            }],
794            ret: Some(RustType::Raw("&'static str".into())),
795            body: vec![RustStmt::Raw(match_body)],
796            is_pub: false,
797            attrs: vec!["allow(dead_code)".into()],
798            doc: vec![
799                format!("Compile-time exhaustiveness check for `{}`.", e.name),
800                "Adding a variant without updating all match sites causes a build error.".into(),
801            ],
802            ..RustFn::default()
803        }));
804    }
805
806    items
807}
808
809/// Render a single HIR item to a code string (without prettyplease formatting).
810///
811/// Used by migration code that appends HIR items to an existing string buffer.
812pub fn render_item_raw(item: &RustItem) -> String {
813    let mut out = String::new();
814    render_item(item, &mut out);
815    out
816}
817
818/// Like `render_item_raw` but passes `RenderOpts` to control assertion style.
819pub(crate) fn render_item_raw_with_opts(item: &RustItem, opts: &RenderOpts) -> String {
820    let mut out = String::new();
821    render_item_with_opts(item, &mut out, opts);
822    out
823}
824
825fn render_item_with_opts(item: &RustItem, out: &mut String, opts: &RenderOpts) {
826    match item {
827        RustItem::Fn(f) => render_fn_opts(f, out, opts),
828        // Only functions contain assertions; other items delegate to the
829        // regular renderer.
830        other => render_item(other, out),
831    }
832}
833
834// ---------------------------------------------------------------------------
835// Builders: AST -> HIR
836// ---------------------------------------------------------------------------
837
838/// Build HIR items for an Assura `TypeDef`.
839///
840/// Returns a `Vec<RustItem>` because some type bodies produce multiple items
841/// (e.g., a refined type produces a newtype struct).
842pub fn build_type_def(t: &assura_ast::TypeDef) -> Vec<RustItem> {
843    use assura_ast::TypeBody;
844    let tps = t.type_params.clone();
845
846    match &t.body {
847        TypeBody::Struct(fields) => {
848            // Always emit `pub` fields: contract modules live in child `mod`s
849            // and must access struct fields (e.g. IR `field` → `slot.y`).
850            // Private Assura fields would be unreadable outside the crate root.
851            let rust_fields: Vec<RustField> = fields
852                .iter()
853                .map(|f| {
854                    let ty_tokens = f.ty.as_ref().map(|t| t.to_tokens()).unwrap_or_default();
855                    RustField {
856                        name: f.name.clone(),
857                        ty: RustType::Raw(crate::map_type_tokens(&ty_tokens)),
858                        is_pub: true,
859                    }
860                })
861                .collect();
862            // Emit a cfg(test) Arbitrary impl so proptest can invent values for
863            // user structs (fields are all pub primitive-mapped types).
864            let mut items = vec![RustItem::Struct(RustStruct {
865                name: t.name.clone(),
866                type_params: tps.clone(),
867                fields: rust_fields.clone(),
868                derives: vec!["Debug".into(), "Clone".into(), "PartialEq".into()],
869                ..RustStruct::default()
870            })];
871            if tps.is_empty() && !rust_fields.is_empty() {
872                let field_names: Vec<&str> = rust_fields.iter().map(|f| f.name.as_str()).collect();
873                let field_tys: Vec<String> = rust_fields
874                    .iter()
875                    .map(|f| match &f.ty {
876                        RustType::Raw(s) => s.clone(),
877                        other => format!("{other:?}"),
878                    })
879                    .collect();
880                const PRIMS: &[&str] = &[
881                    "i64", "u64", "i32", "u32", "bool", "f64", "f32", "i8", "u8", "i16", "u16",
882                    "isize", "usize", "String",
883                ];
884                // Primitives or peer user structs (already emit Arbitrary when
885                // declared earlier in the file, e.g. Outer { inner: Inner }).
886                let arb_field =
887                    |ty: &str| PRIMS.contains(&ty) || crate::types_gen::is_user_type_name(ty);
888                if field_tys.iter().all(|ty| arb_field(ty.as_str())) {
889                    // Use the same strategies as contract proptest (i32-range for i64)
890                    // so field-bearing structs do not re-introduce full-range overflow.
891                    let field_strats: Vec<String> = field_tys
892                        .iter()
893                        .map(|ty| crate::contract::proptest_strategy_for_type(ty))
894                        .collect();
895                    let destructure = field_names.join(", ");
896                    let construct = field_names.join(", ");
897                    let strategy = if field_names.len() == 1 {
898                        format!(
899                            "({}).prop_map(|{destructure}| {} {{ {construct} }})",
900                            field_strats[0], t.name
901                        )
902                    } else {
903                        format!(
904                            "({}).prop_map(|({destructure})| {} {{ {construct} }})",
905                            field_strats.join(", "),
906                            t.name
907                        )
908                    };
909                    items.push(RustItem::Raw(format!(
910                        "#[cfg(test)]\n\
911                         impl proptest::prelude::Arbitrary for {} {{\n\
912                             type Parameters = ();\n\
913                             type Strategy = proptest::strategy::BoxedStrategy<Self>;\n\
914                             fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {{\n\
915                                 use proptest::prelude::*;\n\
916                                 {strategy}.boxed()\n\
917                             }}\n\
918                         }}\n",
919                        t.name
920                    )));
921                }
922            }
923            items
924        }
925        TypeBody::Alias(tokens) => {
926            let rust_ty = crate::map_type_tokens(tokens);
927            vec![RustItem::Raw(format!("pub type {} = {rust_ty};\n", t.name))]
928        }
929        TypeBody::Refined(tokens) => {
930            let base = crate::extract_base_type_from_refined(tokens);
931            let tps_str = if tps.is_empty() {
932                String::new()
933            } else {
934                format!("<{}>", tps.join(", "))
935            };
936            vec![RustItem::Raw(format!(
937                "#[derive(Debug, Clone, PartialEq)]\npub struct {}{tps_str}(pub {base});\n",
938                t.name
939            ))]
940        }
941        TypeBody::Empty => {
942            let tps_str = if tps.is_empty() {
943                String::new()
944            } else {
945                format!("<{}>", tps.join(", "))
946            };
947            if tps.is_empty() {
948                vec![RustItem::Raw(format!(
949                    "#[derive(Debug, Clone, PartialEq)]\npub struct {}{};\n",
950                    t.name, tps_str
951                ))]
952            } else {
953                let phantoms: Vec<String> = tps
954                    .iter()
955                    .map(|p| format!("std::marker::PhantomData<{p}>"))
956                    .collect();
957                vec![RustItem::Raw(format!(
958                    "#[derive(Debug, Clone, PartialEq)]\npub struct {}{tps_str}({});\n",
959                    t.name,
960                    phantoms.join(", ")
961                ))]
962            }
963        }
964    }
965}
966
967// ---------------------------------------------------------------------------
968// Tests
969// ---------------------------------------------------------------------------
970
971#[cfg(test)]
972mod tests {
973    use super::*;
974
975    #[test]
976    fn render_simple_fn() {
977        let f = RustFn {
978            name: "add".into(),
979            params: vec![
980                RustParam {
981                    name: "a".into(),
982                    ty: RustType::i64(),
983                },
984                RustParam {
985                    name: "b".into(),
986                    ty: RustType::i64(),
987                },
988            ],
989            ret: Some(RustType::i64()),
990            body: vec![RustStmt::Expr(RustExpr::Raw("a + b".into()))],
991            ..RustFn::default()
992        };
993        let code = render_items(&[RustItem::Fn(f)]);
994        assert!(code.contains("pub fn add(a: i64, b: i64) -> i64"));
995        assert!(code.contains("a + b"));
996    }
997
998    #[test]
999    fn render_unit_struct() {
1000        let s = RustStruct {
1001            name: "Marker".into(),
1002            ..RustStruct::default()
1003        };
1004        let code = render_items(&[RustItem::Struct(s)]);
1005        assert!(code.contains("#[derive(Debug, Clone, PartialEq)]"));
1006        assert!(code.contains("pub struct Marker;"));
1007    }
1008
1009    #[test]
1010    fn render_struct_with_fields() {
1011        let s = RustStruct {
1012            name: "Point".into(),
1013            fields: vec![
1014                RustField {
1015                    name: "x".into(),
1016                    ty: RustType::f64(),
1017                    is_pub: true,
1018                },
1019                RustField {
1020                    name: "y".into(),
1021                    ty: RustType::f64(),
1022                    is_pub: true,
1023                },
1024            ],
1025            ..RustStruct::default()
1026        };
1027        let code = render_items(&[RustItem::Struct(s)]);
1028        assert!(code.contains("pub struct Point"));
1029        assert!(code.contains("pub x: f64"));
1030        assert!(code.contains("pub y: f64"));
1031    }
1032
1033    #[test]
1034    fn render_error_enum() {
1035        let e = RustEnum {
1036            name: "SafeDivisionError".into(),
1037            variants: vec![
1038                RustVariant {
1039                    name: "DivByZero".into(),
1040                    fields: vec![],
1041                    attrs: vec!["error(\"DivByZero\")".into()],
1042                },
1043                RustVariant {
1044                    name: "Overflow".into(),
1045                    fields: vec![],
1046                    attrs: vec!["error(\"Overflow\")".into()],
1047                },
1048            ],
1049            derives: vec!["Debug".into(), "thiserror::Error".into()],
1050            ..RustEnum::default()
1051        };
1052        let code = render_items(&[RustItem::Enum(e)]);
1053        assert!(code.contains("#[derive(Debug, thiserror::Error)]"));
1054        assert!(code.contains("pub enum SafeDivisionError"));
1055        assert!(code.contains("#[error(\"DivByZero\")]"));
1056        assert!(code.contains("DivByZero,"));
1057        assert!(code.contains("#[error(\"Overflow\")]"));
1058        assert!(code.contains("Overflow,"));
1059    }
1060
1061    #[test]
1062    fn render_fn_with_assert() {
1063        let f = RustFn {
1064            name: "check".into(),
1065            params: vec![RustParam {
1066                name: "x".into(),
1067                ty: RustType::i64(),
1068            }],
1069            ret: Some(RustType::i64()),
1070            body: vec![
1071                RustStmt::Assert {
1072                    cond: "(x > 0)".into(),
1073                    label: "requires".into(),
1074                },
1075                RustStmt::Let {
1076                    name: "__assura_result".into(),
1077                    ty: Some(RustType::i64()),
1078                    init: RustExpr::Todo("implementation provided by AI agent".into()),
1079                },
1080                RustStmt::Expr(RustExpr::Ident("__assura_result".into())),
1081            ],
1082            ..RustFn::default()
1083        };
1084        let code = render_items(&[RustItem::Fn(f)]);
1085        assert!(code.contains("debug_assert!"));
1086        assert!(code.contains("requires"));
1087        assert!(code.contains("todo!"));
1088    }
1089
1090    #[test]
1091    fn render_const() {
1092        let c = RustConst {
1093            name: "MAX_SIZE".into(),
1094            ty: RustType::u64(),
1095            value: "1024".into(),
1096            is_pub: true,
1097            doc: vec![],
1098        };
1099        let code = render_items(&[RustItem::Const(c)]);
1100        assert!(code.contains("pub const MAX_SIZE: u64 = 1024;"));
1101    }
1102
1103    #[test]
1104    fn render_generic_type() {
1105        let ty = RustType::result(RustType::i64(), RustType::Named("MyError".into()));
1106        assert_eq!(render_type(&ty), "Result<i64, MyError>");
1107    }
1108
1109    #[test]
1110    fn render_module() {
1111        let m = RustMod {
1112            name: "contract_foo".into(),
1113            items: vec![RustItem::Comment("inner".into())],
1114            is_pub: true,
1115            doc: vec!["Contract: Foo".into()],
1116        };
1117        let code = render_items(&[RustItem::Mod(m)]);
1118        assert!(code.contains("/// Contract: Foo"));
1119        assert!(code.contains("pub mod contract_foo"));
1120    }
1121
1122    #[test]
1123    fn render_impl_block() {
1124        let i = RustImpl {
1125            trait_name: Some("Display".into()),
1126            target: "Color".into(),
1127            type_params: vec![],
1128            methods: vec![RustFn {
1129                name: "fmt".into(),
1130                params: vec![
1131                    RustParam {
1132                        name: "&self".into(),
1133                        ty: RustType::Raw("&Self".into()),
1134                    },
1135                    RustParam {
1136                        name: "f".into(),
1137                        ty: RustType::Raw("&mut std::fmt::Formatter<'_>".into()),
1138                    },
1139                ],
1140                ret: Some(RustType::Raw("std::fmt::Result".into())),
1141                body: vec![RustStmt::Raw(
1142                    "match self { _ => write!(f, \"Color\") }".into(),
1143                )],
1144                is_pub: false,
1145                ..RustFn::default()
1146            }],
1147        };
1148        let code = render_items(&[RustItem::Impl(i)]);
1149        assert!(code.contains("impl Display for Color"));
1150    }
1151
1152    #[test]
1153    fn type_shorthand_helpers() {
1154        assert_eq!(render_type(&RustType::i64()), "i64");
1155        assert_eq!(render_type(&RustType::u64()), "u64");
1156        assert_eq!(render_type(&RustType::f64()), "f64");
1157        assert_eq!(render_type(&RustType::bool()), "bool");
1158        assert_eq!(render_type(&RustType::string()), "String");
1159        assert_eq!(render_type(&RustType::bytes()), "Vec<u8>");
1160        assert_eq!(render_type(&RustType::Unit), "()");
1161        assert_eq!(render_type(&RustType::Never), "!");
1162    }
1163
1164    #[test]
1165    fn render_trait_def() {
1166        let t = RustTrait {
1167            name: "Serializable".into(),
1168            type_params: vec![],
1169            supertraits: vec!["Clone".into()],
1170            methods: vec![RustFn {
1171                name: "serialize".into(),
1172                params: vec![RustParam {
1173                    name: "&self".into(),
1174                    ty: RustType::Raw("&Self".into()),
1175                }],
1176                ret: Some(RustType::bytes()),
1177                body: vec![],
1178                is_pub: false,
1179                ..RustFn::default()
1180            }],
1181            is_pub: true,
1182            doc: vec![],
1183        };
1184        let code = render_items(&[RustItem::Trait(t)]);
1185        assert!(code.contains("pub trait Serializable: Clone"));
1186    }
1187
1188    // --- Structural tests ---
1189
1190    #[test]
1191    fn count_asserts_in_fn() {
1192        let f = RustFn {
1193            name: "check".into(),
1194            params: vec![RustParam {
1195                name: "x".into(),
1196                ty: RustType::i64(),
1197            }],
1198            ret: Some(RustType::i64()),
1199            body: vec![
1200                RustStmt::Assert {
1201                    cond: "(x > 0)".into(),
1202                    label: "requires".into(),
1203                },
1204                RustStmt::Assert {
1205                    cond: "(x < 100)".into(),
1206                    label: "requires".into(),
1207                },
1208                RustStmt::Expr(RustExpr::Ident("x".into())),
1209            ],
1210            ..RustFn::default()
1211        };
1212        let assert_count = f
1213            .body
1214            .iter()
1215            .filter(|s| matches!(s, RustStmt::Assert { .. }))
1216            .count();
1217        assert_eq!(assert_count, 2, "should have exactly 2 assert statements");
1218    }
1219
1220    // --- HIR builder structural tests ---
1221
1222    #[test]
1223    fn build_error_enum_structure() {
1224        let item = build_error_enum("SafeDivision", &["DivByZero".into(), "Overflow".into()]);
1225        if let RustItem::Enum(e) = &item {
1226            assert_eq!(e.name, "SafeDivisionError");
1227            assert_eq!(e.variants.len(), 2);
1228            assert_eq!(e.variants[0].name, "DivByZero");
1229            assert_eq!(e.variants[1].name, "Overflow");
1230            assert!(e.derives.contains(&"thiserror::Error".to_string()));
1231            // Each variant has an #[error("...")] attribute
1232            for v in &e.variants {
1233                assert_eq!(v.attrs.len(), 1);
1234                assert!(v.attrs[0].starts_with("error(\""));
1235            }
1236        } else {
1237            panic!("Expected RustItem::Enum");
1238        }
1239    }
1240
1241    #[test]
1242    fn build_enum_def_structure() {
1243        let enum_def = assura_ast::EnumDef {
1244            name: "Color".into(),
1245            type_params: vec![],
1246            variants: vec![
1247                assura_ast::EnumVariant {
1248                    name: "Red".into(),
1249                    fields: vec![],
1250                },
1251                assura_ast::EnumVariant {
1252                    name: "Blue".into(),
1253                    fields: vec![],
1254                },
1255                assura_ast::EnumVariant {
1256                    name: "Custom".into(),
1257                    fields: vec!["Int".into()],
1258                },
1259            ],
1260        };
1261        let items = build_enum_def(&enum_def);
1262        // Should produce: enum, Display impl, exhaustiveness check
1263        assert_eq!(items.len(), 3, "enum + Display + exhaustive");
1264
1265        // First item: the enum
1266        if let RustItem::Enum(e) = &items[0] {
1267            assert_eq!(e.name, "Color");
1268            assert_eq!(e.variants.len(), 3);
1269            assert_eq!(e.variants[2].name, "Custom");
1270            assert_eq!(e.variants[2].fields.len(), 1);
1271        } else {
1272            panic!("Expected RustItem::Enum");
1273        }
1274
1275        // Second item: Display impl
1276        assert!(
1277            matches!(&items[1], RustItem::Impl(i) if i.trait_name.as_deref() == Some("std::fmt::Display"))
1278        );
1279
1280        // Third item: exhaustiveness check fn
1281        assert!(matches!(&items[2], RustItem::Fn(f) if f.name == "__exhaustive_check_color"));
1282    }
1283
1284    #[test]
1285    fn build_enum_def_generic_no_display() {
1286        let enum_def = assura_ast::EnumDef {
1287            name: "Option".into(),
1288            type_params: vec!["T".into()],
1289            variants: vec![
1290                assura_ast::EnumVariant {
1291                    name: "Some".into(),
1292                    fields: vec!["T".into()],
1293                },
1294                assura_ast::EnumVariant {
1295                    name: "None".into(),
1296                    fields: vec![],
1297                },
1298            ],
1299        };
1300        let items = build_enum_def(&enum_def);
1301        // Generic enums: no Display impl, no exhaustiveness check
1302        assert_eq!(
1303            items.len(),
1304            1,
1305            "only enum, no Display or exhaustive for generic"
1306        );
1307        assert!(matches!(&items[0], RustItem::Enum(e) if e.type_params == vec!["T"]));
1308    }
1309
1310    // --- Round-trip test ---
1311
1312    #[test]
1313    fn round_trip_parse_format() {
1314        // Build HIR -> render -> parse via syn -> re-format via prettyplease
1315        let items = vec![
1316            RustItem::Fn(RustFn {
1317                name: "add".into(),
1318                params: vec![
1319                    RustParam {
1320                        name: "a".into(),
1321                        ty: RustType::i64(),
1322                    },
1323                    RustParam {
1324                        name: "b".into(),
1325                        ty: RustType::i64(),
1326                    },
1327                ],
1328                ret: Some(RustType::i64()),
1329                body: vec![RustStmt::Expr(RustExpr::Raw("a + b".into()))],
1330                ..RustFn::default()
1331            }),
1332            RustItem::Struct(RustStruct {
1333                name: "Point".into(),
1334                fields: vec![
1335                    RustField {
1336                        name: "x".into(),
1337                        ty: RustType::f64(),
1338                        is_pub: true,
1339                    },
1340                    RustField {
1341                        name: "y".into(),
1342                        ty: RustType::f64(),
1343                        is_pub: true,
1344                    },
1345                ],
1346                ..RustStruct::default()
1347            }),
1348        ];
1349
1350        let rendered = render_items(&items);
1351        // Parse the rendered output back with syn
1352        let parsed = syn::parse_file(&rendered);
1353        assert!(
1354            parsed.is_ok(),
1355            "Rendered HIR should parse as valid Rust: {:?}",
1356            parsed.err()
1357        );
1358
1359        // Re-format the parsed result
1360        let reparsed = prettyplease::unparse(&parsed.unwrap());
1361        // The double-formatted result should be identical (idempotent)
1362        assert_eq!(rendered, reparsed, "Formatting should be idempotent");
1363    }
1364
1365    // --- Type mapping coverage test ---
1366
1367    #[test]
1368    fn type_rendering_coverage() {
1369        // Test all RustType variants render correctly
1370        let cases = vec![
1371            (RustType::Named("MyStruct".into()), "MyStruct"),
1372            (
1373                RustType::Generic("Vec".into(), vec![RustType::Named("u8".into())]),
1374                "Vec<u8>",
1375            ),
1376            (
1377                RustType::Ref(Box::new(RustType::Named("str".into()))),
1378                "&str",
1379            ),
1380            (RustType::RefMut(Box::new(RustType::i64())), "&mut i64"),
1381            (
1382                RustType::Tuple(vec![RustType::i64(), RustType::bool()]),
1383                "(i64, bool)",
1384            ),
1385            (RustType::Unit, "()"),
1386            (RustType::Never, "!"),
1387            (RustType::Raw("Box<dyn Fn()>".into()), "Box<dyn Fn()>"),
1388        ];
1389        for (ty, expected) in cases {
1390            assert_eq!(
1391                render_type(&ty),
1392                expected,
1393                "type rendering mismatch for {:?}",
1394                ty
1395            );
1396        }
1397    }
1398
1399    // --- Malformed HIR rejection test ---
1400
1401    #[test]
1402    fn malformed_hir_produces_warning() {
1403        // A function with syntax errors in its body should produce a warning comment
1404        let items = vec![RustItem::Fn(RustFn {
1405            name: "bad".into(),
1406            body: vec![RustStmt::Raw("this is {{ not valid {{ rust".into())],
1407            ..RustFn::default()
1408        })];
1409        let rendered = render_items(&items);
1410        assert!(
1411            rendered.contains("WARNING"),
1412            "Invalid HIR should produce a warning in output"
1413        );
1414    }
1415}