Skip to content

Commit 5ba3812

Browse files
dtolnaydtolnay
committed
Partition FromStr behaviors into checked and unchecked
1 parent ea101de commit 5ba3812

File tree

3 files changed

+104
-99
lines changed

3 files changed

+104
-99
lines changed

‎src/fallback.rs

+71-55
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,9 @@ use core::mem::ManuallyDrop;
1515
use core::ops::Range;
1616
use core::ops::RangeBounds;
1717
use core::ptr;
18-
use core::str::{self, FromStr};
18+
use core::str;
19+
#[cfg(feature = "proc-macro")]
20+
use core::str::FromStr;
1921
use std::ffi::CStr;
2022
#[cfg(procmacro2_semver_exempt)]
2123
use std::path::PathBuf;
@@ -63,6 +65,24 @@ impl TokenStream {
6365
}
6466
}
6567

68+
pub(crate) fn from_str_checked(src: &str) -> Result<Self, LexError> {
69+
// Create a dummy file & add it to the source map
70+
let mut cursor = get_cursor(src);
71+
72+
// Strip a byte order mark if present
73+
const BYTE_ORDER_MARK: &str = "\u{feff}";
74+
if cursor.starts_with(BYTE_ORDER_MARK) {
75+
cursor = cursor.advance(BYTE_ORDER_MARK.len());
76+
}
77+
78+
parse::token_stream(cursor)
79+
}
80+
81+
#[cfg(feature = "proc-macro")]
82+
pub(crate) fn from_str_unchecked(src: &str) -> Self {
83+
Self::from_str_checked(src).unwrap()
84+
}
85+
6686
pub(crate) fn is_empty(&self) -> bool {
6787
self.inner.len() == 0
6888
}
@@ -170,23 +190,6 @@ fn get_cursor(src: &str) -> Cursor {
170190
Cursor { rest: src }
171191
}
172192

173-
impl FromStr for TokenStream {
174-
type Err = LexError;
175-
176-
fn from_str(src: &str) -> Result<TokenStream, LexError> {
177-
// Create a dummy file & add it to the source map
178-
let mut cursor = get_cursor(src);
179-
180-
// Strip a byte order mark if present
181-
const BYTE_ORDER_MARK: &str = "\u{feff}";
182-
if cursor.starts_with(BYTE_ORDER_MARK) {
183-
cursor = cursor.advance(BYTE_ORDER_MARK.len());
184-
}
185-
186-
parse::token_stream(cursor)
187-
}
188-
}
189-
190193
impl Display for LexError {
191194
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
192195
f.write_str("cannot parse string into token stream")
@@ -226,15 +229,14 @@ impl Debug for TokenStream {
226229
#[cfg(feature = "proc-macro")]
227230
impl From<proc_macro::TokenStream> for TokenStream {
228231
fn from(inner: proc_macro::TokenStream) -> Self {
229-
TokenStream::from_str(&inner.to_string()).expect("compiler token stream parse failed")
232+
TokenStream::from_str_unchecked(&inner.to_string())
230233
}
231234
}
232235

233236
#[cfg(feature = "proc-macro")]
234237
impl From<TokenStream> for proc_macro::TokenStream {
235238
fn from(inner: TokenStream) -> Self {
236-
proc_macro::TokenStream::from_str(&inner.to_string())
237-
.expect("failed to parse to compiler tokens")
239+
proc_macro::TokenStream::from_str_unchecked(&inner.to_string())
238240
}
239241
}
240242

@@ -950,6 +952,36 @@ impl Literal {
950952
}
951953
}
952954

955+
pub(crate) fn from_str_checked(repr: &str) -> Result<Self, LexError> {
956+
let mut cursor = get_cursor(repr);
957+
#[cfg(span_locations)]
958+
let lo = cursor.off;
959+
960+
let negative = cursor.starts_with_char('-');
961+
if negative {
962+
cursor = cursor.advance(1);
963+
if !cursor.starts_with_fn(|ch| ch.is_ascii_digit()) {
964+
return Err(LexError::call_site());
965+
}
966+
}
967+
968+
if let Ok((rest, mut literal)) = parse::literal(cursor) {
969+
if rest.is_empty() {
970+
if negative {
971+
literal.repr.insert(0, '-');
972+
}
973+
literal.span = Span {
974+
#[cfg(span_locations)]
975+
lo,
976+
#[cfg(span_locations)]
977+
hi: rest.off,
978+
};
979+
return Ok(literal);
980+
}
981+
}
982+
Err(LexError::call_site())
983+
}
984+
953985
pub(crate) unsafe fn from_str_unchecked(repr: &str) -> Self {
954986
Literal::_new(repr.to_owned())
955987
}
@@ -1147,40 +1179,6 @@ impl Literal {
11471179
}
11481180
}
11491181

1150-
impl FromStr for Literal {
1151-
type Err = LexError;
1152-
1153-
fn from_str(repr: &str) -> Result<Self, Self::Err> {
1154-
let mut cursor = get_cursor(repr);
1155-
#[cfg(span_locations)]
1156-
let lo = cursor.off;
1157-
1158-
let negative = cursor.starts_with_char('-');
1159-
if negative {
1160-
cursor = cursor.advance(1);
1161-
if !cursor.starts_with_fn(|ch| ch.is_ascii_digit()) {
1162-
return Err(LexError::call_site());
1163-
}
1164-
}
1165-
1166-
if let Ok((rest, mut literal)) = parse::literal(cursor) {
1167-
if rest.is_empty() {
1168-
if negative {
1169-
literal.repr.insert(0, '-');
1170-
}
1171-
literal.span = Span {
1172-
#[cfg(span_locations)]
1173-
lo,
1174-
#[cfg(span_locations)]
1175-
hi: rest.off,
1176-
};
1177-
return Ok(literal);
1178-
}
1179-
}
1180-
Err(LexError::call_site())
1181-
}
1182-
}
1183-
11841182
impl Display for Literal {
11851183
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
11861184
Display::fmt(&self.repr, f)
@@ -1219,3 +1217,21 @@ fn escape_utf8(string: &str, repr: &mut String) {
12191217
}
12201218
}
12211219
}
1220+
1221+
#[cfg(feature = "proc-macro")]
1222+
pub(crate) trait FromStr2: FromStr<Err = proc_macro::LexError> {
1223+
#[cfg(wrap_proc_macro)]
1224+
fn from_str_checked(src: &str) -> Result<Self, Self::Err> {
1225+
Self::from_str(src)
1226+
}
1227+
1228+
fn from_str_unchecked(src: &str) -> Self {
1229+
Self::from_str(src).unwrap()
1230+
}
1231+
}
1232+
1233+
#[cfg(feature = "proc-macro")]
1234+
impl FromStr2 for proc_macro::TokenStream {}
1235+
1236+
#[cfg(feature = "proc-macro")]
1237+
impl FromStr2 for proc_macro::Literal {}

‎src/lib.rs

+2-2
Original file line numberDiff line numberDiff line change
@@ -245,7 +245,7 @@ impl FromStr for TokenStream {
245245
type Err = LexError;
246246

247247
fn from_str(src: &str) -> Result<TokenStream, LexError> {
248-
match imp::TokenStream::from_str(src) {
248+
match imp::TokenStream::from_str_checked(src) {
249249
Ok(tokens) => Ok(TokenStream::_new(tokens)),
250250
Err(lex) => Err(LexError {
251251
inner: lex,
@@ -1309,7 +1309,7 @@ impl FromStr for Literal {
13091309
type Err = LexError;
13101310

13111311
fn from_str(repr: &str) -> Result<Self, LexError> {
1312-
match imp::Literal::from_str(repr) {
1312+
match imp::Literal::from_str_checked(repr) {
13131313
Ok(lit) => Ok(Literal::_new(lit)),
13141314
Err(lex) => Err(LexError {
13151315
inner: lex,

‎src/wrapper.rs

+31-42
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
use crate::detection::inside_proc_macro;
2+
use crate::fallback::{self, FromStr2 as _};
23
#[cfg(span_locations)]
34
use crate::location::LineColumn;
4-
use crate::{fallback, Delimiter, Punct, Spacing, TokenTree};
5+
use crate::{Delimiter, Punct, Spacing, TokenTree};
56
use core::fmt::{self, Debug, Display};
67
#[cfg(span_locations)]
78
use core::ops::Range;
89
use core::ops::RangeBounds;
9-
use core::str::FromStr;
1010
use std::ffi::CStr;
1111
use std::panic;
1212
#[cfg(super_unstable)]
@@ -86,6 +86,21 @@ impl TokenStream {
8686
}
8787
}
8888

89+
pub(crate) fn from_str_checked(src: &str) -> Result<Self, LexError> {
90+
if inside_proc_macro() {
91+
// Catch panic to work around https://github.com/rust-lang/rust/issues/58736.
92+
match panic::catch_unwind(|| proc_macro::TokenStream::from_str_checked(src)) {
93+
Ok(Ok(tokens)) => Ok(TokenStream::Compiler(DeferredTokenStream::new(tokens))),
94+
Ok(Err(lex)) => Err(LexError::Compiler(lex)),
95+
Err(_panic) => Err(LexError::CompilerPanic),
96+
}
97+
} else {
98+
Ok(TokenStream::Fallback(
99+
fallback::TokenStream::from_str_checked(src)?,
100+
))
101+
}
102+
}
103+
89104
pub(crate) fn is_empty(&self) -> bool {
90105
match self {
91106
TokenStream::Compiler(tts) => tts.is_empty(),
@@ -108,28 +123,6 @@ impl TokenStream {
108123
}
109124
}
110125

111-
impl FromStr for TokenStream {
112-
type Err = LexError;
113-
114-
fn from_str(src: &str) -> Result<TokenStream, LexError> {
115-
if inside_proc_macro() {
116-
Ok(TokenStream::Compiler(DeferredTokenStream::new(
117-
proc_macro_parse(src)?,
118-
)))
119-
} else {
120-
Ok(TokenStream::Fallback(fallback::TokenStream::from_str(src)?))
121-
}
122-
}
123-
}
124-
125-
// Work around https://github.com/rust-lang/rust/issues/58736.
126-
fn proc_macro_parse(src: &str) -> Result<proc_macro::TokenStream, LexError> {
127-
match panic::catch_unwind(|| proc_macro::TokenStream::from_str(src)) {
128-
Ok(result) => result.map_err(LexError::Compiler),
129-
Err(_) => Err(LexError::CompilerPanic),
130-
}
131-
}
132-
133126
impl Display for TokenStream {
134127
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
135128
match self {
@@ -150,7 +143,7 @@ impl From<TokenStream> for proc_macro::TokenStream {
150143
match inner {
151144
TokenStream::Compiler(inner) => inner.into_token_stream(),
152145
TokenStream::Fallback(inner) => {
153-
proc_macro::TokenStream::from_str(&inner.to_string()).unwrap()
146+
proc_macro::TokenStream::from_str_unchecked(&inner.to_string())
154147
}
155148
}
156149
}
@@ -817,9 +810,19 @@ macro_rules! unsuffixed_integers {
817810
}
818811

819812
impl Literal {
813+
pub(crate) fn from_str_checked(repr: &str) -> Result<Self, LexError> {
814+
if inside_proc_macro() {
815+
let literal = proc_macro::Literal::from_str_checked(repr)?;
816+
Ok(Literal::Compiler(literal))
817+
} else {
818+
let literal = fallback::Literal::from_str_checked(repr)?;
819+
Ok(Literal::Fallback(literal))
820+
}
821+
}
822+
820823
pub(crate) unsafe fn from_str_unchecked(repr: &str) -> Self {
821824
if inside_proc_macro() {
822-
Literal::Compiler(proc_macro::Literal::from_str(repr).expect("invalid literal"))
825+
Literal::Compiler(proc_macro::Literal::from_str_unchecked(repr))
823826
} else {
824827
Literal::Fallback(unsafe { fallback::Literal::from_str_unchecked(repr) })
825828
}
@@ -901,7 +904,7 @@ impl Literal {
901904
#[cfg(no_literal_byte_character)]
902905
{
903906
let fallback = fallback::Literal::byte_character(byte);
904-
proc_macro::Literal::from_str(&fallback.repr).unwrap()
907+
proc_macro::Literal::from_str_unchecked(&fallback.repr)
905908
}
906909
})
907910
} else {
@@ -928,7 +931,7 @@ impl Literal {
928931
#[cfg(no_literal_c_string)]
929932
{
930933
let fallback = fallback::Literal::c_string(string);
931-
proc_macro::Literal::from_str(&fallback.repr).unwrap()
934+
proc_macro::Literal::from_str_unchecked(&fallback.repr)
932935
}
933936
})
934937
} else {
@@ -976,20 +979,6 @@ impl From<fallback::Literal> for Literal {
976979
}
977980
}
978981

979-
impl FromStr for Literal {
980-
type Err = LexError;
981-
982-
fn from_str(repr: &str) -> Result<Self, Self::Err> {
983-
if inside_proc_macro() {
984-
let literal = proc_macro::Literal::from_str(repr)?;
985-
Ok(Literal::Compiler(literal))
986-
} else {
987-
let literal = fallback::Literal::from_str(repr)?;
988-
Ok(Literal::Fallback(literal))
989-
}
990-
}
991-
}
992-
993982
impl Display for Literal {
994983
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
995984
match self {

0 commit comments

Comments
 (0)