rustgym_util/
nested_integer.rs

1#[derive(Debug, PartialEq, Eq)]
2pub enum NestedInteger {
3    Int(i32),
4    List(Vec<NestedInteger>),
5}
6
7#[macro_export]
8macro_rules! nested_integer {
9    ($string:expr) => {{
10        if let Some(res) = NestedIntegerParser::new($string).parse() {
11            res
12        } else {
13            panic!()
14        }
15    }};
16}
17
18use std::iter::Peekable;
19use std::vec::IntoIter;
20use NestedInteger::*;
21
22#[derive(PartialEq, Eq, Clone, Copy, Debug)]
23enum Tok {
24    Num(i32),
25    Op(char),
26}
27
28use Tok::*;
29
30pub struct NestedIntegerParser {
31    it: Peekable<IntoIter<Tok>>,
32}
33
34impl NestedIntegerParser {
35    pub fn new<T: Into<String>>(s: T) -> Self {
36        let s: String = s.into();
37        let mut tokens: Vec<Tok> = vec![];
38        let mut it = s.chars().peekable();
39        while let Some(c) = it.next() {
40            match c {
41                '0'..='9' => {
42                    let mut val = (c as u8 - b'0') as i32;
43                    while let Some(&c) = it.peek() {
44                        match c {
45                            '0'..='9' => {
46                                it.next();
47                                val *= 10;
48                                val += (c as u8 - b'0') as i32;
49                            }
50                            _ => {
51                                break;
52                            }
53                        }
54                    }
55                    tokens.push(Num(val));
56                }
57                _ => {
58                    tokens.push(Op(c));
59                }
60            }
61        }
62        NestedIntegerParser {
63            it: tokens.into_iter().peekable(),
64        }
65    }
66
67    pub fn parse(&mut self) -> Option<NestedInteger> {
68        if !self.eat('[') {
69            return self.parse_int();
70        }
71        let mut list: Vec<NestedInteger> = vec![];
72        while let Some(x) = self.parse() {
73            list.push(x);
74            self.eat(',');
75        }
76        if !self.eat(']') {
77            return None;
78        }
79        Some(List(list))
80    }
81
82    fn parse_int(&mut self) -> Option<NestedInteger> {
83        let sign = if self.eat('-') { -1 } else { 1 };
84        if let Some(&Tok::Num(num)) = self.it.peek() {
85            self.it.next();
86            return Some(Int(sign * num));
87        }
88        None
89    }
90
91    fn eat(&mut self, c: char) -> bool {
92        if let Some(&Tok::Op(t)) = self.it.peek() {
93            if t == c {
94                self.it.next();
95                true
96            } else {
97                false
98            }
99        } else {
100            false
101        }
102    }
103}