1use super::{FastInput, ScanSource};
2use std::str::from_utf8_unchecked;
3
4macro_rules! impl_fast_scan_integer {
5 ($($ty:ty, $read:ident, $method:ident);* $(;)?) => {$(
6 #[inline]
7 fn $read(&mut self) -> Option<$ty> {
8 Some(unsafe { self.$method() })
9 }
10 )*};
11}
12
13impl ScanSource for FastInput {
14 #[inline]
15 fn skip_whitespace(&mut self) {
16 self.skip_whitespace();
17 }
18
19 #[inline]
20 fn next_token(&mut self) -> Option<&str> {
21 Some(unsafe { from_utf8_unchecked(self.bytes()) })
22 }
23
24 impl_fast_scan_integer!(
25 u8, read_u8, u8;
26 u16, read_u16, u16;
27 u32, read_u32, u32;
28 u64, read_u64, u64;
29 u128, read_u128, u128;
30 usize, read_usize, usize;
31 i8, read_i8, i8;
32 i16, read_i16, i16;
33 i32, read_i32, i32;
34 i64, read_i64, i64;
35 i128, read_i128, i128;
36 isize, read_isize, isize;
37 );
38}
39
40#[cfg(test)]
41mod tests {
42 use super::*;
43 use crate::tools::testutil::{integer_boundary_values, sample_usize};
44 use crate::tools::{
45 BufferedInput, Bytes, Chars, Scan, Scanner, SizedCollect, Usize1, Xorshift,
46 };
47 use std::array;
48
49 #[test]
50 fn test_integer_tokens() {
51 let mut rng = Xorshift::default();
52 macro_rules! check {
53 ($($ty:ty),* $(,)?) => {$(
54 let mut values = integer_boundary_values!($ty);
55 values.extend((0..=255).map(|x| x as $ty));
56 values.extend((0..512).map(|_| ((rng.rand64() as u128) << 64 | rng.rand64() as u128) as $ty));
57 let width = <$ty>::MAX.to_string().len();
58 let tokens: Vec<_> = values.iter().flat_map(|value| {
59 let token = value.to_string();
60 let digits = token.trim_start_matches('-');
61 let padded = format!("{}{}{}", if token.starts_with('-') { "-" } else { "" }, "0".repeat(width - digits.len()), digits);
62 [token, padded]
63 }).collect();
64 for sep in [" ", "\n", "\t", "\x0c"] {
65 let input = format!("{} ", tokens.join(sep));
66 let mut scanner = unsafe { FastInput::from_slice(input.as_bytes()) };
67 for &value in &values {
68 assert_eq!(<$ty as Scan>::scan(&mut scanner), Some(value));
69 assert_eq!(<$ty as Scan>::scan(&mut scanner), Some(value));
70 }
71 let mut scanner = unsafe { BufferedInput::new(input.as_bytes()) };
72 for &value in &values {
73 assert_eq!(<$ty as Scan>::scan(&mut scanner), Some(value));
74 assert_eq!(<$ty as Scan>::scan(&mut scanner), Some(value));
75 }
76 }
77 )*};
78 }
79 check!(
80 u8, u16, u32, u64, u128, usize, i8, i16, i32, i64, i128, isize
81 );
82 }
83
84 #[test]
85 fn test_composite_scan() {
86 crate::define_enum_scan! {
87 enum Query: raw {
88 "ADD" => Add { v: Usize1, n: usize, xs: [(i64, u32); n] },
89 "END" => End,
90 }
91 }
92 let mut rng = Xorshift::default();
93 for _ in 0..1000 {
94 let v = rng.random(1..=100usize);
95 let n = rng.random(0..=16);
96 let xs: Vec<(i64, u32)> = rng.random_iter((.., ..)).take(n).collect();
97 let pair: [i32; 2] = array::from_fn(|_| rng.random(..));
98 let len = rng.random(0..=16);
99 let values: Vec<u64> = rng.random_iter(..).take(len).collect();
100 let token: String = (0..rng.random(1..=16))
101 .map(|_| char::from_u32(rng.random(0x3041..=0x3096)).unwrap())
102 .collect();
103 let input = format!(
104 "ADD {v} {n} {} {token} {token} {} {} {len} {} END ",
105 xs.iter()
106 .map(|(a, b)| format!("{a} {b}"))
107 .collect::<Vec<_>>()
108 .join(" "),
109 pair[0],
110 pair[1],
111 values
112 .iter()
113 .map(ToString::to_string)
114 .collect::<Vec<_>>()
115 .join(" ")
116 );
117 macro_rules! check {
118 ($scanner:expr) => {{
119 let mut scanner = $scanner;
120 crate::scan!(scanner, q: Query, bytes: Bytes, chars: Chars, actual_pair: [i32; const 2], actual_values: SizedCollect<u64>);
121 match q {
122 Query::Add { v: actual_v, n: actual_n, xs: actual_xs } => {
123 assert_eq!((actual_v, actual_n, actual_xs), (v - 1, n, xs.clone()));
124 }
125 Query::End => panic!("unexpected query"),
126 }
127 assert_eq!(bytes, token.as_bytes());
128 assert_eq!(chars, token.chars().collect::<Vec<_>>());
129 assert_eq!(actual_pair, pair);
130 assert_eq!(actual_values, values);
131 assert!(matches!(scanner.scan::<Query>(), Query::End));
132 }};
133 }
134 check!(unsafe { FastInput::from_slice(input.as_bytes()) });
135 check!(Scanner::new(&input));
136 check!(unsafe { BufferedInput::new(input.as_bytes()) });
137 }
138 }
139
140 #[test]
141 fn test_borrowed_tokens() {
142 let mut rng = Xorshift::new_with_seed(197283);
143 for _ in 0..256 {
144 let offset = rng.random(0..64);
145 let mut input = "!".repeat(offset);
146 let mut tokens = Vec::new();
147 for _ in 0..rng.random(1..32) {
148 let start = input.len();
149 for _ in 0..rng.random(1..256) {
150 let ch = loop {
151 if let Some(ch) = char::from_u32(rng.random(0..0x110000))
152 && !ch.is_ascii_whitespace()
153 {
154 break ch;
155 }
156 };
157 input.push(ch);
158 }
159 tokens.push(start..input.len());
160 input.push([' ', '\t', '\n', '\r', '\x0c'][rng.random(0usize..5)]);
161 }
162 input.push_str(" ");
163 let mut scanner = Scanner::new(&input[offset..]);
164 let mut fast = unsafe { FastInput::from_slice(&input.as_bytes()[offset..]) };
166 for (i, range) in tokens.into_iter().enumerate() {
167 let expected = &input[range];
168 let a = if i % 2 == 0 {
169 crate::scan!(scanner, token: &str);
170 token
171 } else {
172 crate::scan_value!(scanner, &str)
173 };
174 let b = if i % 2 == 0 {
175 crate::scan!(fast, token: &str);
176 token
177 } else {
178 crate::scan_value!(fast, &str)
179 };
180 assert_eq!(a, expected);
181 assert_eq!(b, expected);
182 assert_eq!(a.as_ptr(), expected.as_ptr());
183 assert_eq!(b.as_ptr(), expected.as_ptr());
184 }
185 }
186
187 struct ChunkedRead<'a> {
188 input: &'a [u8],
189 width: usize,
190 until_error: Option<usize>,
191 interrupted: bool,
192 }
193 impl std::io::Read for ChunkedRead<'_> {
194 fn read(&mut self, output: &mut [u8]) -> std::io::Result<usize> {
195 self.interrupted = !self.interrupted;
196 if self.interrupted {
197 return Err(std::io::ErrorKind::Interrupted.into());
198 }
199 if self.until_error == Some(0) {
200 self.until_error = None;
201 return Err(std::io::ErrorKind::Other.into());
202 }
203 let n = output
204 .len()
205 .min(self.width)
206 .min(self.until_error.unwrap_or(usize::MAX));
207 let n = self.input.read(&mut output[..n])?;
208 if let Some(until_error) = &mut self.until_error {
209 *until_error -= n;
210 }
211 Ok(n)
212 }
213 }
214 let lengths = sample_usize(&mut rng, 8, 0..=1 << 17, 32)
215 .into_iter()
216 .chain((1 << 16) - 64..=(1 << 16) + 64)
217 .chain([1 << 19, (1 << 19) + 1]);
218 for len in lengths {
219 let mut token = String::new();
220 while token.len() < len {
221 let ch = ['a', 'é', 'あ', '😀'][rng.random(0usize..4)];
222 token.push(if token.len() + ch.len_utf8() <= len {
223 ch
224 } else {
225 'a'
226 });
227 }
228 for whitespace in [false, true] {
229 let mut input = if whitespace {
230 " ".repeat(len)
231 } else {
232 token.clone()
233 };
234 if !whitespace && len != 0 {
235 for _ in 0..rng.random(0..=32) {
236 input.push([' ', '\n', '\t', '\r', '\x0c'][rng.random(0usize..5)]);
237 let value: u128 = rng.random(..);
238 input.push_str(&value.to_string());
239 }
240 }
241 for width in sample_usize(&mut rng, 8, 1..=1 << 16, 16) {
242 let until_error = if width == 1 << 16 && len >= 1 << 16 && !whitespace {
243 Some(if len >= 1 << 19 { 1 << 19 } else { 1 << 16 })
244 } else if rng.rand(4) == 0 {
245 Some(rng.random(0..=input.len()))
246 } else {
247 None
248 };
249 let mut scanner = unsafe {
250 BufferedInput::new(ChunkedRead {
251 input: input.as_bytes(),
252 width,
253 until_error,
254 interrupted: false,
255 })
256 };
257 let mut failed = false;
258 for expected in input.split_ascii_whitespace().map(Some).chain([None, None]) {
259 let actual = loop {
260 let result =
261 std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
262 scanner.skip_whitespace();
263 scanner.next_token().map(str::to_owned)
264 }));
265 match result {
266 Ok(token) => break token,
267 Err(_) => {
268 assert!(until_error.is_some() && !failed);
269 failed = true;
270 }
271 }
272 };
273 assert_eq!(actual.as_deref(), expected);
274 }
275 assert_eq!(failed, until_error.is_some());
276 assert_eq!(scanner.read_u32(), None);
277 }
278 }
279 }
280 }
281
282 #[test]
283 fn test_prepare_io() {
284 use std::io::Write as _;
285 let mut rng = Xorshift::default();
286 for n in 0..=32 {
287 let values: Vec<_> = (0..n)
288 .map(|_| (rng.rand64(), rng.rand64() as i64))
289 .collect();
290 let mut input = n.to_string();
291 if n == 0 {
292 input.push('\n');
293 }
294 let mut expected = Vec::new();
295 for &(a, b) in &values {
296 input.push_str(&format!("\n{a} {b}"));
297 writeln!(expected, "{a} {b}").unwrap();
298 }
299 input.push_str("\nEND");
300 macro_rules! check {
301 ($($mode:ident;)?) => {{
302 let mut output = Vec::new();
303 {
304 crate::prepare_io!($($mode;)? input.as_bytes(), &mut output);
305 sc!(len: usize);
306 for (a, b) in sv!([(u64, i64); iter len]) {
307 pp!(@tup (a, b));
308 }
309 sc!(end: &str);
310 assert_eq!(end, "END");
311 }
312 assert_eq!(output, expected);
313 }};
314 }
315 check!();
316 check!(buffered;);
317 }
318 }
319
320 #[test]
321 fn test_collection_scan() {
322 use crate::tools::Collect;
323 let mut rng = Xorshift::default();
324 for mask in 0..256 {
325 let rows: Vec<Vec<_>> = (0..4)
326 .map(|i| (0..(mask >> (i * 2)) & 3).map(|_| rng.rand64()).collect())
327 .collect();
328 let mut input = String::new();
329 for row in &rows {
330 input.push_str(
331 &row.iter()
332 .map(ToString::to_string)
333 .collect::<Vec<_>>()
334 .join(" "),
335 );
336 input.push('\n');
337 }
338 input.push_str("END\n ");
339 for lazy in [false, true] {
340 let mut scanner = unsafe { FastInput::from_slice(input.as_bytes()) };
341 for row in &rows {
342 let values = if lazy {
343 crate::scan_value!(scanner, [u64; iter row.len()]).collect::<Vec<_>>()
344 } else {
345 crate::scan_value!(scanner, [u64; row.len()])
346 };
347 assert_eq!(&values, row);
348 }
349 assert_eq!(scanner.scan::<String>(), "END");
350 for _ in 0..32 {
351 assert!(crate::scan_value!(scanner, [u8; 0]).is_empty());
352 assert!(crate::scan_value!(scanner, [u8; const 0]).is_empty());
353 assert!(scanner.scan_vec::<u8>(0).is_empty());
354 assert_eq!(crate::scan_value!(scanner, [u8; iter 0]).count(), 0);
355 assert!(scanner.mscan(Collect::<u8>::new(0)).is_empty());
356 }
357 }
358 }
359 }
360}