From 7f552e4594dcf7b6e4cd80251479d9cf315c8157 Mon Sep 17 00:00:00 2001
From: harupy <hkawamura0130@gmail.com>
Date: Thu, 5 Jan 2023 18:24:54 +0900
Subject: [PATCH] Address comments

---
 parser/src/lexer.rs         |  6 +++---
 parser/src/string_parser.rs |  2 +-
 parser/src/token.rs         | 42 +++++++++++++++++++++++++++----------
 3 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/parser/src/lexer.rs b/parser/src/lexer.rs
index d266f64..8251794 100644
--- a/parser/src/lexer.rs
+++ b/parser/src/lexer.rs
@@ -244,12 +244,12 @@ where
         // Detect potential string like rb'' b'' f'' u'' r''
         match self.window[..3] {
             [Some(c), Some('"' | '\''), ..] => {
-                if let Ok(kind) = StringKind::try_from(c.to_string()) {
+                if let Ok(kind) = StringKind::try_from(c) {
                     return self.lex_string(kind);
                 }
             }
             [Some(c1), Some(c2), Some('"' | '\'')] => {
-                if let Ok(kind) = StringKind::try_from(format!("{c1}{c2}")) {
+                if let Ok(kind) = StringKind::try_from([c1, c2]) {
                     return self.lex_string(kind);
                 }
             }
@@ -459,7 +459,7 @@ where
 
     fn lex_string(&mut self, kind: StringKind) -> LexResult {
         let start_pos = self.get_pos();
-        for _ in 0..kind.to_string().len() {
+        for _ in 0..kind.prefix_len() {
             self.next_char();
         }
         let quote_char = self.next_char().unwrap();
diff --git a/parser/src/string_parser.rs b/parser/src/string_parser.rs
index 4e94347..1d95cff 100644
--- a/parser/src/string_parser.rs
+++ b/parser/src/string_parser.rs
@@ -26,7 +26,7 @@ impl<'a> StringParser<'a> {
         str_start: Location,
         str_end: Location,
     ) -> Self {
-        let offset = kind.to_string().len() + if triple_quoted { 3 } else { 1 };
+        let offset = kind.prefix_len() + if triple_quoted { 3 } else { 1 };
         Self {
             chars: source.chars().peekable(),
             kind,
diff --git a/parser/src/token.rs b/parser/src/token.rs
index e915429..c618539 100644
--- a/parser/src/token.rs
+++ b/parser/src/token.rs
@@ -240,19 +240,30 @@ pub enum StringKind {
     Unicode,
 }
 
-impl TryFrom<String> for StringKind {
+impl TryFrom<char> for StringKind {
     type Error = String;
 
-    fn try_from(value: String) -> Result<Self, String> {
-        match value.as_str() {
-            "" => Ok(StringKind::String),
-            "r" | "R" => Ok(StringKind::RawString),
-            "u" | "U" => Ok(StringKind::Unicode),
-            "b" | "B" => Ok(StringKind::Bytes),
-            "f" | "F" => Ok(StringKind::FString),
-            "fr" | "Fr" | "fR" | "FR" | "rf" | "rF" | "Rf" | "RF" => Ok(StringKind::RawFString),
-            "br" | "Br" | "bR" | "BR" | "rb" | "rB" | "Rb" | "RB" => Ok(StringKind::RawBytes),
-            s => Err(format!("Unexpected string prefix: {s}")),
+    fn try_from(ch: char) -> Result<Self, String> {
+        match ch {
+            'r' | 'R' => Ok(StringKind::RawString),
+            'f' | 'F' => Ok(StringKind::FString),
+            'u' | 'U' => Ok(StringKind::Unicode),
+            'b' | 'B' => Ok(StringKind::Bytes),
+            c => Err(format!("Unexpected string prefix: {c}")),
+        }
+    }
+}
+
+impl TryFrom<[char; 2]> for StringKind {
+    type Error = String;
+
+    fn try_from(chars: [char; 2]) -> Result<Self, String> {
+        match chars {
+            ['r' | 'R', 'f' | 'F'] => Ok(StringKind::RawFString),
+            ['f' | 'F', 'r' | 'R'] => Ok(StringKind::RawFString),
+            ['r' | 'R', 'b' | 'B'] => Ok(StringKind::RawBytes),
+            ['b' | 'B', 'r' | 'R'] => Ok(StringKind::RawBytes),
+            [c1, c2] => Err(format!("Unexpected string prefix: {c1}{c2}")),
         }
     }
 }
@@ -291,4 +302,13 @@ impl StringKind {
     pub fn is_unicode(&self) -> bool {
         matches!(self, StringKind::Unicode)
     }
+
+    pub fn prefix_len(&self) -> usize {
+        use StringKind::*;
+        match self {
+            String => 0,
+            RawString | FString | Unicode | Bytes => 1,
+            RawFString | RawBytes => 2,
+        }
+    }
 }