From ebb1e900d3b3edcac0e96f23cd15509f1fbe77d9 Mon Sep 17 00:00:00 2001
From: River Riddle <riddleriver@gmail.com>
Date: Tue, 26 Apr 2022 13:27:03 -0700
Subject: [PATCH] [mlir:PDLL] Fix error handling of eof within a string literal

We currently aren't handling this properly, and in the case
of a string block just crash. This commit adds proper error handling
and detection for eof.

Differential Revision: https://reviews.llvm.org/D124585
---
 mlir/lib/Tools/PDLL/Parser/Lexer.cpp       | 21 +++++++++++++++------
 mlir/lib/Tools/PDLL/Parser/Parser.cpp      |  2 ++
 mlir/test/mlir-pdll/Parser/string-eof.pdll |  9 +++++++++
 3 files changed, 26 insertions(+), 6 deletions(-)
 create mode 100644 mlir/test/mlir-pdll/Parser/string-eof.pdll

diff --git a/mlir/lib/Tools/PDLL/Parser/Lexer.cpp b/mlir/lib/Tools/PDLL/Parser/Lexer.cpp
index 6b338f1..250b21d 100644
--- a/mlir/lib/Tools/PDLL/Parser/Lexer.cpp
+++ b/mlir/lib/Tools/PDLL/Parser/Lexer.cpp
@@ -355,19 +355,28 @@ Token Lexer::lexString(const char *tokStart, bool isStringBlock) {
       case '"':
         // If this is a string block, we only end the string when we encounter a
         // `}]`.
-        if (!isStringBlock) return formToken(Token::string, tokStart);
+        if (!isStringBlock)
+          return formToken(Token::string, tokStart);
         continue;
       case '}':
         // If this is a string block, we only end the string when we encounter a
         // `}]`.
-        if (!isStringBlock || *curPtr != ']') continue;
+        if (!isStringBlock || *curPtr != ']')
+          continue;
         ++curPtr;
         return formToken(Token::string_block, tokStart);
-      case 0:
+      case 0: {
         // If this is a random nul character in the middle of a string, just
-        // include it.  If it is the end of file, then it is an error.
-        if (curPtr - 1 != curBuffer.end()) continue;
-        LLVM_FALLTHROUGH;
+        // include it. If it is the end of file, then it is an error.
+        if (curPtr - 1 != curBuffer.end())
+          continue;
+        --curPtr;
+
+        StringRef expectedEndStr = isStringBlock ? "}]" : "\"";
+        return emitError(curPtr - 1,
+                         "expected '" + expectedEndStr + "' in string literal");
+      }
+
       case '\n':
       case '\v':
       case '\f':
diff --git a/mlir/lib/Tools/PDLL/Parser/Parser.cpp b/mlir/lib/Tools/PDLL/Parser/Parser.cpp
index 3397b30..53d5417 100644
--- a/mlir/lib/Tools/PDLL/Parser/Parser.cpp
+++ b/mlir/lib/Tools/PDLL/Parser/Parser.cpp
@@ -1245,6 +1245,8 @@ FailureOr<T *> Parser::parseUserNativeConstraintOrRewriteDecl(
   } else if (isInline) {
     return emitError(name.getLoc(),
                      "external declarations must be declared in global scope");
+  } else if (curToken.is(Token::error)) {
+    return failure();
   }
   if (failed(parseToken(Token::semicolon,
                         "expected `;` after native declaration")))
diff --git a/mlir/test/mlir-pdll/Parser/string-eof.pdll b/mlir/test/mlir-pdll/Parser/string-eof.pdll
new file mode 100644
index 0000000..5265d33
--- /dev/null
+++ b/mlir/test/mlir-pdll/Parser/string-eof.pdll
@@ -0,0 +1,9 @@
+// RUN: not mlir-pdll %s -I %S -split-input-file 2>&1 | FileCheck %s
+
+// CHECK: expected '}]' in string literal
+Constraint Cst() [{
+
+// -----
+
+// CHECK: expected '"' in string literal
+Constraint Cst() "
\ No newline at end of file
-- 
2.7.4