set_span/
lib.rs

1#![doc = include_str!("../README.md")]
2#![allow(clippy::needless_doctest_main)]
3
4use std::convert::identity;
5
6use proc_macro::{Delimiter, Group, Span, TokenStream, TokenTree as TT};
7use proc_macro_tool::{
8    err, rerr, try_pfunc, GetSpan, ParseIter, ParseIterExt,
9    SetSpan, TokenStreamExt, TokenTreeExt, WalkExt,
10};
11
12fn eoi(iter: impl IntoIterator<Item = TT>) -> Result<(), TokenStream> {
13    iter.into_iter().next().map_or(Ok(()), |tt| {
14        Err(err("unexpected token, expected end of input", tt))
15    })
16}
17
18fn fmt(g: Group) -> String {
19    if g.is_delimiter(Delimiter::None) {
20        format!("Ø{g}Ø")
21    } else {
22        g.to_string()
23    }
24}
25
26const NUM_SUFFIXS: &[&str] = &[
27    "i8", "i16", "i32", "i64", "i128",
28    "u8", "u16", "u32", "u64", "u128",
29];
30trait StrExt {
31    fn remove_number_suffix(&self) -> &Self;
32}
33impl StrExt for str {
34    fn remove_number_suffix(&self) -> &Self {
35        for &suffix in NUM_SUFFIXS {
36            if let Some(s) = self.strip_suffix(suffix) {
37                return s;
38            }
39        }
40        self
41    }
42}
43
44#[must_use]
45fn index_tt<I>(mut tt: TT, iter: &mut ParseIter<I>) -> Result<TT, TokenStream>
46where I: Iterator<Item = TT>,
47{
48    while let Some((mut span, mut param)) = iter
49        .next_if(|tt| tt.is_delimiter(Delimiter::Bracket))
50        .map(|tt| tt.into_group().unwrap())
51        .map(|g| (g.span_close(), g.stream().into_iter()))
52    {
53        let i = param.next()
54            .ok_or_else(|| err!("unexpected token, expected literal", span))?
55            .span_as(&mut span)
56            .into_literal()
57            .map_err(|_| err!("unexpected token, expected literal", span))?
58            .to_string()
59            .remove_number_suffix()
60            .parse()
61            .map_err(|e| err!(@("parse number {e}"), span))?;
62        let g = tt.into_group()
63            .map_err(|t| err!(@("cannot index {t}, e.g [...]"), span))?;
64        tt = g.stream().into_iter().nth(i)
65            .ok_or_else(|| err!(@("index {i} out of range, of {}", fmt(g)), span))?
66    };
67    Ok(tt)
68}
69
70fn parse_input_span<I>(iter: I) -> Result<TT, TokenStream>
71where I: Iterator<Item = TT>,
72{
73    let mut iter = iter.parse_iter();
74    if iter.peek_puncts("#").is_some()
75    && iter.peek_i_is(1, |t| t.is_keyword("mixed"))
76    {
77        return Ok(iter.next().unwrap().set_spaned(Span::mixed_site()));
78    }
79
80    let mut tt = iter.next()
81        .ok_or_else(|| err!("unexpected comma of input start"))?;
82
83    tt = index_tt(tt, &mut iter)?;
84
85    if let Some(end) = iter.next() {
86        rerr!("unexpected token, expected [...] or comma", end)
87    }
88
89    Ok(tt)
90}
91
92fn extract_expand_body<I>(
93    input: &mut ParseIter<I>,
94    span: Span,
95) -> Result<TokenStream, TokenStream>
96where I: Iterator<Item = TT>,
97{
98    let Some(tt) = input.next() else {
99        rerr!("unexpected end of input, expected a brace {...}", span)
100    };
101    let Some(group) = tt.as_group() else {
102        rerr!("unexpected token, expected a brace {...}", tt)
103    };
104
105    let out = if group.is_solid_group() {
106        group.stream()
107    } else {
108        extract_expand_body(input, group.span())?
109    };
110
111    eoi(input)?;
112
113    Ok(out)
114}
115
116fn do_operation(
117    input: TokenStream,
118    spant: TT,
119) -> Result<TokenStream, TokenStream> {
120    try_pfunc(input, false, [
121        "set_span",
122        "set_index_span",
123    ], |i, param| {
124        Ok(match &*i.to_string() {
125            "set_span" => {
126                param.stream()
127                    .walk(|tt| tt.set_spaned(spant.span()))
128            },
129            "set_index_span" => {
130                let iter = &mut param.stream().parse_iter();
131                let spant = index_tt(spant.clone(), iter)?;
132                let result = iter.next()
133                    .ok_or_else(|| err!("unexpected end of input, expected {...}", param))?
134                    .into_group()
135                    .map_err(|t| err!("expected {...}", t))?
136                    .stream()
137                    .walk(|tt| tt.set_spaned(spant.span()));
138                eoi(iter)?;
139                result
140            },
141            _ => unreachable!(),
142        })
143    })
144}
145
146fn set_span_impl(input: TokenStream) -> Result<TokenStream, TokenStream> {
147    let Some((spani, mut input)) = input.split_puncts(",") else {
148        rerr!("unexpected end of input, expected comma");
149    };
150    let span = parse_input_span(spani.into_iter())?;
151    let input = extract_expand_body(&mut input, Span::call_site())?;
152    do_operation(input, span)
153}
154
155fn set_span_all_impl(input: TokenStream) -> Result<TokenStream, TokenStream> {
156    let iter = &mut input.parse_iter();
157    let spant = iter.next()
158        .ok_or_else(|| err!("unexpected end of input, expected any token"))?;
159    let spant = index_tt(spant, iter)?;
160
161    iter.next_puncts(",");
162
163    let result = iter.next()
164        .ok_or_else(|| err!("unexpected end of input, expected {...}", Span::call_site()))?
165        .into_group()
166        .map_err(|t| err!("expected {...}", t))?
167        .stream()
168        .walk(|tt| tt.set_spaned(spant.span()));
169    eoi(iter)?;
170
171    Ok(result)
172}
173
174/// Set the span of certain tokens in the code block to the span of the input token
175///
176/// - grammar := span `,` `{` code\* `}`
177/// - span := `#mixed` / any-token index\*
178/// - index := `[` num-literal `]`
179/// - code := `#set_span` `{` ... `}`<br>
180///         / `#set_index_span` `{` index\* `{` ... `}` `}`<br>
181///         / any-token
182///
183/// `{`, `}` is general bracket, contain `()` `[]` `{}`
184///
185/// - `set_span!((a, b, (c))[0], {...})` set `a` span
186/// - `set_span!((a, b, (c))[1], {...})` set `b` span
187/// - `set_span!((a, b, (c))[2], {...})` set `(c)` span
188/// - `set_span!((a, b, (c))[2][0], {...})` set `c` span
189///
190/// Similarly, there is also
191/// `set_span!((a), {#set_index_span([0]{...})})` set `a` span
192///
193///
194/// # Example
195/// ```
196/// macro_rules! foo {
197///     ($t:tt) => {
198///         foo! { ($t) ($t) }
199///     };
200///     ($t:tt (0)) => {
201///         set_span::set_span! {$t[0], {
202///             #set_span {
203///                 compile_error! {"input by zero"}
204///             }
205///         }}
206///     };
207///     ($_:tt ($lit:literal)) => { /*...*/ };
208/// }
209/// // foo!(0); // test error msg
210/// ```
211#[proc_macro]
212pub fn set_span(input: TokenStream) -> TokenStream {
213    set_span_impl(input).unwrap_or_else(identity)
214}
215
216/// Like `set_span!(tt[0], {#set_span{...}})`
217///
218/// # Example
219/// ```
220/// macro_rules! foo {
221///     ($t:tt) => {
222///         foo! { ($t) ($t) }
223///     };
224///     ($t:tt (0)) => {
225///         set_span::set_span_all! {$t[0], {
226///             compile_error! {"input by zero"}
227///         }}
228///     };
229///     ($_:tt ($lit:literal)) => { /*...*/ };
230/// }
231/// // foo!(0); // test error msg
232/// ```
233#[proc_macro]
234pub fn set_span_all(input: TokenStream) -> TokenStream {
235    set_span_all_impl(input).unwrap_or_else(identity)
236}
set_span/lib.rs

set_span/
lib.rs
