pub struct Lexer<'a> {
    pub buf: &'a str,
}

impl<'a> Lexer<'a> {
    fn skip_whitespace(&mut self) {
        while let Some(b' ') = self.buf.as_bytes().first() {
            self.buf = &self.buf[1..]
        }
    }

    fn parse_num(&mut self) -> Option<&'a str> {
        let mut i = 0;
        while matches!(self.buf.as_bytes().get(i), Some(&b) if b.is_ascii_digit()) {
            i += 1;
        }
        if let Some(b'.') = self.buf.as_bytes().get(i) {
            i += 1;
        }
        while matches!(self.buf.as_bytes().get(i), Some(&b) if b.is_ascii_digit()) {
            i += 1;
        }
        let (Some(b'e') | Some(b'E')) = self.buf.as_bytes().get(i) else {
            let t = &self.buf[..i];
            self.buf = &self.buf[i..];
            return Some(t);
        };
        i += 1;
        if let Some(b'-') = self.buf.as_bytes().get(i) {
            i += 1;
        }
        while matches!(self.buf.as_bytes().get(i), Some(&b) if b.is_ascii_digit()) {
            i += 1;
        }
        let t = &self.buf[..i];
        self.buf = &self.buf[i..];
        Some(t)
    }
}

impl<'a> Iterator for Lexer<'a> {
    type Item = &'a str;

    fn next(&mut self) -> Option<&'a str> {
        self.skip_whitespace();
        let &first_byte = self.buf.as_bytes().first()?;
        if first_byte.is_ascii_digit() || first_byte == b'.' {
            return self.parse_num();
        }
        match first_byte {
            b'/' => {
                if let Some(b'/') = self.buf.as_bytes().get(1) {
                    self.buf = &self.buf[2..];
                    return Some("//");
                }
                self.buf = &self.buf[1..];
                return Some("/");
            }
            b'+' | b'-' | b'*' | b'^' | b'(' | b')' => {
                let t = &self.buf[..1];
                self.buf = &self.buf[1..];
                return Some(t);
            }
            _ => (),
        }
        let mut i = 0;
        while let Some(b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9') = self.buf.as_bytes().get(i) {
            i += 1;
        }
        let t = &self.buf[..i];
        self.buf = &self.buf[i..];
        Some(t)
    }
}