Skip to content

Commit fd239f9

Browse files
committed
Fix proc macro panicked with Invalid type in const_random! in Rust 1.77.
With Rust nightly 1.77, any invocation of `const_random!` with a `u8` array causes a compile-time panic. This can be seen when running `cargo test`: ``` error: proc macro panicked --> tests/tests.rs:52:28 | 52 | const VALUE1: &[u8] = &const_random!([u8; 30]); | ^^^^^^^^^^^^^^^^^^^^^^^ | = help: message: Invalid type ``` This is because the proc macro starts by calling `to_string` on the input token stream, and then uses substring matching to "parse" it. In Rust 1.77 the `Display` impl for `TokenStream` has changed, and what used to be converted to the string `"[u8 ; 30]"` is now `"[u8; 30]"`. As a result, the `byte_array.starts_with("[u8 ; ")` call fails. Note that substring matching is inherently flawed because the whitespace in the output of `to_string` is not guaranteed. This commit rewrites the proc macro to be robust in the face of `to_string` whitespace changes, by iterating over the individual `TokenTrees`s. Note: I ran `cargo fmt` within `macro/` to format the changes to `macro/src/lib.rs` and it made some minor changes to `macro/src/span.rs` as well.
1 parent e9f560b commit fd239f9

File tree

2 files changed

+86
-44
lines changed

2 files changed

+86
-44
lines changed

macro/src/lib.rs

+83-40
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,9 @@
1-
#[allow(unused_extern_crates)]
21
extern crate proc_macro;
32

43
use proc_macro::*;
54
use std::iter::once;
65
mod span;
7-
use crate::span::{gen_random_bytes, gen_random};
8-
6+
use crate::span::{gen_random, gen_random_bytes};
97

108
/// Create a TokenStream of an identifier out of a string
119
fn ident(ident: &str) -> TokenStream {
@@ -14,44 +12,89 @@ fn ident(ident: &str) -> TokenStream {
1412

1513
#[proc_macro]
1614
pub fn const_random(input: TokenStream) -> TokenStream {
17-
match &input.to_string()[..] {
18-
"u8" => TokenTree::from(Literal::u8_suffixed(gen_random())).into(),
19-
"u16" => TokenTree::from(Literal::u16_suffixed(gen_random())).into(),
20-
"u32" => TokenTree::from(Literal::u32_suffixed(gen_random())).into(),
21-
"u64" => TokenTree::from(Literal::u64_suffixed(gen_random())).into(),
22-
"u128" => TokenTree::from(Literal::u128_suffixed(gen_random())).into(),
23-
"i8" => TokenTree::from(Literal::i8_suffixed(gen_random())).into(),
24-
"i16" => TokenTree::from(Literal::i16_suffixed(gen_random())).into(),
25-
"i32" => TokenTree::from(Literal::i32_suffixed(gen_random())).into(),
26-
"i64" => TokenTree::from(Literal::i64_suffixed(gen_random())).into(),
27-
"i128" => TokenTree::from(Literal::i128_suffixed(gen_random())).into(),
28-
"usize" => {
29-
let value: TokenStream = TokenTree::from(Literal::u128_suffixed(gen_random())).into();
30-
let type_cast: TokenStream = [value, ident("as"), ident("usize")]
31-
.iter()
32-
.cloned()
33-
.collect();
34-
TokenTree::from(Group::new(Delimiter::Parenthesis, type_cast)).into()
15+
let mut iter = input.into_iter();
16+
let Some(tt) = iter.next() else {
17+
panic!("missing type arg");
18+
};
19+
20+
let result = match &tt {
21+
TokenTree::Ident(id) => {
22+
let s = id.to_string();
23+
match s.as_str() {
24+
"u8" => TokenTree::from(Literal::u8_suffixed(gen_random())).into(),
25+
"u16" => TokenTree::from(Literal::u16_suffixed(gen_random())).into(),
26+
"u32" => TokenTree::from(Literal::u32_suffixed(gen_random())).into(),
27+
"u64" => TokenTree::from(Literal::u64_suffixed(gen_random())).into(),
28+
"u128" => TokenTree::from(Literal::u128_suffixed(gen_random())).into(),
29+
"i8" => TokenTree::from(Literal::i8_suffixed(gen_random())).into(),
30+
"i16" => TokenTree::from(Literal::i16_suffixed(gen_random())).into(),
31+
"i32" => TokenTree::from(Literal::i32_suffixed(gen_random())).into(),
32+
"i64" => TokenTree::from(Literal::i64_suffixed(gen_random())).into(),
33+
"i128" => TokenTree::from(Literal::i128_suffixed(gen_random())).into(),
34+
"usize" => {
35+
// Note: usize does not implement `Random` and follow the pattern above. If it
36+
// did, when cross-compiling from a 32-bit host to a 64-bit target,
37+
// `usize::random()` would produce a 32-bit random usize which would then be
38+
// turned into a suffixed literal (e.g. `0x1234_5678usize`). On the 64-bit
39+
// target that literal would always have the upper 32 bits as zero, which would
40+
// be bad. Instead we produce code that will generate a 128-bit integer literal
41+
// (on the target) and then truncate it to usize (on the host).
42+
let value: TokenStream =
43+
TokenTree::from(Literal::u128_suffixed(gen_random())).into();
44+
let type_cast: TokenStream = [value, ident("as"), ident("usize")]
45+
.iter()
46+
.cloned()
47+
.collect();
48+
TokenTree::from(Group::new(Delimiter::Parenthesis, type_cast)).into()
49+
}
50+
"isize" => {
51+
// The same reasoning as `usize` applies for `isize`.
52+
let value: TokenStream =
53+
TokenTree::from(Literal::i128_suffixed(gen_random())).into();
54+
let type_cast: TokenStream = [value, ident("as"), ident("isize")]
55+
.iter()
56+
.cloned()
57+
.collect();
58+
TokenTree::from(Group::new(Delimiter::Parenthesis, type_cast)).into()
59+
}
60+
_ => panic!("invalid integer type arg: `{}`", s),
61+
}
3562
}
36-
"isize" => {
37-
let value: TokenStream = TokenTree::from(Literal::i128_suffixed(gen_random())).into();
38-
let type_cast: TokenStream = [value, ident("as"), ident("isize")]
39-
.iter()
40-
.cloned()
41-
.collect();
42-
TokenTree::from(Group::new(Delimiter::Parenthesis, type_cast)).into()
63+
TokenTree::Group(group) if group.delimiter() == Delimiter::Bracket => {
64+
let mut iter = group.stream().into_iter();
65+
match (&iter.next(), &iter.next(), &iter.next(), &iter.next()) {
66+
(
67+
Some(TokenTree::Ident(ident)),
68+
Some(TokenTree::Punct(punct)),
69+
Some(TokenTree::Literal(literal)),
70+
None,
71+
) if ident.to_string().as_str() == "u8" && punct.as_char() == ';' => {
72+
let Ok(len) = literal.to_string().parse() else {
73+
panic!("invalid array length: `{}`", literal);
74+
};
75+
let mut random_bytes = vec![0; len];
76+
gen_random_bytes(&mut random_bytes);
77+
let array_parts: TokenStream = random_bytes
78+
.into_iter()
79+
.flat_map(|byte| {
80+
let val = TokenTree::from(Literal::u8_suffixed(byte));
81+
let comma = TokenTree::from(Punct::new(',', Spacing::Alone));
82+
once(val).chain(once(comma))
83+
})
84+
.collect();
85+
TokenTree::from(Group::new(Delimiter::Bracket, array_parts)).into()
86+
}
87+
_ => panic!("invalid array type arg: `{}`", tt),
88+
}
4389
}
44-
byte_array if byte_array.starts_with("[u8 ; ") && byte_array.ends_with(']')=> {
45-
let len = byte_array[6..byte_array.len()-1].parse().unwrap();
46-
let mut random_bytes = vec![0; len];
47-
gen_random_bytes(&mut random_bytes);
48-
let array_parts: TokenStream = random_bytes.into_iter().flat_map(|byte| {
49-
let val = TokenTree::from(Literal::u8_suffixed(byte));
50-
let comma = TokenTree::from(Punct::new(',', Spacing::Alone));
51-
once(val).chain(once(comma))
52-
}).collect();
53-
TokenTree::from(Group::new(Delimiter::Bracket, array_parts)).into()
90+
_ => {
91+
panic!("invalid type arg: `{}`", tt);
5492
}
55-
_ => panic!("Invalid type"),
56-
}
93+
};
94+
95+
if let Some(tt) = iter.next() {
96+
panic!("invalid trailing token tree: `{}`", tt);
97+
};
98+
99+
result
57100
}

macro/src/span.rs

+3-4
Original file line numberDiff line numberDiff line change
@@ -2,16 +2,15 @@ use proc_macro::Span;
22
use std::option_env;
33

44
use once_cell::race::OnceBox;
5-
use tiny_keccak::{Xof, Hasher, Shake};
6-
5+
use tiny_keccak::{Hasher, Shake, Xof};
76

87
static SEED: OnceBox<Vec<u8>> = OnceBox::new();
98

109
fn get_seed() -> &'static [u8] {
1110
&SEED.get_or_init(|| {
1211
if let Some(value) = option_env!("CONST_RANDOM_SEED") {
13-
Box::new(value.as_bytes().to_vec())
14-
} else {
12+
Box::new(value.as_bytes().to_vec())
13+
} else {
1514
let mut value = [0u8; 32];
1615
getrandom::getrandom(&mut value).unwrap();
1716
Box::new(value.to_vec())

0 commit comments

Comments
 (0)