1use proc_macro2::TokenStream as TokenStream2;
2use proc_macro::{TokenStream, TokenTree};
3use std::cmp::Ordering;
4use quote::quote;
5
6#[proc_macro]
9pub fn input_checked(cnt: TokenStream) -> TokenStream {
10 let v = cnt
11 .into_iter()
12 .next()
13 .expect("Input macro expects a string");
14
15 match v {
16 TokenTree::Literal(l) => {
17 let s = l.to_string();
18 let mut gs = parse_groups(&s[1..][..s.len() - 2]);
19
20 let out = match gs.len().cmp(&1) {
21 Ordering::Equal => {
22 let Group { ty, radix } = gs.pop().unwrap();
23
24 quote! {
25 <#ty as ::saneput::FromStdin>::read_cin(&mut _cin, #radix)
26 }
27 }
28 Ordering::Greater => {
29 let tupitems = gs.into_iter().map(|Group { ty, radix }| {
30 quote! {
31 <#ty as ::saneput::FromStdin>::read_cin(&mut _cin, #radix)
32 }
33 });
34
35 quote! {
36 (
37 #(#tupitems),*
38 )
39 }
40 }
41 _ => panic!("Input string must contain at least one group"),
42 };
43
44 (quote! {
45 {
46 <::std::io::Stdout as ::std::io::Write>::flush(&mut std::io::stdout()).unwrap();
47 let mut _cin = ::std::io::stdin();
48 #out
49 }
50 })
51 .into()
52 }
53 _ => panic!("Input macro expects a string"),
54 }
55}
56
57#[proc_macro]
60pub fn input(cnt: TokenStream) -> TokenStream {
61 let v = cnt
62 .into_iter()
63 .next()
64 .expect("Input macro expects a string");
65
66 match v {
67 TokenTree::Literal(l) => {
68 let s = l.to_string();
69 let mut gs = parse_groups(&s[1..][..s.len() - 2]);
70
71 let out = match gs.len().cmp(&1) {
72 Ordering::Equal => {
73 let Group { ty, radix } = gs.pop().unwrap();
74
75 quote! {
76 <#ty as ::saneput::FromStdin>::read_cin(&mut _cin, #radix).unwrap()
77 }
78 }
79 Ordering::Greater => {
80 let tupitems = gs.into_iter().map(|Group { ty, radix }| {
81 quote! {
82 <#ty as ::saneput::FromStdin>::read_cin(&mut _cin, #radix).unwrap()
83 }
84 });
85
86 quote! {
87 (
88 #(#tupitems),*
89 )
90 }
91 }
92 _ => panic!("Input string must contain at least one group"),
93 };
94
95 (quote! {
96 {
97 <::std::io::Stdout as ::std::io::Write>::flush(&mut std::io::stdout()).unwrap();
98 let mut _cin = ::std::io::stdin();
99 #out
100 }
101 })
102 .into()
103 }
104 _ => panic!("Input macro expects a string"),
105 }
106}
107
108struct Group {
109 ty: TokenStream2,
110 radix: TokenStream2,
111}
112
113fn parse_groups(s: &str) -> Vec<Group> {
114 let mut current_group = None;
115 let mut groups = vec![];
116
117 for (i, c) in s.char_indices() {
118 if c == '{' {
119 if current_group.is_none() {
120 current_group = Some(i + 1);
121 } else {
122 panic!("Unexpected `{{`");
123 }
124 } else if c == '}' {
125 if let Some(cg) = current_group {
126 groups.push(parse_single_group(&s[cg..i]));
127 current_group = None;
128 } else {
129 panic!("Unexpected `}}`");
130 }
131 } else if current_group.is_none() {
132 panic!("Unexpected character: `{c}`");
133 }
134 }
135
136 groups
137}
138
139fn parse_single_group(mut s: &str) -> Group {
140 if let Some((mut ty, radix)) = s.split_once(':') {
141 if ty.is_empty() {
142 ty = "i32";
143 }
144 Group {
145 ty: ty.parse().unwrap(),
146 radix: radix_from_str(radix),
147 }
148 } else {
149 if s.is_empty() {
150 s = "i32";
151 }
152 Group {
153 ty: s.parse().unwrap(),
154 radix: quote!(::std::option::Option::None),
155 }
156 }
157}
158
159fn radix_from_str(s: &str) -> TokenStream2 {
160 let v = match s {
161 "b" => quote!(2),
162 "o" => quote!(8),
163 "d" => quote!(10),
164 "x" => quote!(16),
165 _ => panic!("Invalid radix. Expected: `b`, `o`, `d`, `x`"),
166 };
167
168 quote!(::std::option::Option::Some(#v))
169}