unsafe_access/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::{parse_macro_input, visit_mut::VisitMut, Expr, ExprArray, ExprIndex, ItemFn};
6
7#[proc_macro]
8pub fn unchecked_indices_ref(input: TokenStream) -> TokenStream {
9    let expr = parse_macro_input!(input as ExprArray);
10
11    let transformed_elements = expr.elems.into_iter().map(|element| {
12        if let Expr::Index(ExprIndex { expr, index, .. }) = element {
13            quote! {
14                unsafe { #expr.get_unchecked(#index) }
15            }
16        } else {
17            quote! {
18                #element
19            }
20        }
21    });
22
23    let output = quote! {
24        [#(#transformed_elements),*]
25    };
26
27    TokenStream::from(output)
28}
29
30#[proc_macro]
31pub fn unchecked_indices_clone(input: TokenStream) -> TokenStream {
32    let expr = parse_macro_input!(input as ExprArray);
33
34    let transformed_elements = expr.elems.into_iter().map(|element| {
35        if let Expr::Index(ExprIndex { expr, index, .. }) = element {
36            quote! {
37                unsafe { #expr.get_unchecked(#index).clone() }
38            }
39        } else {
40            quote! {
41                #element
42            }
43        }
44    });
45
46    let output = quote! {
47        [#(#transformed_elements),*]
48    };
49
50    TokenStream::from(output)
51}
52
53#[proc_macro]
54pub fn unchecked_indices_copy(input: TokenStream) -> TokenStream {
55    let expr = parse_macro_input!(input as ExprArray);
56
57    let transformed_elements = expr.elems.into_iter().map(|element| {
58        if let Expr::Index(ExprIndex { expr, index, .. }) = element {
59            quote! {
60                unsafe { *#expr.get_unchecked(#index) }
61            }
62        } else {
63            quote! {
64                #element
65            }
66        }
67    });
68
69    let output = quote! {
70            [#(#transformed_elements),*]
71    };
72
73    TokenStream::from(output)
74}
75
76struct IndexAccessVisitor;
77
78impl VisitMut for IndexAccessVisitor {
79    fn visit_expr_mut(&mut self, expr: &mut Expr) {
80        // Use a recursive approach to handle nested expressions
81        if let Expr::Index(index_expr) = expr {
82            let expr_clone = *index_expr.expr.clone();
83            let index_clone = *index_expr.index.clone();
84
85            // Replace the indexed expression with a .get_unchecked() call inside an unsafe block
86            *expr = syn::parse_quote! {
87                unsafe { *#expr_clone.get_unchecked(#index_clone) }
88            };
89        } else {
90            // Continue visiting recursively
91            syn::visit_mut::visit_expr_mut(self, expr);
92        }
93    }
94}
95
96#[proc_macro_attribute]
97pub fn unsafe_access_fn(_attr: TokenStream, input: TokenStream) -> TokenStream {
98    let mut func = parse_macro_input!(input as ItemFn);
99
100    // Ensure the function is marked unsafe
101    if func.sig.unsafety.is_none() {
102        return syn::Error::new_spanned(&func.sig.fn_token, "Function must be declared unsafe")
103            .to_compile_error()
104            .into();
105    }
106
107    let mut visitor = IndexAccessVisitor;
108    visitor.visit_item_fn_mut(&mut func);
109
110    let output = quote! {
111        #func
112    };
113
114    TokenStream::from(output)
115}
unsafe_access/lib.rs

unsafe_access/
lib.rs
