saneput_proc/
lib.rs

1use proc_macro2::TokenStream as TokenStream2;
2use proc_macro::{TokenStream, TokenTree};
3use std::cmp::Ordering;
4use quote::quote;
5
6/// Parses standard input.
7/// Flushes standard out when called. For more info see crate's root.
8#[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/// Parses standard input, panicing when error occurs.
58/// Flushes standard out when called. For more info see crate's root.
59#[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}