[flang-commits] [flang] [flang] Fix crash due to overly broad assertion (PR #95809)

[Peter Klausler via flang-commits]

https://github.com/klausler created https://github.com/llvm/llvm-project/pull/95809

Fix https://github.com/llvm/llvm-project/issues/95689 and add a regression test.

>From 6737a541ceb9d6cc0577262da24684ae574477ea Mon Sep 17 00:00:00 2001
From: Peter Klausler <pklausler at nvidia.com>
Date: Mon, 17 Jun 2024 09:17:04 -0700
Subject: [PATCH] [flang] Fix crash due to overly broad assertion

Fix https://github.com/llvm/llvm-project/issues/95689 and add
a regression test.
---
 flang/lib/Parser/prescan.cpp            | 8 +++-----
 flang/test/Preprocessing/multi-cont.F90 | 6 ++++++
 2 files changed, 9 insertions(+), 5 deletions(-)
 create mode 100644 flang/test/Preprocessing/multi-cont.F90

diff --git a/flang/lib/Parser/prescan.cpp b/flang/lib/Parser/prescan.cpp
index 8efcd617cf0f9..2a6ecfbb0830e 100644
--- a/flang/lib/Parser/prescan.cpp
+++ b/flang/lib/Parser/prescan.cpp
@@ -185,11 +185,9 @@ void Prescanner::Statement() {
       // a comment marker or directive sentinel.  If so, disable line
       // continuation, so that NextToken() won't consume anything from
       // following lines.
-      if (IsLegalIdentifierStart(*at_)) {
-        CHECK(NextToken(tokens));
-        CHECK(tokens.SizeInTokens() == 1);
-        CharBlock id{tokens.TokenAt(0)};
-        if (preprocessor_.IsNameDefined(id) &&
+      if (IsLegalIdentifierStart(*at_) && NextToken(tokens) &&
+          tokens.SizeInTokens() > 0) {
+        if (CharBlock id{tokens.TokenAt(0)}; preprocessor_.IsNameDefined(id) &&
             !preprocessor_.IsFunctionLikeDefinition(id)) {
           if (auto replaced{preprocessor_.MacroReplacement(tokens, *this)}) {
             auto newLineClass{ClassifyLine(*replaced, GetCurrentProvenance())};
diff --git a/flang/test/Preprocessing/multi-cont.F90 b/flang/test/Preprocessing/multi-cont.F90
new file mode 100644
index 0000000000000..5a2c0b9c47976
--- /dev/null
+++ b/flang/test/Preprocessing/multi-cont.F90
@@ -0,0 +1,6 @@
+! RUN: %flang -E %s 2>&1 | FileCheck --strict-whitespace %s
+! CHECK:      print *, 666
+pr&
+&i&
+&nt *, 666
+end