From c5f3936eb243aef206a516f8e1845f2e3d29d1ba Mon Sep 17 00:00:00 2001
From: ConnorSkees <39542938+ConnorSkees@users.noreply.github.com>
Date: Sat, 21 Mar 2020 23:57:36 -0400
Subject: [PATCH] refactor lexing of at-rules

---
 src/lexer.rs   | 19 +++++++------------
 src/lib.rs     |  1 +
 tests/error.rs |  4 ++++
 3 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/src/lexer.rs b/src/lexer.rs
index d952464..6eef5ea 100644
--- a/src/lexer.rs
+++ b/src/lexer.rs
@@ -195,20 +195,15 @@ impl<'a> Lexer<'a> {
     fn lex_at_rule(&mut self) -> TokenKind {
         self.buf.next();
         self.pos.next_char();
-        let mut string = String::with_capacity(99);
-        while let Some(c) = self.buf.peek() {
-            if !c.is_alphabetic() && c != &'-' && c != &'_' {
-                break;
+        if let TokenKind::Ident(s) = self.lex_ident() {
+            if !s.is_empty() {
+                TokenKind::AtRule(AtRuleKind::from(s.as_ref()))
+            } else {
+                TokenKind::Error("Expected identifier.".into())
             }
-            let tok = self
-                .buf
-                .next()
-                .expect("this is impossible because we have already peeked");
-            self.pos.next_char();
-            string.push(tok);
+        } else {
+            TokenKind::Error("Expected identifier.".into())
         }
-
-        TokenKind::AtRule(AtRuleKind::from(string.as_ref()))
     }
 
     fn lex_forward_slash(&mut self) -> TokenKind {
diff --git a/src/lib.rs b/src/lib.rs
index d1f92b4..f61657f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -426,6 +426,7 @@ impl<'a> StyleSheetParser<'a> {
                         "Base-level rules cannot contain the parent-selector-referencing character '&'.".into(),
                     )
                 }
+                TokenKind::Error(e) => return Err(e.clone()),
                 _ => match dbg!(self.lexer.next()) {
                     Some(Token { pos, .. }) => self.error(pos, "unexpected toplevel token"),
                     _ => unsafe { std::hint::unreachable_unchecked() },
diff --git a/tests/error.rs b/tests/error.rs
index 44df98b..34c91e1 100644
--- a/tests/error.rs
+++ b/tests/error.rs
@@ -11,3 +11,7 @@ error!(
     ascii_control_character,
     "a {color: ;}", "Error: Expected expression."
 );
+error!(
+    toplevel_invalid_atrule_ident,
+    "@`or $i from 1 through 3 {}", "Error: Expected identifier."
+);