use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, Data, DeriveInput, Error, Expr, ExprRange, Fields, Ident, Index as SynIndex, Lit, LitInt, RangeLimits, Type};
use std::collections::HashSet;
enum FieldAccess {
Named(Ident),
Tuple(SynIndex),
Array { ident: Ident, offset: usize },
}
fn collect_indices(input: &DeriveInput) -> Result<(Vec<(usize, FieldAccess)>, Type), Error> {
if let Data::Struct(data) = &input.data {
if let Fields::Named(fields) = &data.fields {
let mut index_field_pairs = Vec::new();
let mut indices = HashSet::new();
let mut types = Vec::new();
for field in &fields.named {
let ident = field.ident.as_ref().unwrap().clone();
for attr in &field.attrs {
if !attr.path().is_ident("index") { continue; }
if let Ok(expr_range) = attr.parse_args::<ExprRange>() {
let start_expr = expr_range.start.as_ref().ok_or_else(||Error::new_spanned(&expr_range, "Missing start in range"))?;
let end_expr = expr_range.end.as_ref().ok_or_else(||Error::new_spanned(&expr_range, "Missing end in range"))?;
let start = if let Expr::Lit(lit_expr) = &**start_expr {
if let Lit::Int(li) = &lit_expr.lit {
li.base10_parse::<usize>()?
} else {
return Err(Error::new_spanned(lit_expr, "Start must be integer literal"));
}
} else {
return Err(Error::new_spanned(start_expr, "Unsupported start expression"));
};
let end_val = if let Expr::Lit(lit_expr) = &**end_expr {
if let Lit::Int(li) = &lit_expr.lit {
li.base10_parse::<usize>()?
} else {
return Err(Error::new_spanned(lit_expr, "End must be integer literal"));
}
} else {
return Err(Error::new_spanned(end_expr, "Unsupported end expression"));
};
let (end_exclusive, span_len) = match expr_range.limits {
RangeLimits::HalfOpen(_) => (
end_val,
end_val.checked_sub(start)
.ok_or_else(||Error::new_spanned(&expr_range, "End must be >= start"))?
),
RangeLimits::Closed(_) => (
end_val.checked_add(1)
.ok_or_else(||Error::new_spanned(&expr_range, "Range end overflow"))?,
end_val.checked_sub(start)
.and_then(|d| d.checked_add(1))
.ok_or_else(||Error::new_spanned(&expr_range, "Invalid inclusive range bounds"))?
),
};
let (elem_ty, arr_len) = if let Type::Array(arr) = &field.ty {
let len_expr = &arr.len;
if let Expr::Lit(lit_expr) = &*len_expr {
if let Lit::Int(li) = &lit_expr.lit {
((*arr.elem).clone(), li.base10_parse::<usize>()?)
} else {
return Err(Error::new_spanned(lit_expr, "Array length must be integer literal"));
}
} else {
return Err(Error::new_spanned(len_expr, "Unsupported array length expression"));
}
} else {
return Err(Error::new_spanned(&field.ty, "Ranges only allowed on arrays"));
};
if span_len != arr_len {
return Err(Error::new_spanned(&expr_range,
format!("Range covers {} elements but array has length {}", span_len, arr_len)
));
}
for idx in start..end_exclusive {
if !indices.insert(idx) {
return Err(Error::new_spanned(&attr, format!("Duplicate index {}", idx)));
}
index_field_pairs.push((idx, FieldAccess::Array { ident: ident.clone(), offset: start }));
types.push(elem_ty.clone());
}
}
else if let Ok(lit) = attr.parse_args::<LitInt>() {
let idx = lit.base10_parse::<usize>()?;
if !indices.insert(idx) {
return Err(Error::new_spanned(&lit, format!("Duplicate index {}", idx)));
}
types.push(field.ty.clone());
index_field_pairs.push((idx, FieldAccess::Named(ident.clone())));
} else {
return Err(Error::new_spanned(&attr, "Expected usize or range start..end or start..=end"));
}
}
}
if index_field_pairs.is_empty() {
return Err(Error::new_spanned(&input, "No #[index] attributes found"));
}
let first_ty = &types[0];
for ty in &types[1..] {
if ty != first_ty {
return Err(Error::new_spanned(&input, "Indexed fields must yield same element type"));
}
}
index_field_pairs.sort_by_key(|(i, _)| *i);
for expected in 0..index_field_pairs.len() {
if !indices.contains(&expected) {
return Err(Error::new_spanned(&input, format!("Missing index {}", expected)));
}
}
Ok((index_field_pairs, first_ty.clone()))
} else if let Fields::Unnamed(fields) = &data.fields {
let types_: Vec<_> = fields.unnamed.iter().map(|f| f.ty.clone()).collect();
let first_ty = types_.first().cloned().ok_or_else(|| Error::new_spanned(&input, "Tuple struct must have fields"))?;
for ty in &types_[1..] {
if ty != &first_ty {
return Err(Error::new_spanned(&input, "All tuple fields must have same type"));
}
}
let pairs = types_.into_iter().enumerate().map(|(i, _)| (i, FieldAccess::Tuple(SynIndex::from(i)))).collect();
Ok((pairs, first_ty))
} else {
Err(Error::new_spanned(&input, "Only structs with named or tuple fields supported"))
}
} else {
Err(Error::new_spanned(&input, "Index only derivable on structs"))
}
}
#[proc_macro_derive(Index, attributes(index))]
pub fn derive_index(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let name = &input.ident;
let (pairs, out_ty) = match collect_indices(&input) {
Ok(r) => r,
Err(e) => return e.to_compile_error().into(),
};
let arms = pairs.iter().map(|(i, fa)| match fa {
FieldAccess::Named(id) => quote! { #i => &self.#id, },
FieldAccess::Tuple(idx) => quote! { #i => &self.#idx, },
FieldAccess::Array { ident, offset } => quote! { #i => &self.#ident[#i - #offset], },
});
quote! {
impl std::ops::Index<usize> for #name {
type Output = #out_ty;
#[inline]
fn index(&self, index: usize) -> &Self::Output {
match index {
#(#arms)*
_ => panic!("Index out of bounds: {}", index),
}
}
}
}.into()
}
#[proc_macro_derive(IndexMut, attributes(index))]
pub fn derive_index_mut(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let name = &input.ident;
let (pairs, _) = match collect_indices(&input) {
Ok(r) => r,
Err(e) => return e.to_compile_error().into(),
};
let arms = pairs.iter().map(|(i, fa)| match fa {
FieldAccess::Named(id) => quote! { #i => &mut self.#id, },
FieldAccess::Tuple(idx) => quote! { #i => &mut self.#idx, },
FieldAccess::Array { ident, offset } => quote! { #i => &mut self.#ident[#i - #offset], },
});
quote! {
impl std::ops::IndexMut<usize> for #name {
#[inline]
fn index_mut(&mut self, index: usize) -> &mut <Self as std::ops::Index<usize>>::Output {
match index {
#(#arms)*
_ => panic!("Index out of bounds: {}", index),
}
}
}
}.into()
}
#[proc_macro_derive(Iter, attributes(index))]
pub fn derive_iter(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let name = &input.ident;
let (pairs, output_type) = match collect_indices(&input) {
Ok(p) => p,
Err(err) => return err.to_compile_error().into(),
};
let n = pairs.len();
let accesses = pairs.iter().map(|(i, fa)| match fa {
FieldAccess::Named(ident) => quote! { &self.#ident },
FieldAccess::Tuple(idx) => quote! { &self.#idx },
FieldAccess::Array { ident, offset } => quote! { &self.#ident[#i - #offset] },
});
let expanded = quote! {
impl<'a> IntoIterator for &'a #name {
type Item = &'a #output_type;
type IntoIter = std::array::IntoIter<&'a #output_type, #n>;
fn into_iter(self) -> Self::IntoIter {
[ #(#accesses,)* ].into_iter()
}
}
};
expanded.into()
}
#[proc_macro_derive(IterMut, attributes(index))]
pub fn derive_iter_mut(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
let name = &input.ident;
let (pairs, _output_type) = match collect_indices(&input) {
Ok(p) => p,
Err(err) => return err.to_compile_error().into(),
};
let n = pairs.len();
let accesses = pairs.iter().map(|(i, fa)| match fa {
FieldAccess::Named(ident) => quote! { &mut self.#ident },
FieldAccess::Tuple(idx) => quote! { &mut self.#idx },
FieldAccess::Array { ident, offset } => quote! { &mut self.#ident[#i - #offset] },
});
let expanded = quote! {
impl<'a> IntoIterator for &'a mut #name {
type Item = &'a mut <#name as std::ops::Index<usize>>::Output;
type IntoIter = std::array::IntoIter<&'a mut <#name as std::ops::Index<usize>>::Output, #n>;
fn into_iter(self) -> Self::IntoIter {
[ #(#accesses,)* ].into_iter()
}
}
};
expanded.into()
}