amalie_macros/
lib.rs

1extern crate proc_macro;
2use proc_macro::TokenStream;
3use quote::quote;
4use syn::{parse_macro_input, Expr, ExprLit, Lit, visit_mut::VisitMut, ExprArray, Token};
5use syn::parse::{Parse, ParseStream};
6
7#[proc_macro]
8pub fn zz(input: TokenStream) -> TokenStream {
9    let mut ast = parse_macro_input!(input as Expr);
10    ZZTransformer.visit_expr_mut(&mut ast);
11    TokenStream::from(quote! { #ast })
12}
13
14
15struct ZZTransformer;
16
17impl VisitMut for ZZTransformer {
18    fn visit_expr_mut(&mut self, node: &mut Expr) {
19        match node {
20            Expr::Lit(ExprLit { lit: Lit::Int(lit_int), .. }) => {
21                let value = lit_int.base10_digits();
22                if value == "1" {
23                    *node = syn::parse_quote! {
24                       ZZ::one()
25                    };
26                }
27                else if value == "0" {
28                    *node = syn::parse_quote! {
29                       ZZ::zero()
30                    };
31                }
32                else {
33                    *node = syn::parse_quote! {
34                       ZZ::zz_from_str(#value).unwrap()
35                    };
36                }
37            },
38            _ => syn::visit_mut::visit_expr_mut(self, node),
39        }
40    }
41}
42
43struct VecZZInput {
44    values: Vec<Expr>,
45}
46
47impl Parse for VecZZInput {
48    fn parse(input: ParseStream) -> syn::Result<Self> {
49        let mut values = Vec::new();
50        while !input.is_empty() {
51            let value: Expr = input.parse()?;
52            if !input.is_empty() {
53                let _: Token![,] = input.parse()?;
54            }
55            values.push(value);
56        }
57        Ok(VecZZInput { values })
58    }
59}
60
61#[proc_macro]
62pub fn vec_zz(input: TokenStream) -> TokenStream {
63    let VecZZInput { mut values } = parse_macro_input!(input as VecZZInput);
64
65    let vector = values.iter_mut().map(|value| {
66        quote! {
67            zz!(#value)
68        }
69    });
70
71    let output = quote! {
72        vec![#(#vector),*]
73    };
74    output.into()
75}
76
77struct MatrixInput {
78    rows: Vec<ExprArray>,
79}
80
81impl Parse for MatrixInput {
82    fn parse(input: ParseStream) -> syn::Result<Self> {
83        let mut rows = Vec::new();
84        while !input.is_empty() {
85            let row: ExprArray = input.parse()?;
86            if !input.is_empty() {
87                let _: Token![,] = input.parse()?;
88            }
89            rows.push(row);
90        }
91        Ok(MatrixInput { rows })
92    }
93}
94
95#[proc_macro]
96pub fn matrix(input: TokenStream) -> TokenStream {
97    let MatrixInput { mut rows } = parse_macro_input!(input as MatrixInput);
98
99    let matrix = rows.iter_mut().map(|row| {
100        let elements = row.elems.iter_mut().map(|mut num| {
101            ZZTransformer.visit_expr_mut(&mut num);
102            quote! {
103                #num
104            }
105        });
106        quote! {
107            vec![#(#elements),*]
108        }
109    });
110
111    let output = quote! {
112        {
113            let m: Matrix = vec![#(zz!(#matrix)),*].try_into().unwrap(); 
114            m 
115        }
116    };
117    output.into()
118}
119
amalie_macros/lib.rs

amalie_macros/
lib.rs
