nalgebra_macros/
stack_impl.rs

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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
use crate::Matrix;
use proc_macro2::{Span, TokenStream as TokenStream2};
use quote::{format_ident, quote, quote_spanned};
use syn::spanned::Spanned;
use syn::{Error, Expr, Lit};

#[allow(clippy::too_many_lines)]
pub fn stack_impl(matrix: Matrix) -> syn::Result<TokenStream2> {
    // The prefix is used to construct variable names
    // that are extremely unlikely to collide with variable names used in e.g. expressions
    // by the user. Although we could use a long, pseudo-random string, this makes the generated
    // code very painful to parse, so we settle for something more semantic that is still
    // very unlikely to collide
    let prefix = "___na";
    let n_block_rows = matrix.nrows();
    let n_block_cols = matrix.ncols();

    let mut output = quote! {};

    // First assign data and shape for each matrix entry to variables
    // (this is important so that we, for example, don't evaluate an expression more than once)
    for i in 0..n_block_rows {
        for j in 0..n_block_cols {
            let expr = &matrix[(i, j)];
            if !is_literal_zero(expr) {
                let ident_block = format_ident!("{prefix}_stack_{i}_{j}_block");
                let ident_shape = format_ident!("{prefix}_stack_{i}_{j}_shape");
                output.extend(std::iter::once(quote_spanned! {expr.span()=>
                    let ref #ident_block = #expr;
                    let #ident_shape = #ident_block.shape_generic();
                }));
            }
        }
    }

    // Determine the number of rows (dimension) in each block row,
    // and write out variables that define block row dimensions and offsets into the
    // output matrix
    for i in 0..n_block_rows {
        // The dimension of the block row is the result of trying to unify the row shape of
        // all blocks in the block row
        let dim = (0 ..n_block_cols)
            .filter_map(|j| {
                let expr = &matrix[(i, j)];
                if !is_literal_zero(expr) {
                    let mut ident_shape = format_ident!("{prefix}_stack_{i}_{j}_shape");
                    ident_shape.set_span(ident_shape.span().located_at(expr.span()));
                    Some(quote_spanned!{expr.span()=> #ident_shape.0 })
                } else {
                    None
                }
            }).reduce(|a, b| {
                let expect_msg = format!("All blocks in block row {i} must have the same number of rows");
                quote_spanned!{b.span()=>
                    <nalgebra::constraint::ShapeConstraint as nalgebra::constraint::SameNumberOfRows<_, _>>::representative(#a, #b)
                        .expect(#expect_msg)
                }
            }).ok_or(Error::new(Span::call_site(), format!("Block row {i} cannot consist entirely of implicit zero blocks.")))?;

        let dim_ident = format_ident!("{prefix}_stack_row_{i}_dim");
        let offset_ident = format_ident!("{prefix}_stack_row_{i}_offset");

        let offset = if i == 0 {
            quote! { 0 }
        } else {
            let prev_offset_ident = format_ident!("{prefix}_stack_row_{}_offset", i - 1);
            let prev_dim_ident = format_ident!("{prefix}_stack_row_{}_dim", i - 1);
            quote! { #prev_offset_ident + <_ as nalgebra::Dim>::value(&#prev_dim_ident) }
        };

        output.extend(std::iter::once(quote! {
            let #dim_ident = #dim;
            let #offset_ident = #offset;
        }));
    }

    // Do the same thing for the block columns
    for j in 0..n_block_cols {
        let dim = (0 ..n_block_rows)
            .filter_map(|i| {
                let expr = &matrix[(i, j)];
                if !is_literal_zero(expr) {
                    let mut ident_shape = format_ident!("{prefix}_stack_{i}_{j}_shape");
                    ident_shape.set_span(ident_shape.span().located_at(expr.span()));
                    Some(quote_spanned!{expr.span()=> #ident_shape.1 })
                } else {
                    None
                }
            }).reduce(|a, b| {
                let expect_msg = format!("All blocks in block column {j} must have the same number of columns");
                quote_spanned!{b.span()=>
                        <nalgebra::constraint::ShapeConstraint as nalgebra::constraint::SameNumberOfColumns<_, _>>::representative(#a, #b)
                            .expect(#expect_msg)
                }
            }).ok_or(Error::new(Span::call_site(), format!("Block column {j} cannot consist entirely of implicit zero blocks.")))?;

        let dim_ident = format_ident!("{prefix}_stack_col_{j}_dim");
        let offset_ident = format_ident!("{prefix}_stack_col_{j}_offset");

        let offset = if j == 0 {
            quote! { 0 }
        } else {
            let prev_offset_ident = format_ident!("{prefix}_stack_col_{}_offset", j - 1);
            let prev_dim_ident = format_ident!("{prefix}_stack_col_{}_dim", j - 1);
            quote! { #prev_offset_ident + <_ as nalgebra::Dim>::value(&#prev_dim_ident) }
        };

        output.extend(std::iter::once(quote! {
            let #dim_ident = #dim;
            let #offset_ident = #offset;
        }));
    }

    // Determine number of rows and cols in output matrix,
    // by adding together dimensions of all block rows/cols
    let num_rows = (0..n_block_rows)
        .map(|i| {
            let ident = format_ident!("{prefix}_stack_row_{i}_dim");
            quote! { #ident }
        })
        .reduce(|a, b| {
            quote! {
                <_ as nalgebra::DimAdd<_>>::add(#a, #b)
            }
        })
        .unwrap_or(quote! { nalgebra::dimension::U0 });

    let num_cols = (0..n_block_cols)
        .map(|j| {
            let ident = format_ident!("{prefix}_stack_col_{j}_dim");
            quote! { #ident }
        })
        .reduce(|a, b| {
            quote! {
                <_ as nalgebra::DimAdd<_>>::add(#a, #b)
            }
        })
        .unwrap_or(quote! { nalgebra::dimension::U0 });

    // It should be possible to use `uninitialized_generic` here instead
    // however that would mean that the macro needs to generate unsafe code
    // which does not seem like a great idea.
    output.extend(std::iter::once(quote! {
        let mut matrix = nalgebra::Matrix::zeros_generic(#num_rows, #num_cols);
    }));

    for i in 0..n_block_rows {
        for j in 0..n_block_cols {
            let row_dim = format_ident!("{prefix}_stack_row_{i}_dim");
            let col_dim = format_ident!("{prefix}_stack_col_{j}_dim");
            let row_offset = format_ident!("{prefix}_stack_row_{i}_offset");
            let col_offset = format_ident!("{prefix}_stack_col_{j}_offset");
            let expr = &matrix[(i, j)];
            if !is_literal_zero(expr) {
                let expr_ident = format_ident!("{prefix}_stack_{i}_{j}_block");
                output.extend(std::iter::once(quote! {
                    let start = (#row_offset, #col_offset);
                    let shape = (#row_dim, #col_dim);
                    let input_view = #expr_ident.generic_view((0, 0), shape);
                    let mut output_view = matrix.generic_view_mut(start, shape);
                    output_view.copy_from(&input_view);
                }));
            }
        }
    }

    Ok(quote! {
        {
            #output
            matrix
        }
    })
}

fn is_literal_zero(expr: &Expr) -> bool {
    matches!(expr,
        Expr::Lit(syn::ExprLit { lit: Lit::Int(integer_literal), .. })
        if integer_literal.base10_digits() == "0")
}

#[cfg(test)]
mod tests {
    use crate::stack_impl::stack_impl;
    use crate::Matrix;
    use quote::quote;

    #[test]
    fn stack_simple_generation() {
        let input: Matrix = syn::parse_quote![
            a, 0;
            0, b;
        ];

        let result = stack_impl(input).unwrap();

        let expected = quote! {{
            let ref ___na_stack_0_0_block = a;
            let ___na_stack_0_0_shape = ___na_stack_0_0_block.shape_generic();
            let ref ___na_stack_1_1_block = b;
            let ___na_stack_1_1_shape = ___na_stack_1_1_block.shape_generic();
            let ___na_stack_row_0_dim = ___na_stack_0_0_shape.0;
            let ___na_stack_row_0_offset = 0;
            let ___na_stack_row_1_dim = ___na_stack_1_1_shape.0;
            let ___na_stack_row_1_offset = ___na_stack_row_0_offset + <_ as nalgebra::Dim>::value(&___na_stack_row_0_dim);
            let ___na_stack_col_0_dim = ___na_stack_0_0_shape.1;
            let ___na_stack_col_0_offset = 0;
            let ___na_stack_col_1_dim = ___na_stack_1_1_shape.1;
            let ___na_stack_col_1_offset = ___na_stack_col_0_offset + <_ as nalgebra::Dim>::value(&___na_stack_col_0_dim);
            let mut matrix = nalgebra::Matrix::zeros_generic(
                <_ as nalgebra::DimAdd<_>>::add(___na_stack_row_0_dim, ___na_stack_row_1_dim),
                <_ as nalgebra::DimAdd<_>>::add(___na_stack_col_0_dim, ___na_stack_col_1_dim)
            );
            let start = (___na_stack_row_0_offset, ___na_stack_col_0_offset);
            let shape = (___na_stack_row_0_dim, ___na_stack_col_0_dim);
            let input_view = ___na_stack_0_0_block.generic_view((0,0), shape);
            let mut output_view = matrix.generic_view_mut(start, shape);
            output_view.copy_from(&input_view);
            let start = (___na_stack_row_1_offset, ___na_stack_col_1_offset);
            let shape = (___na_stack_row_1_dim, ___na_stack_col_1_dim);
            let input_view = ___na_stack_1_1_block.generic_view((0,0), shape);
            let mut output_view = matrix.generic_view_mut(start, shape);
            output_view.copy_from(&input_view);
            matrix
        }};

        assert_eq!(format!("{result}"), format!("{}", expected));
    }

    #[test]
    fn stack_complex_generation() {
        let input: Matrix = syn::parse_quote![
            a, 0, b;
            0, c, d;
            e, 0, 0;
        ];

        let result = stack_impl(input).unwrap();

        let expected = quote! {{
            let ref ___na_stack_0_0_block = a;
            let ___na_stack_0_0_shape = ___na_stack_0_0_block.shape_generic();
            let ref ___na_stack_0_2_block = b;
            let ___na_stack_0_2_shape = ___na_stack_0_2_block.shape_generic();
            let ref ___na_stack_1_1_block = c;
            let ___na_stack_1_1_shape = ___na_stack_1_1_block.shape_generic();
            let ref ___na_stack_1_2_block = d;
            let ___na_stack_1_2_shape = ___na_stack_1_2_block.shape_generic();
            let ref ___na_stack_2_0_block = e;
            let ___na_stack_2_0_shape = ___na_stack_2_0_block.shape_generic();
            let ___na_stack_row_0_dim = < nalgebra :: constraint :: ShapeConstraint as nalgebra :: constraint :: SameNumberOfRows < _ , _ >> :: representative (___na_stack_0_0_shape . 0 , ___na_stack_0_2_shape . 0) . expect ("All blocks in block row 0 must have the same number of rows") ;
            let ___na_stack_row_0_offset = 0;
            let ___na_stack_row_1_dim = < nalgebra :: constraint :: ShapeConstraint as nalgebra :: constraint :: SameNumberOfRows < _ , _ >> :: representative (___na_stack_1_1_shape . 0 , ___na_stack_1_2_shape . 0) . expect ("All blocks in block row 1 must have the same number of rows") ;
            let ___na_stack_row_1_offset = ___na_stack_row_0_offset + <_ as nalgebra::Dim>::value(&___na_stack_row_0_dim);
            let ___na_stack_row_2_dim = ___na_stack_2_0_shape.0;
            let ___na_stack_row_2_offset = ___na_stack_row_1_offset + <_ as nalgebra::Dim>::value(&___na_stack_row_1_dim);
            let ___na_stack_col_0_dim = < nalgebra :: constraint :: ShapeConstraint as nalgebra :: constraint :: SameNumberOfColumns < _ , _ >> :: representative (___na_stack_0_0_shape . 1 , ___na_stack_2_0_shape . 1) . expect ("All blocks in block column 0 must have the same number of columns") ;
            let ___na_stack_col_0_offset = 0;
            let ___na_stack_col_1_dim = ___na_stack_1_1_shape.1;
            let ___na_stack_col_1_offset = ___na_stack_col_0_offset + <_ as nalgebra::Dim>::value(&___na_stack_col_0_dim);
            let ___na_stack_col_2_dim = < nalgebra :: constraint :: ShapeConstraint as nalgebra :: constraint :: SameNumberOfColumns < _ , _ >> :: representative (___na_stack_0_2_shape . 1 , ___na_stack_1_2_shape . 1) . expect ("All blocks in block column 2 must have the same number of columns") ;
            let ___na_stack_col_2_offset = ___na_stack_col_1_offset + <_ as nalgebra::Dim>::value(&___na_stack_col_1_dim);
            let mut matrix = nalgebra::Matrix::zeros_generic(
                <_ as nalgebra::DimAdd<_>>::add(
                    <_ as nalgebra::DimAdd<_>>::add(___na_stack_row_0_dim, ___na_stack_row_1_dim),
                    ___na_stack_row_2_dim
                ),
                <_ as nalgebra::DimAdd<_>>::add(
                    <_ as nalgebra::DimAdd<_>>::add(___na_stack_col_0_dim, ___na_stack_col_1_dim),
                    ___na_stack_col_2_dim
                )
            );
            let start = (___na_stack_row_0_offset, ___na_stack_col_0_offset);
            let shape = (___na_stack_row_0_dim, ___na_stack_col_0_dim);
            let input_view = ___na_stack_0_0_block.generic_view((0,0), shape);
            let mut output_view = matrix.generic_view_mut(start, shape);
            output_view.copy_from(&input_view);
            let start = (___na_stack_row_0_offset, ___na_stack_col_2_offset);
            let shape = (___na_stack_row_0_dim, ___na_stack_col_2_dim);
            let input_view = ___na_stack_0_2_block.generic_view((0,0), shape);
            let mut output_view = matrix.generic_view_mut(start, shape);
            output_view.copy_from(&input_view);
            let start = (___na_stack_row_1_offset, ___na_stack_col_1_offset);
            let shape = (___na_stack_row_1_dim, ___na_stack_col_1_dim);
            let input_view = ___na_stack_1_1_block.generic_view((0,0), shape);
            let mut output_view = matrix.generic_view_mut(start, shape);
            output_view.copy_from(&input_view);
            let start = (___na_stack_row_1_offset, ___na_stack_col_2_offset);
            let shape = (___na_stack_row_1_dim, ___na_stack_col_2_dim);
            let input_view = ___na_stack_1_2_block.generic_view((0,0), shape);
            let mut output_view = matrix.generic_view_mut(start, shape);
            output_view.copy_from(&input_view);
            let start = (___na_stack_row_2_offset, ___na_stack_col_0_offset);
            let shape = (___na_stack_row_2_dim, ___na_stack_col_0_dim);
            let input_view = ___na_stack_2_0_block.generic_view((0,0), shape);
            let mut output_view = matrix.generic_view_mut(start, shape);
            output_view.copy_from(&input_view);
            matrix
        }};

        assert_eq!(format!("{result}"), format!("{}", expected));
    }
}