Added a state in the string evaluation

1 files changed, 51 insertions, 6 deletions

diff --git a/lib/qparser/lexer.ml b/lib/qparser/lexer.ml
index f2d04bf..75072c6 100644
--- a/lib/qparser/lexer.ml
+++ b/lib/qparser/lexer.ml
@@ -53,6 +53,21 @@ let location_ident = [%sedlex.regexp? letters | digit]
 let location_prefix = [%sedlex.regexp? '!' | '$' | '#' | '^']
 let location = [%sedlex.regexp? Opt location_prefix, Plus location_ident]
 
+let end_string : Lexbuf.t -> token =
+ fun buffer ->
+  let lexbuf = Lexbuf.buffer buffer in
+  match%sedlex lexbuf with
+  | '}' ->
+      Lexbuf.leave_state buffer;
+      TEXT_MARKER
+  | '\'' ->
+      Lexbuf.leave_state buffer;
+      TEXT_MARKER
+  | '"' ->
+      Lexbuf.leave_state buffer;
+      TEXT_MARKER
+  | _ -> raise Not_found
+
 let rec read_long_string level buf buffer =
   let lexbuf = Lexbuf.buffer buffer in
   match%sedlex lexbuf with
@@ -61,7 +76,11 @@ let rec read_long_string level buf buffer =
       read_long_string (level + 1) buf buffer
   | '}' -> (
       match level with
-      | 0 -> Buffer.contents buf
+      | 0 ->
+          Lexbuf.leave_state buffer;
+          Lexbuf.enter_state buffer Lexbuf.EndString;
+          Lexbuf.rollback buffer;
+          LITERAL (Buffer.contents buf)
       | _ ->
           Buffer.add_string buf (Sedlexing.Utf8.lexeme lexbuf);
           read_long_string (level - 1) buf buffer)
@@ -79,7 +98,11 @@ let rec read_dquoted_string buf buffer =
   | "\"\"" ->
       Buffer.add_char buf '"';
       read_dquoted_string buf buffer
-  | '"' -> Buffer.contents buf
+  | '"' ->
+      Lexbuf.leave_state buffer;
+      Lexbuf.enter_state buffer Lexbuf.EndString;
+      Lexbuf.rollback buffer;
+      LITERAL (Buffer.contents buf)
   | any ->
       Buffer.add_string buf (Sedlexing.Utf8.lexeme lexbuf);
       read_dquoted_string buf buffer
@@ -91,7 +114,11 @@ let rec read_quoted_string buf buffer =
   | "''" ->
       Buffer.add_char buf '\'';
       read_quoted_string buf buffer
-  | '\'' -> Buffer.contents buf
+  | '\'' ->
+      Lexbuf.leave_state buffer;
+      Lexbuf.enter_state buffer Lexbuf.EndString;
+      Lexbuf.rollback buffer;
+      LITERAL (Buffer.contents buf)
   | eol ->
       Buffer.add_string buf (Lexbuf.content buffer);
       read_quoted_string buf buffer
@@ -105,12 +132,15 @@ let rec skip_comment buffer =
   match%sedlex lexbuf with
   | '{' ->
       let _ = wait_balance (read_long_string 0) buffer in
+      let _ = end_string buffer in
       skip_comment buffer
   | '\'' ->
       let _ = wait_balance read_quoted_string buffer in
+      let _ = end_string buffer in
       skip_comment buffer
   | '"' ->
       let _ = wait_balance read_dquoted_string buffer in
+      let _ = end_string buffer in
       skip_comment buffer
   | eol ->
       (* Ugly hack used in order to put the eol in the front of the next
@@ -180,9 +210,16 @@ let rec token : Lexbuf.t -> token =
       AMPERSAND
   | '!' -> if Lexbuf.level buffer > 0 then EXCLAMATION else skip_comment buffer
   | spaces -> token buffer
-  | '\'' -> LITERAL (wait_balance read_quoted_string buffer)
-  | '"' -> LITERAL (wait_balance read_dquoted_string buffer)
-  | '{' -> LITERAL (wait_balance (read_long_string 0) buffer)
+  | '\'' ->
+      Lexbuf.enter_state buffer Lexbuf.String;
+      TEXT_MARKER
+  | '"' ->
+      Lexbuf.enter_state buffer Lexbuf.DString;
+      TEXT_MARKER
+  | '{' ->
+      Lexbuf.enter_state buffer (Lexbuf.MString 0);
+      TEXT_MARKER
+  | '}' -> TEXT_MARKER
   | eof -> raise EOF
   | _ ->
       let tok = Lexbuf.content buffer in
@@ -190,6 +227,14 @@ let rec token : Lexbuf.t -> token =
 
       raise @@ LexError msg
 
+let main buffer =
+  match Lexbuf.state buffer with
+  | Some Lexbuf.String -> wait_balance read_quoted_string buffer
+  | Some Lexbuf.DString -> wait_balance read_dquoted_string buffer
+  | Some (Lexbuf.MString level) -> wait_balance (read_long_string level) buffer
+  | Some Lexbuf.EndString -> end_string buffer
+  | _ -> token buffer
+
 let rec discard buffer =
   let lexbuf = Lexbuf.buffer buffer in
   match%sedlex lexbuf with