ifmt.rs 7.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228
  1. use ::quote::{quote, ToTokens};
  2. use ::std::ops::Not;
  3. use ::syn::{
  4. parse::{Parse, ParseStream},
  5. punctuated::Punctuated,
  6. *,
  7. };
  8. use proc_macro2::TokenStream;
  9. pub fn format_args_f_impl(input: IfmtInput) -> Result<TokenStream> {
  10. let IfmtInput {
  11. mut format_literal,
  12. mut positional_args,
  13. mut named_args,
  14. } = input;
  15. let s = format_literal.value();
  16. let out_format_literal = &mut String::with_capacity(s.len());
  17. let mut iterator = s.char_indices().peekable();
  18. while let Some((i, c)) = iterator.next() {
  19. out_format_literal.push(c);
  20. if c != '{' {
  21. continue;
  22. }
  23. // encountered `{`, let's see if it was `{{`
  24. if let Some(&(_, '{')) = iterator.peek() {
  25. let _ = iterator.next();
  26. out_format_literal.push('{');
  27. continue;
  28. }
  29. let (end, colon_or_closing_brace) =
  30. iterator
  31. .find(|&(_, c)| c == '}' || c == ':')
  32. .expect(concat!(
  33. "Invalid format string literal\n",
  34. "note: if you intended to print `{`, ",
  35. "you can escape it using `{{`",
  36. ));
  37. // We use defer to ensure all the `continue`s append the closing char.
  38. let mut out_format_literal = defer(&mut *out_format_literal, |it| {
  39. it.push(colon_or_closing_brace)
  40. });
  41. let out_format_literal: &mut String = *out_format_literal;
  42. let mut arg = s[i + 1..end].trim();
  43. if let Some("=") = arg.get(arg.len().saturating_sub(1)..) {
  44. assert_eq!(
  45. out_format_literal.pop(), // Remove the opening brace
  46. Some('{'),
  47. );
  48. arg = &arg[..arg.len() - 1];
  49. out_format_literal.push_str(arg);
  50. out_format_literal.push_str(" = {");
  51. }
  52. if arg.is_empty() {
  53. continue;
  54. }
  55. enum Segment {
  56. Ident(Ident),
  57. LitInt(LitInt),
  58. }
  59. let segments: Vec<Segment> = {
  60. impl Parse for Segment {
  61. fn parse(input: ParseStream<'_>) -> Result<Self> {
  62. let lookahead = input.lookahead1();
  63. if lookahead.peek(Ident) {
  64. input.parse().map(Segment::Ident)
  65. } else if lookahead.peek(LitInt) {
  66. input.parse().map(Segment::LitInt)
  67. } else {
  68. Err(lookahead.error())
  69. }
  70. }
  71. }
  72. match ::syn::parse::Parser::parse_str(
  73. Punctuated::<Segment, Token![.]>::parse_separated_nonempty,
  74. arg,
  75. ) {
  76. Ok(segments) => segments.into_iter().collect(),
  77. Err(err) => return Err(err),
  78. }
  79. };
  80. match segments.len() {
  81. 0 => unreachable!("`parse_separated_nonempty` returned empty"),
  82. 1 => {
  83. out_format_literal.push_str(arg);
  84. match { segments }.pop().unwrap() {
  85. Segment::LitInt(_) => {
  86. // found something like `{0}`, let `format_args!`
  87. // handle it.
  88. continue;
  89. }
  90. Segment::Ident(ident) => {
  91. // if `ident = ...` is not yet among the extra args
  92. if named_args.iter().all(|(it, _)| *it != ident) {
  93. named_args.push((
  94. ident.clone(),
  95. parse_quote!(#ident), // Expr
  96. ));
  97. }
  98. }
  99. }
  100. }
  101. _ => {
  102. ::std::fmt::Write::write_fmt(
  103. out_format_literal,
  104. format_args!("{}", positional_args.len()),
  105. )
  106. .expect("`usize` or `char` Display impl cannot panic");
  107. let segments: Punctuated<TokenStream, Token![.]> = segments
  108. .into_iter()
  109. .map(|it| match it {
  110. Segment::Ident(ident) => ident.into_token_stream(),
  111. Segment::LitInt(literal) => literal.into_token_stream(),
  112. })
  113. .collect();
  114. positional_args.push(parse_quote! {
  115. #segments
  116. })
  117. }
  118. }
  119. }
  120. let named_args = named_args.into_iter().map(|(ident, expr)| {
  121. quote! {
  122. #ident = #expr
  123. }
  124. });
  125. format_literal = LitStr::new(out_format_literal, format_literal.span());
  126. Ok(quote! {
  127. format_args!(
  128. #format_literal
  129. #(, #positional_args)*
  130. #(, #named_args)*
  131. )
  132. })
  133. }
  134. #[allow(dead_code)] // dumb compiler does not see the struct being used...
  135. pub struct IfmtInput {
  136. format_literal: LitStr,
  137. positional_args: Vec<Expr>,
  138. named_args: Vec<(Ident, Expr)>,
  139. }
  140. impl Parse for IfmtInput {
  141. fn parse(input: ParseStream) -> Result<Self> {
  142. let format_literal = input.parse()?;
  143. let mut positional_args = vec![];
  144. loop {
  145. if input.parse::<Option<Token![,]>>()?.is_none() {
  146. return Ok(Self {
  147. format_literal,
  148. positional_args,
  149. named_args: vec![],
  150. });
  151. }
  152. if input.peek(Ident) && input.peek2(Token![=]) && input.peek3(Token![=]).not() {
  153. // Found a positional parameter
  154. break;
  155. }
  156. positional_args.push(input.parse()?);
  157. }
  158. let named_args = Punctuated::<_, Token![,]>::parse_terminated_with(input, |input| {
  159. Ok({
  160. let name: Ident = input.parse()?;
  161. let _: Token![=] = input.parse()?;
  162. let expr: Expr = input.parse()?;
  163. (name, expr)
  164. })
  165. })?
  166. .into_iter()
  167. .collect();
  168. Ok(Self {
  169. format_literal,
  170. positional_args,
  171. named_args,
  172. })
  173. }
  174. }
  175. pub fn defer<'a, T: 'a, Drop: 'a>(x: T, drop: Drop) -> impl ::core::ops::DerefMut<Target = T> + 'a
  176. where
  177. Drop: FnOnce(T),
  178. {
  179. use ::core::mem::ManuallyDrop;
  180. struct Ret<T, Drop>(ManuallyDrop<T>, ManuallyDrop<Drop>)
  181. where
  182. Drop: FnOnce(T);
  183. impl<T, Drop> ::core::ops::Drop for Ret<T, Drop>
  184. where
  185. Drop: FnOnce(T),
  186. {
  187. fn drop(&'_ mut self) {
  188. use ::core::ptr;
  189. unsafe {
  190. // # Safety
  191. //
  192. // - This is the canonical example of using `ManuallyDrop`.
  193. let value = ManuallyDrop::into_inner(ptr::read(&self.0));
  194. let drop = ManuallyDrop::into_inner(ptr::read(&self.1));
  195. drop(value);
  196. }
  197. }
  198. }
  199. impl<T, Drop> ::core::ops::Deref for Ret<T, Drop>
  200. where
  201. Drop: FnOnce(T),
  202. {
  203. type Target = T;
  204. #[inline]
  205. fn deref(&'_ self) -> &'_ Self::Target {
  206. &self.0
  207. }
  208. }
  209. impl<T, Drop> ::core::ops::DerefMut for Ret<T, Drop>
  210. where
  211. Drop: FnOnce(T),
  212. {
  213. #[inline]
  214. fn deref_mut(&'_ mut self) -> &'_ mut Self::Target {
  215. &mut self.0
  216. }
  217. }
  218. Ret(ManuallyDrop::new(x), ManuallyDrop::new(drop))
  219. }