1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use crate::dialect::operation::{OperationElement, OperationField, VariadicKind};
use proc_macro2::{Span, TokenStream};
use quote::quote;
use syn::Ident;

pub fn generate_element_getter(
    field: &(impl OperationField + OperationElement),
    singular_kind: &str,
    plural_kind: &str,
    error_variant: &Ident,
    index: usize,
    length: usize,
) -> TokenStream {
    let singular_kind_identifier = Ident::new(singular_kind, Span::call_site());
    let plural_kind_identifier = Ident::new(plural_kind, Span::call_site());
    let count = Ident::new(&format!("{singular_kind}_count"), Span::call_site());
    let name = field.name();

    let body = match field.variadic_kind() {
        VariadicKind::Simple { unfixed_seen } => {
            if field.is_optional() {
                // Optional element, and some singular elements.
                // Only present if the amount of groups is at least the number of
                // elements.
                quote! {
                    if self.operation.#count() < #length {
                        Err(::melior::Error::#error_variant(#name))
                    } else {
                        self.operation.#singular_kind_identifier(#index)
                    }
                }
            } else if field.is_variadic() {
                // A unfixed group
                // Length computed by subtracting the amount of other
                // singular elements from the number of elements.
                quote! {
                    let group_length = self.operation.#count() - #length + 1;
                    self.operation.#plural_kind_identifier().skip(#index).take(group_length)
                }
            } else if *unfixed_seen {
                // Single element after unfixed group
                // Compute the length of that variable group and take the next element
                quote! {
                    let group_length = self.operation.#count() - #length + 1;
                    self.operation.#singular_kind_identifier(#index + group_length - 1)
                }
            } else {
                // All elements so far are singular
                quote! {
                    self.operation.#singular_kind_identifier(#index)
                }
            }
        }
        VariadicKind::SameSize {
            unfixed_count,
            preceding_simple_count,
            preceding_variadic_count,
        } => {
            let get_elements = if field.is_unfixed() {
                quote! {
                    self.operation.#plural_kind_identifier().skip(start).take(group_len)
                }
            } else {
                quote! {
                    self.operation.#singular_kind_identifier(start)
                }
            };

            quote! {
                let total_var_len = self.operation.#count() - #unfixed_count + 1;
                let group_len = total_var_len / #unfixed_count;
                let start = #preceding_simple_count + #preceding_variadic_count * group_len;

                #get_elements
            }
        }
        VariadicKind::AttributeSized => {
            let segment_size_attribute = format!("{singular_kind}_segment_sizes");
            let get_elements = if !field.is_unfixed() {
                quote! {
                    self.operation.#singular_kind_identifier(start)
                }
            } else if field.is_optional() {
                quote! {
                    if group_len == 0 {
                        Err(::melior::Error::#error_variant(#name))
                    } else {
                        self.operation.#singular_kind_identifier(start)
                    }
                }
            } else {
                quote! {
                    Ok(self.operation.#plural_kind_identifier().skip(start).take(group_len))
                }
            };

            quote! {
                let attribute =
                    ::melior::ir::attribute::DenseI32ArrayAttribute::<'c>::try_from(
                        self.operation
                        .attribute(#segment_size_attribute)?
                    )?;
                let start = (0..#index)
                    .map(|index| attribute.element(index))
                    .collect::<Result<Vec<_>, _>>()?
                    .into_iter()
                    .sum::<i32>() as usize;
                let group_len = attribute.element(#index)? as usize;

                #get_elements
            }
        }
    };

    let identifier = field.singular_identifier();
    let return_type = field.return_type();

    quote! {
        pub fn #identifier(&self) -> #return_type {
            #body
        }
    }
}