ref_mut_n/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::__private::Span;
4use syn::parse::{Parse, ParseStream};
5use syn::token::Comma;
6use syn::{parse_macro_input, Ident, LitInt};
7
8struct RefMutN {
9    referred: Ident,
10    n: usize,
11}
12
13impl Parse for RefMutN {
14    fn parse(input: ParseStream) -> Result<Self, syn::Error> {
15        let referred: Ident = input.parse()?;
16        input.parse::<Comma>()?;
17        let n = input.parse::<LitInt>()?.base10_parse::<usize>()?;
18
19        Ok(Self { referred, n })
20    }
21}
22
23/// Defines a nested mutable reference.
24///
25/// # Example
26///
27/// ```
28/// use ref_mut_n::ref_mut_n;
29///
30/// let mut data = [[0u8; 3]; 4];
31/// let _ref_data: &mut [&mut [u8]] = ref_mut_n!(data, 4);
32/// ```
33#[proc_macro]
34pub fn ref_mut_n(input: TokenStream) -> TokenStream {
35    let RefMutN { referred, n } = parse_macro_input!(input as RefMutN);
36
37    let assign = (0..n)
38        .into_iter()
39        .map(|n| Ident::new(&format!("a{}", n), Span::call_site()))
40        .collect::<Vec<_>>();
41
42    let token = quote! {
43        {
44            let [#(ref mut #assign),*] = #referred;
45            &mut [#(#assign),*]
46        }
47    };
48    token.into()
49}