melior_macro/dialect/generation/
element_accessor.rs

1use crate::dialect::operation::{OperationElement, VariadicKind};
2use proc_macro2::{Span, TokenStream};
3use quote::quote;
4use syn::Ident;
5
6pub fn generate_element_getter(
7    field: &impl OperationElement,
8    singular_kind: &str,
9    plural_kind: &str,
10    error_variant: &Ident,
11    index: usize,
12    length: usize,
13) -> TokenStream {
14    let singular_kind_identifier = Ident::new(singular_kind, Span::call_site());
15    let plural_kind_identifier = Ident::new(plural_kind, Span::call_site());
16    let count = Ident::new(&format!("{singular_kind}_count"), Span::call_site());
17    let name = field.name();
18
19    let body = match field.variadic_kind() {
20        VariadicKind::Simple { unfixed_seen } => {
21            if field.is_optional() {
22                // Optional element, and some singular elements.
23                // Only present if the amount of groups is at least the number of
24                // elements.
25                quote! {
26                    if self.operation.#count() < #length {
27                        Err(::melior::Error::#error_variant(#name))
28                    } else {
29                        self.operation.#singular_kind_identifier(#index)
30                    }
31                }
32            } else if field.is_variadic() {
33                // A unfixed group
34                // Length computed by subtracting the amount of other
35                // singular elements from the number of elements.
36                quote! {
37                    let group_length = self.operation.#count() - #length + 1;
38                    self.operation.#plural_kind_identifier().skip(#index).take(group_length)
39                }
40            } else if *unfixed_seen {
41                // Single element after unfixed group
42                // Compute the length of that variable group and take the next element
43                quote! {
44                    let group_length = self.operation.#count() - #length + 1;
45                    self.operation.#singular_kind_identifier(#index + group_length - 1)
46                }
47            } else {
48                // All elements so far are singular
49                quote! {
50                    self.operation.#singular_kind_identifier(#index)
51                }
52            }
53        }
54        VariadicKind::SameSize {
55            unfixed_count,
56            preceding_simple_count,
57            preceding_variadic_count,
58        } => {
59            let get_elements = if field.is_unfixed() {
60                quote! {
61                    self.operation.#plural_kind_identifier().skip(start).take(group_len)
62                }
63            } else {
64                quote! {
65                    self.operation.#singular_kind_identifier(start)
66                }
67            };
68
69            quote! {
70                let total_var_len = self.operation.#count() - #unfixed_count + 1;
71                let group_len = total_var_len / #unfixed_count;
72                let start = #preceding_simple_count + #preceding_variadic_count * group_len;
73
74                #get_elements
75            }
76        }
77        VariadicKind::AttributeSized => {
78            let segment_size_attribute = format!("{singular_kind}_segment_sizes");
79            let get_elements = if !field.is_unfixed() {
80                quote! {
81                    self.operation.#singular_kind_identifier(start)
82                }
83            } else if field.is_optional() {
84                quote! {
85                    if group_len == 0 {
86                        Err(::melior::Error::#error_variant(#name))
87                    } else {
88                        self.operation.#singular_kind_identifier(start)
89                    }
90                }
91            } else {
92                quote! {
93                    Ok(self.operation.#plural_kind_identifier().skip(start).take(group_len))
94                }
95            };
96
97            quote! {
98                let attribute =
99                    ::melior::ir::attribute::DenseI32ArrayAttribute::<'c>::try_from(
100                        self.operation
101                        .attribute(#segment_size_attribute)?
102                    )?;
103                let start = (0..#index)
104                    .map(|index| attribute.element(index))
105                    .collect::<Result<Vec<_>, _>>()?
106                    .into_iter()
107                    .sum::<i32>() as usize;
108                let group_len = attribute.element(#index)? as usize;
109
110                #get_elements
111            }
112        }
113    };
114
115    let identifier = field.singular_identifier();
116    let return_type = field.return_type();
117
118    quote! {
119        pub fn #identifier(&self) -> #return_type {
120            #body
121        }
122    }
123}
melior_macro/dialect/generation/element_accessor.rs

melior_macro/dialect/generation/
element_accessor.rs
