Skip to main content

melior_macro/dialect/
utility.rs

1use super::error::Error;
2use comrak::{arena_tree::NodeEdge, format_commonmark, nodes::NodeValue, parse_document, Arena};
3use convert_case::{Case, Casing};
4use proc_macro2::Ident;
5use quote::format_ident;
6use syn::{parse_quote, Type};
7
8const RESERVED_NAMES: &[&str] = &["name", "operation", "builder"];
9
10pub fn generate_result_type(r#type: Type) -> Type {
11    parse_quote!(Result<#r#type, ::melior::Error>)
12}
13
14pub fn generate_iterator_type(r#type: Type) -> Type {
15    parse_quote!(impl Iterator<Item = #r#type>)
16}
17
18pub fn sanitize_snake_case_identifier(name: &str) -> Result<Ident, Error> {
19    sanitize_name(&name.to_case(Case::Snake))
20}
21
22fn sanitize_name(name: &str) -> Result<Ident, Error> {
23    // Replace any "." with "_".
24    let mut name = name.replace('.', "_");
25
26    // Add "_" suffix to avoid conflicts with existing methods.
27    if RESERVED_NAMES.contains(&name.as_str())
28        || name
29            .chars()
30            .next()
31            .ok_or_else(|| Error::InvalidIdentifier(name.clone()))?
32            .is_numeric()
33    {
34        name = format!("_{}", name);
35    }
36
37    // Try to parse the string as an ident, and prefix the identifier
38    // with "r#" if it is not a valid identifier.
39    Ok(syn::parse_str::<Ident>(&name).unwrap_or_else(|_| format_ident!("r#{}", name)))
40}
41
42pub fn sanitize_documentation(string: &str) -> Result<String, Error> {
43    let arena = Arena::new();
44    let node = parse_document(&arena, &unindent::unindent(string), &Default::default());
45
46    for node in node.traverse() {
47        let NodeEdge::Start(node) = node else {
48            continue;
49        };
50        let mut ast = node.data.borrow_mut();
51        let NodeValue::CodeBlock(block) = &mut ast.value else {
52            continue;
53        };
54
55        if block.info.is_empty() {
56            // Mark them not in Rust to prevent documentation tests.
57            block.info = "text".into();
58        }
59    }
60
61    let mut buffer = Vec::with_capacity(string.len());
62
63    format_commonmark(node, &Default::default(), &mut buffer)?;
64
65    Ok(String::from_utf8(buffer)?)
66}
67
68pub fn capitalize_string(string: &str) -> String {
69    if string.is_empty() {
70        "".into()
71    } else {
72        string[..1].to_uppercase() + &string[1..]
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79    use pretty_assertions::assert_eq;
80
81    #[test]
82    fn sanitize_name_with_dot() {
83        assert_eq!(
84            sanitize_snake_case_identifier("foo.bar").unwrap(),
85            "foo_bar"
86        );
87    }
88
89    #[test]
90    fn sanitize_name_with_dot_and_underscore() {
91        assert_eq!(
92            sanitize_snake_case_identifier("foo.bar_baz").unwrap(),
93            "foo_bar_baz"
94        );
95    }
96
97    #[test]
98    fn sanitize_reserved_name() {
99        assert_eq!(
100            sanitize_snake_case_identifier("builder").unwrap(),
101            "_builder"
102        );
103    }
104
105    #[test]
106    fn sanitize_code_block() {
107        assert_eq!(
108            &sanitize_documentation("```\nfoo\n```\n").unwrap(),
109            "``` text\nfoo\n```\n"
110        );
111    }
112
113    #[test]
114    fn sanitize_code_blocks() {
115        assert_eq!(
116            &sanitize_documentation("```\nfoo\n```\n\n```\nbar\n```\n").unwrap(),
117            "``` text\nfoo\n```\n\n``` text\nbar\n```\n"
118        );
119    }
120
121    #[test]
122    fn capitalize() {
123        assert_eq!(&capitalize_string("foo"), "Foo");
124    }
125}