Skip to content

Add enumerating version of all_tuples #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 8, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 96 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
#![cfg_attr(docsrs, feature(doc_auto_cfg))]

use proc_macro::TokenStream;
use proc_macro2::{Span as Span2, TokenStream as TokenStream2};
use proc_macro2::{Literal, Span as Span2, TokenStream as TokenStream2};
use quote::{format_ident, quote};
use syn::{
parse::{Parse, ParseStream},
Expand Down Expand Up @@ -191,6 +191,81 @@ pub fn all_tuples(input: TokenStream) -> TokenStream {
})
}

/// A variant of [`all_tuples!`] that enumerates its output.
///
/// In particular, the tuples used by the inner macro will themselves be composed
/// of tuples which contain the index.
///
/// For example, with a single parameter:
/// ```
/// # use variadics_please::all_tuples_enumerated;
///
/// trait Squawk {
/// fn squawk(&self);
/// }
///
/// // If every type in a tuple is `Squawk`, the tuple can squawk by having its
/// // constituents squawk sequentially:
/// macro_rules! impl_squawk {
/// ($(($n:tt, $T:ident)),*) => {
/// impl<$($T: Squawk),*> Squawk for ($($T,)*) {
/// fn squawk(&self) {
/// $(
/// self.$n.squawk();
/// )*
/// }
/// }
/// };
/// }
///
/// all_tuples_enumerated!(impl_squawk, 1, 15, T);
/// // impl_squawk!((0, T0));
/// // impl_squawk!((0, T0), (1, T1));
/// // ..
/// // impl_squawk!((0, T0) .. (14, T14));
/// ```
///
/// With multiple parameters, the result is similar, but with the additional parameters
/// included in each tuple; e.g.:
/// ```ignore
/// all_tuples_enumerated!(impl_squawk, 1, 15, P, p);
/// // impl_squawk!((0, P0, p0));
/// // impl_squawk!((0, P0, p0), (1, P1, p1));
/// // ..
/// // impl_squawk!((0, P0, p0) .. (14, P14, p14));
/// ```
#[proc_macro]
pub fn all_tuples_enumerated(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as AllTuples);
let len = 1 + input.end - input.start;
let mut ident_tuples = Vec::with_capacity(len);
for i in 0..=len {
let idents = input
.idents
.iter()
.map(|ident| format_ident!("{}{}", ident, i));
ident_tuples.push(to_ident_tuple_enumerated(idents, i));
}

let macro_ident = &input.macro_ident;
let invocations = (input.start..=input.end).map(|i| {
let ident_tuples = choose_ident_tuples_enumerated(&input, &ident_tuples, i);
let attrs = if input.fake_variadic {
fake_variadic_attrs(len, i)
} else {
TokenStream2::default()
};
quote! {
#macro_ident!(#attrs #ident_tuples);
}
});
TokenStream::from(quote! {
#(
#invocations
)*
})
}

/// Helper macro to generate tuple pyramids with their length. Useful to generate scaffolding to
/// work around Rust lacking variadics. Invoking `all_tuples_with_size!(impl_foo, start, end, P, Q, ..)`
/// invokes `impl_foo` providing ident tuples through arity `start..=end` preceded by their length.
Expand Down Expand Up @@ -372,6 +447,20 @@ fn choose_ident_tuples(input: &AllTuples, ident_tuples: &[TokenStream2], i: usiz
}
}

fn choose_ident_tuples_enumerated(
input: &AllTuples,
ident_tuples: &[TokenStream2],
i: usize,
) -> TokenStream2 {
if input.fake_variadic && i == 1 {
let ident_tuple = to_ident_tuple_enumerated(input.idents.iter().cloned(), 0);
quote! { #ident_tuple }
} else {
let ident_tuples = &ident_tuples[..i];
quote! { #(#ident_tuples),* }
}
}

fn to_ident_tuple(idents: impl Iterator<Item = Ident>, len: usize) -> TokenStream2 {
if len < 2 {
quote! { #(#idents)* }
Expand All @@ -380,6 +469,12 @@ fn to_ident_tuple(idents: impl Iterator<Item = Ident>, len: usize) -> TokenStrea
}
}

/// Like `to_ident_tuple`, but it enumerates the identifiers
fn to_ident_tuple_enumerated(idents: impl Iterator<Item = Ident>, idx: usize) -> TokenStream2 {
let idx = Literal::usize_unsuffixed(idx);
quote! { (#idx, #(#idents),*) }
}

fn fake_variadic_attrs(len: usize, i: usize) -> TokenStream2 {
let cfg = quote! { any(docsrs, docsrs_dep) };
match i {
Expand Down