1 use std::io::{self, Read};
2 use std::process;
3
4 use csv_core::{ReadFieldResult, Reader};
5
run(mut data: &[u8]) -> Option<u64>6 fn run(mut data: &[u8]) -> Option<u64> {
7 let mut rdr = Reader::new();
8
9 // Count the number of records in Massachusetts.
10 let mut count = 0;
11 // Indicates the current field index. Reset to 0 at start of each record.
12 let mut fieldidx = 0;
13 // True when the current record is in the United States.
14 let mut inus = false;
15 // Buffer for field data. Must be big enough to hold the largest field.
16 let mut field = [0; 1024];
17 loop {
18 // Attempt to incrementally read the next CSV field.
19 let (result, nread, nwrite) = rdr.read_field(data, &mut field);
20 // nread is the number of bytes read from our input. We should never
21 // pass those bytes to read_field again.
22 data = &data[nread..];
23 // nwrite is the number of bytes written to the output buffer `field`.
24 // The contents of the buffer after this point is unspecified.
25 let field = &field[..nwrite];
26
27 match result {
28 // We don't need to handle this case because we read all of the
29 // data up front. If we were reading data incrementally, then this
30 // would be a signal to read more.
31 ReadFieldResult::InputEmpty => {}
32 // If we get this case, then we found a field that contains more
33 // than 1024 bytes. We keep this example simple and just fail.
34 ReadFieldResult::OutputFull => {
35 return None;
36 }
37 // This case happens when we've successfully read a field. If the
38 // field is the last field in a record, then `record_end` is true.
39 ReadFieldResult::Field { record_end } => {
40 if fieldidx == 0 && field == b"us" {
41 inus = true;
42 } else if inus && fieldidx == 3 && field == b"MA" {
43 count += 1;
44 }
45 if record_end {
46 fieldidx = 0;
47 inus = false;
48 } else {
49 fieldidx += 1;
50 }
51 }
52 // This case happens when the CSV reader has successfully exhausted
53 // all input.
54 ReadFieldResult::End => {
55 break;
56 }
57 }
58 }
59 Some(count)
60 }
61
main()62 fn main() {
63 // Read the entire contents of stdin up front.
64 let mut data = vec![];
65 if let Err(err) = io::stdin().read_to_end(&mut data) {
66 println!("{}", err);
67 process::exit(1);
68 }
69 match run(&data) {
70 None => {
71 println!("error: could not count records, buffer too small");
72 process::exit(1);
73 }
74 Some(count) => {
75 println!("{}", count);
76 }
77 }
78 }
79