count_macro/lib.rs
1use std::collections::HashMap;
2
3use once_cell::sync::Lazy;
4use proc_macro::{Group, Ident, Literal, TokenStream, TokenTree};
5use regex::Regex;
6
7static RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"_int(?:_([^_ \s]+))?_").unwrap());
8
9fn map_tokens(token: TokenTree, counters: &mut HashMap<String, usize>) -> TokenTree {
10 match token {
11 TokenTree::Ident(v) => {
12 let mut ident = v.to_string();
13 let mut has_changed = false;
14
15 #[allow(clippy::redundant_clone)]
16 for captures in RE.captures_iter(&ident.clone()) {
17 let id = match captures.get(1) {
18 Some(id) => id.as_str().to_string(),
19
20 _ => "!@".to_string(),
21 };
22
23 let counter = match counters.get_mut(&id) {
24 Some(v) => v,
25
26 _ => {
27 counters.insert(id.clone(), 0);
28 counters.get_mut(&id).unwrap()
29 }
30 };
31
32 let full_match = &captures[0];
33 if full_match == ident {
34 let to_ret = TokenTree::Literal(Literal::usize_unsuffixed(*counter));
35 *counter += 1;
36 return to_ret;
37 } else {
38 has_changed = true;
39 ident = ident.replace(&captures[0], &counter.to_string());
40 *counter += 1;
41 }
42 }
43
44 if has_changed {
45 return TokenTree::Ident(Ident::new(&ident, v.span()));
46 }
47
48 TokenTree::Ident(v)
49 }
50
51 TokenTree::Group(v) => TokenTree::Group(Group::new(
52 v.delimiter(),
53 v.stream()
54 .into_iter()
55 .map(|token| map_tokens(token, counters))
56 .collect(),
57 )),
58
59 v => v,
60 }
61}
62
63/// Count without wrapping (panic if counter exceeds usize).
64/// Every instance of `_int_` and `_int_countername_` will be replaced with the counter value and then the counter will be increased.
65///
66/// # Examples
67///
68/// ## Ident to literal
69/// ```rust
70/// use count_macro::count;
71///
72/// let a = count!(vec![_int_, _int_, _int_]);
73/// assert_eq!(a, vec![0, 1, 2]);
74///
75/// ```
76///
77/// ## Ident to ident
78/// ```rust
79/// use count_macro::count;
80///
81/// count! {
82/// let a_int_ = "Hello";
83/// let a_int_ = "World";
84/// }
85///
86/// assert_eq!(a0, "Hello");
87/// assert_eq!(a1, "World");
88///
89/// ```
90///
91/// ## In macro
92/// ```rust
93/// use count_macro::count;
94///
95/// macro_rules! my_macro {
96/// ($($v:expr),*) => {
97/// count!{
98/// $(
99/// let _ = $v; // Ignoring $v
100///
101/// println!("{}", _int_);
102/// )*
103/// }
104/// };
105/// }
106///
107/// my_macro!('@', '@', '@', '@'); // Will print from 0 to 3
108///
109/// ```
110/// ## Multiple counters
111/// ```rust
112/// use count_macro::count;
113///
114/// // With two different counters
115/// // _int_ does not increment _int_name_
116/// count! {
117/// let a_int_ = _int_name_;
118/// let a_int_ = _int_name_;
119/// let a_int_ = _int_name_;
120/// }
121///
122/// assert_eq!(a0, 0);
123/// assert_eq!(a1, 1);
124/// assert_eq!(a2, 2);
125///
126/// ```
127#[proc_macro]
128pub fn count(item: TokenStream) -> TokenStream {
129 let mut counters = HashMap::new();
130
131 item.into_iter()
132 .map(|token| map_tokens(token, &mut counters))
133 .collect()
134}
135
136fn wrapping_map_tokens(token: TokenTree, counters: &mut HashMap<String, usize>) -> TokenTree {
137 match token {
138 TokenTree::Ident(v) => {
139 let mut ident = v.to_string();
140 let mut has_changed = false;
141
142 #[allow(clippy::redundant_clone)]
143 for captures in RE.captures_iter(&ident.clone()) {
144 let id = match captures.get(1) {
145 Some(id) => id.as_str().to_string(),
146
147 _ => "!@".to_string(),
148 };
149
150 let counter = match counters.get_mut(&id) {
151 Some(v) => v,
152
153 _ => {
154 counters.insert(id.clone(), 0);
155 counters.get_mut(&id).unwrap()
156 }
157 };
158
159 let full_match = &captures[0];
160 if full_match == ident {
161 let to_ret = TokenTree::Literal(Literal::usize_unsuffixed(*counter));
162 *counter = counter.wrapping_add(1);
163 return to_ret;
164 } else {
165 has_changed = true;
166 ident = ident.replace(&captures[0], &counter.to_string());
167 *counter = counter.wrapping_add(1);
168 }
169 }
170
171 if has_changed {
172 return TokenTree::Ident(Ident::new(&ident, v.span()));
173 }
174
175 TokenTree::Ident(v)
176 }
177
178 TokenTree::Group(v) => TokenTree::Group(Group::new(
179 v.delimiter(),
180 v.stream()
181 .into_iter()
182 .map(|token| map_tokens(token, counters))
183 .collect(),
184 )),
185
186 v => v,
187 }
188}
189
190
191/// Count with wrapping (wraps to 0 if counter exceeds usize).
192/// Every instance of `_int_` and `_int_countername_` will be replaced with the counter value and then the counter will be increased.
193///
194/// # Examples
195///
196/// ## Ident to literal
197/// ```rust
198/// use count_macro::wrapping_count;
199///
200/// let a = wrapping_count!(vec![_int_, _int_, _int_]);
201/// assert_eq!(a, vec![0, 1, 2]);
202///
203/// ```
204///
205/// ## Ident to ident
206/// ```rust
207/// use count_macro::wrapping_count;
208///
209/// wrapping_count! {
210/// let a_int_ = "Hello";
211/// let a_int_ = "World";
212/// }
213///
214/// assert_eq!(a0, "Hello");
215/// assert_eq!(a1, "World");
216///
217/// ```
218///
219/// ## In macro
220/// ```rust
221/// use count_macro::wrapping_count;
222///
223/// macro_rules! my_macro {
224/// ($($v:expr),*) => {
225/// wrapping_count!{
226/// $(
227/// let _ = $v; // Ignoring $v
228///
229/// println!("{}", _int_);
230/// )*
231/// }
232/// };
233/// }
234///
235/// my_macro!('@', '@', '@', '@'); // Will print from 0 to 3
236///
237/// ```
238///
239/// ## Multiple counters
240/// ```rust
241/// use count_macro::wrapping_count;
242///
243/// // With two different counters
244/// // _int_ does not increment _int_name_
245/// wrapping_count! {
246/// let a_int_ = _int_name_;
247/// let a_int_ = _int_name_;
248/// let a_int_ = _int_name_;
249/// }
250///
251/// assert_eq!(a0, 0);
252/// assert_eq!(a1, 1);
253/// assert_eq!(a2, 2);
254///
255/// ```
256#[proc_macro]
257pub fn wrapping_count(item: TokenStream) -> TokenStream {
258 let mut counters = HashMap::new();
259
260 item.into_iter()
261 .map(|token| wrapping_map_tokens(token, &mut counters))
262 .collect()
263}