Skip to main content

cairo_lang_plugins/plugins/
panicable.rs

1use cairo_lang_defs::patcher::{PatchBuilder, RewriteNode};
2use cairo_lang_defs::plugin::{
3    MacroPlugin, MacroPluginMetadata, PluginDiagnostic, PluginGeneratedFile, PluginResult,
4};
5use cairo_lang_filesystem::ids::SmolStrId;
6use cairo_lang_syntax::attribute::structured::{
7    Attribute, AttributeArg, AttributeArgVariant, AttributeStructurize,
8};
9use cairo_lang_syntax::node::helpers::{GetIdentifier, QueryAttrs};
10use cairo_lang_syntax::node::{TypedSyntaxNode, ast};
11use cairo_lang_utils::try_extract_matches;
12use indoc::formatdoc;
13use itertools::Itertools;
14use salsa::Database;
15
16#[derive(Debug, Default)]
17#[non_exhaustive]
18pub struct PanicablePlugin;
19
20const PANIC_WITH_ATTR: &str = "panic_with";
21/// The feature gating the deprecated `#[panic_with]` macro. Stored with quotes, as
22/// `allowed_features` holds the raw string-literal text.
23const DEPRECATED_PANIC_WITH_FEATURE: &str = r#""deprecated-panic-with""#;
24
25impl MacroPlugin for PanicablePlugin {
26    fn generate_code<'db>(
27        &self,
28        db: &'db dyn Database,
29        item_ast: ast::ModuleItem<'db>,
30        metadata: &MacroPluginMetadata<'_>,
31    ) -> PluginResult<'db> {
32        let (declaration, attributes, visibility) = match item_ast {
33            ast::ModuleItem::ExternFunction(extern_func_ast) => (
34                extern_func_ast.declaration(db),
35                extern_func_ast.attributes(db),
36                extern_func_ast.visibility(db),
37            ),
38            ast::ModuleItem::FreeFunction(free_func_ast) => (
39                free_func_ast.declaration(db),
40                free_func_ast.attributes(db),
41                free_func_ast.visibility(db),
42            ),
43            _ => return PluginResult::default(),
44        };
45
46        generate_panicable_code(db, declaration, attributes, visibility, metadata)
47    }
48
49    fn declared_attributes<'db>(&self, db: &'db dyn Database) -> Vec<SmolStrId<'db>> {
50        vec![SmolStrId::from(db, PANIC_WITH_ATTR)]
51    }
52}
53
54/// Generate code defining a panicable variant of a function marked with `#[panic_with]` attribute.
55fn generate_panicable_code<'db>(
56    db: &'db dyn Database,
57    declaration: ast::FunctionDeclaration<'db>,
58    attributes: ast::AttributeList<'db>,
59    visibility: ast::Visibility<'db>,
60    metadata: &MacroPluginMetadata<'_>,
61) -> PluginResult<'db> {
62    let mut attrs = attributes.query_attr(db, PANIC_WITH_ATTR);
63    let Some(attr) = attrs.next() else {
64        // No `#[panic_with]` attribute found.
65        return PluginResult::default();
66    };
67    let mut diagnostics = vec![];
68    let deprecation_feature = SmolStrId::from(db, DEPRECATED_PANIC_WITH_FEATURE);
69    if !metadata.allowed_features.contains(&deprecation_feature) {
70        diagnostics.push(PluginDiagnostic::warning(
71            attr.stable_ptr(db),
72            format!(
73                "Usage of deprecated macro `{PANIC_WITH_ATTR}` with no \
74                 `#[feature({DEPRECATED_PANIC_WITH_FEATURE})]` attribute. Use simple calculations \
75                 instead, as these are supported in const context.",
76            ),
77        ));
78    }
79    if let Some(extra_attr) = attrs.next() {
80        diagnostics.push(PluginDiagnostic::error(
81            extra_attr.stable_ptr(db),
82            "`#[panic_with]` cannot be applied multiple times to the same item.".into(),
83        ));
84        return PluginResult { code: None, diagnostics, remove_original_item: false };
85    }
86
87    let signature = declaration.signature(db);
88    let Some((inner_ty, success_variant, failure_variant)) =
89        extract_success_ty_and_variants(db, &signature)
90    else {
91        diagnostics.push(PluginDiagnostic::error(
92            signature.ret_ty(db).stable_ptr(db),
93            "Currently only wrapping functions returning an Option<T> or Result<T, E>".into(),
94        ));
95        return PluginResult { code: None, diagnostics, remove_original_item: false };
96    };
97
98    let mut builder = PatchBuilder::new(db, &attr);
99    let attr = attr.structurize(db);
100
101    let Some((err_value, panicable_name)) = parse_arguments(db, &attr) else {
102        diagnostics.push(PluginDiagnostic::error(
103            attr.stable_ptr,
104            "Failed to extract panic data attribute".into(),
105        ));
106        return PluginResult { code: None, diagnostics, remove_original_item: false };
107    };
108    builder.add_node(visibility.as_syntax_node());
109    builder.add_node(declaration.function_kw(db).as_syntax_node());
110    builder.add_modified(RewriteNode::from_ast_trimmed(&panicable_name));
111    builder.add_node(declaration.generic_params(db).as_syntax_node());
112    builder.add_node(signature.lparen(db).as_syntax_node());
113    builder.add_node(signature.parameters(db).as_syntax_node());
114    builder.add_node(signature.rparen(db).as_syntax_node());
115    let args = signature
116        .parameters(db)
117        .elements(db)
118        .map(|param| {
119            let ref_kw =
120                if let Ok(ast::Modifier::Ref(_)) = param.modifiers(db).elements(db).exactly_one() {
121                    "ref "
122                } else {
123                    ""
124                };
125            format!("{}{}", ref_kw, param.name(db).as_syntax_node().get_text(db))
126        })
127        .join(", ");
128    builder.add_modified(RewriteNode::interpolate_patched(
129        &formatdoc!(
130            r#"
131                -> $inner_ty$ {{
132                    match $function_name$({args}) {{
133                        {success_variant}(v) => v,
134                        {failure_variant}(_) => core::panic_with_const_felt252::<$err_value$>(),
135                    }}
136                }}
137            "#
138        ),
139        &[
140            ("inner_ty".to_string(), RewriteNode::from_ast_trimmed(&inner_ty)),
141            ("function_name".to_string(), RewriteNode::from_ast_trimmed(&declaration.name(db))),
142            ("err_value".to_string(), RewriteNode::from_ast_trimmed(&err_value)),
143        ]
144        .into(),
145    ));
146
147    let (content, code_mappings) = builder.build();
148    PluginResult {
149        code: Some(PluginGeneratedFile {
150            name: "panicable".into(),
151            content,
152            code_mappings,
153            aux_data: None,
154            diagnostics_note: Default::default(),
155            is_unhygienic: false,
156        }),
157        diagnostics,
158        remove_original_item: false,
159    }
160}
161
162/// Given a function signature, if it returns `Option::<T>` or `Result::<T, E>`, returns T and the
163/// variant match strings. Otherwise, returns None.
164fn extract_success_ty_and_variants<'a>(
165    db: &'a dyn Database,
166    signature: &ast::FunctionSignature<'a>,
167) -> Option<(ast::GenericArg<'a>, String, String)> {
168    let ret_ty_expr =
169        try_extract_matches!(signature.ret_ty(db), ast::OptionReturnTypeClause::ReturnTypeClause)?
170            .ty(db);
171    let ret_ty_path = try_extract_matches!(ret_ty_expr, ast::Expr::Path)?;
172
173    // Currently wrapping functions returning an Option<T> or Result<T, E>.
174    let Ok(ast::PathSegment::WithGenericArgs(segment)) =
175        ret_ty_path.segments(db).elements(db).exactly_one()
176    else {
177        return None;
178    };
179    let ty = segment.identifier(db).long(db);
180    if ty == "Option" {
181        let inner = segment.generic_args(db).generic_args(db).elements(db).exactly_one().ok()?;
182        Some((inner.clone(), "Option::Some".to_owned(), "Option::None".to_owned()))
183    } else if ty == "Result" {
184        let [inner, _err] =
185            segment.generic_args(db).generic_args(db).elements(db).collect_array()?;
186        Some((inner.clone(), "Result::Ok".to_owned(), "Result::Err".to_owned()))
187    } else {
188        None
189    }
190}
191
192/// Parse `#[panic_with(...)]` attribute arguments and return a tuple with error value and
193/// panicable function name.
194fn parse_arguments<'a>(
195    db: &'a dyn Database,
196    attr: &Attribute<'a>,
197) -> Option<(ast::TerminalShortString<'a>, ast::TerminalIdentifier<'a>)> {
198    let [
199        AttributeArg {
200            variant: AttributeArgVariant::Unnamed(ast::Expr::ShortString(err_value)),
201            ..
202        },
203        AttributeArg { variant: AttributeArgVariant::Unnamed(ast::Expr::Path(name)), .. },
204    ] = &attr.args[..]
205    else {
206        return None;
207    };
208
209    let Ok(ast::PathSegment::Simple(segment)) = name.segments(db).elements(db).exactly_one() else {
210        return None;
211    };
212
213    Some((err_value.clone(), segment.ident(db)))
214}