[flang][runtime] Handle NAN(...) in namelist input (#153101)

The various per-type functions for list-directed (including namelist)
input editing all call a common function to detect whether the next
token of input is the name of a namelist item. This check simply
determines whether this next token looks like an identifier followed by
'=', '(', or '%', and this fails when the next item of input is a NAN
with parenthesized stuff afterwards. Make the check smarter so that it
ensures that any upcoming possible identifier is actually the name of an
item in the namelist group. (And that's tricky too when the group has an
array item named "nan" and the upcoming input is "nan("; see the
newly-added unit test case.)

Fixes https://github.com/llvm/llvm-project/issues/152538.

more

5 files changed, 97 insertions, 38 deletions

diff --git a/flang-rt/include/flang-rt/runtime/io-stmt.h b/flang-rt/include/flang-rt/runtime/io-stmt.h
index 0b3194d..9f71d51 100644
--- a/flang-rt/include/flang-rt/runtime/io-stmt.h
+++ b/flang-rt/include/flang-rt/runtime/io-stmt.h
@@ -438,7 +438,9 @@ template <>
 class ListDirectedStatementState<Direction::Input>
     : public FormattedIoStatementState<Direction::Input> {
 public:
-  RT_API_ATTRS bool inNamelistSequence() const { return inNamelistSequence_; }
+  RT_API_ATTRS const NamelistGroup *namelistGroup() const {
+    return namelistGroup_;
+  }
   RT_API_ATTRS int EndIoStatement();
 
   // Skips value separators, handles repetition and null values.
@@ -451,18 +453,19 @@ public:
   // input statement.  This member function resets some state so that
   // repetition and null values work correctly for each successive
   // NAMELIST input item.
-  RT_API_ATTRS void ResetForNextNamelistItem(bool inNamelistSequence) {
+  RT_API_ATTRS void ResetForNextNamelistItem(
+      const NamelistGroup *namelistGroup) {
     remaining_ = 0;
     if (repeatPosition_) {
       repeatPosition_->Cancel();
     }
     eatComma_ = false;
     realPart_ = imaginaryPart_ = false;
-    inNamelistSequence_ = inNamelistSequence;
+    namelistGroup_ = namelistGroup;
   }
 
 protected:
-  bool inNamelistSequence_{false};
+  const NamelistGroup *namelistGroup_{nullptr};
 
 private:
   int remaining_{0}; // for "r*" repetition
diff --git a/flang-rt/lib/runtime/edit-input.cpp b/flang-rt/lib/runtime/edit-input.cpp
index 80cc085..4f01623 100644
--- a/flang-rt/lib/runtime/edit-input.cpp
+++ b/flang-rt/lib/runtime/edit-input.cpp
@@ -534,7 +534,7 @@ static RT_API_ATTRS ScannedRealInput ScanRealInput(
       next = io.NextInField(remaining, edit);
     }
     if (!next || *next == ')') { // NextInField fails on separators like ')'
-      std::size_t byteCount{0};
+      std::size_t byteCount{1};
       if (!next) {
         next = io.GetCurrentChar(byteCount);
       }
diff --git a/flang-rt/lib/runtime/io-stmt.cpp b/flang-rt/lib/runtime/io-stmt.cpp
index af44a9d..e08088f 100644
--- a/flang-rt/lib/runtime/io-stmt.cpp
+++ b/flang-rt/lib/runtime/io-stmt.cpp
@@ -1086,7 +1086,7 @@ ChildListIoStatementState<DIR>::ChildListIoStatementState(
   if constexpr (DIR == Direction::Input) {
     if (auto *listInput{child.parent()
                 .get_if<ListDirectedStatementState<Direction::Input>>()}) {
-      this->inNamelistSequence_ = listInput->inNamelistSequence();
+      this->namelistGroup_ = listInput->namelistGroup();
     }
   }
 #else
diff --git a/flang-rt/lib/runtime/namelist.cpp b/flang-rt/lib/runtime/namelist.cpp
index cbc3226..44a8fe2 100644
--- a/flang-rt/lib/runtime/namelist.cpp
+++ b/flang-rt/lib/runtime/namelist.cpp
@@ -44,8 +44,7 @@ bool IODEF(OutputNamelist)(Cookie cookie, const NamelistGroup &group) {
     if ((connection.NeedAdvance(prefixLen) &&
             !(io.AdvanceRecord() && EmitAscii(io, " ", 1))) ||
         !EmitAscii(io, prefix, prefixLen) ||
-        (connection.NeedAdvance(
-             Fortran::runtime::strlen(str) + (suffix != ' ')) &&
+        (connection.NeedAdvance(runtime::strlen(str) + (suffix != ' ')) &&
             !(io.AdvanceRecord() && EmitAscii(io, " ", 1)))) {
       return false;
     }
@@ -102,8 +101,8 @@ static constexpr RT_API_ATTRS char NormalizeIdChar(char32_t ch) {
   return static_cast<char>(ch >= 'A' && ch <= 'Z' ? ch - 'A' + 'a' : ch);
 }
 
-static RT_API_ATTRS bool GetLowerCaseName(
-    IoStatementState &io, char buffer[], std::size_t maxLength) {
+static RT_API_ATTRS bool GetLowerCaseName(IoStatementState &io, char buffer[],
+    std::size_t maxLength, bool crashIfTooLong = true) {
   std::size_t byteLength{0};
   if (auto ch{io.GetNextNonBlank(byteLength)}) {
     if (IsLegalIdStart(*ch)) {
@@ -117,8 +116,10 @@ static RT_API_ATTRS bool GetLowerCaseName(
       if (j <= maxLength) {
         return true;
       }
-      io.GetIoErrorHandler().SignalError(
-          "Identifier '%s...' in NAMELIST input group is too long", buffer);
+      if (crashIfTooLong) {
+        io.GetIoErrorHandler().SignalError(
+            "Identifier '%s...' in NAMELIST input group is too long", buffer);
+      }
     }
   }
   return false;
@@ -356,9 +357,8 @@ static RT_API_ATTRS bool HandleComponent(IoStatementState &io, Descriptor &desc,
     const DescriptorAddendum *addendum{source.Addendum()};
     if (const typeInfo::DerivedType *
         type{addendum ? addendum->derivedType() : nullptr}) {
-      if (const typeInfo::Component *
-          comp{type->FindDataComponent(
-              compName, Fortran::runtime::strlen(compName))}) {
+      if (const typeInfo::Component *comp{
+              type->FindDataComponent(compName, runtime::strlen(compName))}) {
         bool createdDesc{false};
         if (comp->rank() > 0 && source.rank() > 0) {
           // If base and component are both arrays, the component name
@@ -484,7 +484,7 @@ bool IODEF(InputNamelist)(Cookie cookie, const NamelistGroup &group) {
       handler.SignalError("NAMELIST input group has no name");
       return false;
     }
-    if (Fortran::runtime::strcmp(group.groupName, name) == 0) {
+    if (runtime::strcmp(group.groupName, name) == 0) {
       break; // found it
     }
     SkipNamelistGroup(io);
@@ -503,7 +503,7 @@ bool IODEF(InputNamelist)(Cookie cookie, const NamelistGroup &group) {
     }
     std::size_t itemIndex{0};
     for (; itemIndex < group.items; ++itemIndex) {
-      if (Fortran::runtime::strcmp(name, group.item[itemIndex].name) == 0) {
+      if (runtime::strcmp(name, group.item[itemIndex].name) == 0) {
         break;
       }
     }
@@ -577,13 +577,14 @@ bool IODEF(InputNamelist)(Cookie cookie, const NamelistGroup &group) {
     if (const auto *addendum{useDescriptor->Addendum()};
         addendum && addendum->derivedType()) {
       const NonTbpDefinedIoTable *table{group.nonTbpDefinedIo};
-      listInput->ResetForNextNamelistItem(/*inNamelistSequence=*/true);
+      listInput->ResetForNextNamelistItem(&group);
       if (!IONAME(InputDerivedType)(cookie, *useDescriptor, table) &&
           handler.InError()) {
         return false;
       }
     } else {
-      listInput->ResetForNextNamelistItem(useDescriptor->rank() > 0);
+      listInput->ResetForNextNamelistItem(
+          useDescriptor->rank() > 0 ? &group : nullptr);
       if (!descr::DescriptorIO<Direction::Input>(io, *useDescriptor) &&
           handler.InError()) {
         return false;
@@ -607,27 +608,51 @@ bool IODEF(InputNamelist)(Cookie cookie, const NamelistGroup &group) {
 }
 
 RT_API_ATTRS bool IsNamelistNameOrSlash(IoStatementState &io) {
-  if (auto *listInput{
-          io.get_if<ListDirectedStatementState<Direction::Input>>()}) {
-    if (listInput->inNamelistSequence()) {
-      SavedPosition savedPosition{io};
-      std::size_t byteCount{0};
-      if (auto ch{io.GetNextNonBlank(byteCount)}) {
-        if (IsLegalIdStart(*ch)) {
-          do {
-            io.HandleRelativePosition(byteCount);
-            ch = io.GetCurrentChar(byteCount);
-          } while (ch && IsLegalIdChar(*ch));
-          ch = io.GetNextNonBlank(byteCount);
-          // TODO: how to deal with NaN(...) ambiguity?
-          return ch && (*ch == '=' || *ch == '(' || *ch == '%');
-        } else {
-          return *ch == '/' || *ch == '&' || *ch == '$';
-        }
-      }
+  auto *listInput{io.get_if<ListDirectedStatementState<Direction::Input>>()};
+  if (!listInput || !listInput->namelistGroup()) {
+    return false; // not namelist
+  }
+  SavedPosition savedPosition{io};
+  std::size_t byteCount{0};
+  auto ch{io.GetNextNonBlank(byteCount)};
+  if (!ch) {
+    return false;
+  } else if (!IsLegalIdStart(*ch)) {
+    return *ch == '/' || *ch == '&' || *ch == '$';
+  }
+  char id[nameBufferSize];
+  if (!GetLowerCaseName(io, id, sizeof id, /*crashIfTooLong=*/false)) {
+    return true; // long name
+  }
+  // It looks like a name, but might be "inf" or "nan".  Check what
+  // follows.
+  ch = io.GetNextNonBlank(byteCount);
+  if (!ch) {
+    return false;
+  } else if (*ch == '=' || *ch == '%') {
+    return true;
+  } else if (*ch != '(') {
+    return false;
+  } else if (runtime::strcmp(id, "nan") != 0) {
+    return true;
+  }
+  // "nan(" ambiguity
+  int depth{1};
+  while (true) {
+    io.HandleRelativePosition(byteCount);
+    ch = io.GetNextNonBlank(byteCount);
+    if (depth == 0) {
+      // nan(...) followed by '=', '%', or '('?
+      break;
+    } else if (!ch) {
+      return true; // not a valid NaN(...)
+    } else if (*ch == '(') {
+      ++depth;
+    } else if (*ch == ')') {
+      --depth;
     }
   }
-  return false;
+  return ch && (*ch == '=' || *ch == '%' || *ch == '(');
 }
 
 RT_OFFLOAD_API_GROUP_END
diff --git a/flang-rt/unittests/Runtime/Namelist.cpp b/flang-rt/unittests/Runtime/Namelist.cpp
index ee4018e..f190bea 100644
--- a/flang-rt/unittests/Runtime/Namelist.cpp
+++ b/flang-rt/unittests/Runtime/Namelist.cpp
@@ -334,4 +334,35 @@ TEST(NamelistTests, RealValueForInt) {
   EXPECT_EQ(got, expect);
 }
 
+TEST(NamelistTests, NanInputAmbiguity) {
+  OwningPtr<Descriptor> xDesc{// real :: x(5) = 0.
+      MakeArray<TypeCategory::Real, static_cast<int>(sizeof(float))>(
+          std::vector<int>{5}, std::vector<float>{{0, 0, 0, 0, 0}})};
+  OwningPtr<Descriptor> nanDesc{// real :: nan(2) = 0.
+      MakeArray<TypeCategory::Real, static_cast<int>(sizeof(float))>(
+          std::vector<int>{2}, std::vector<float>{{0, 0}})};
+  const NamelistGroup::Item items[]{{"x", *xDesc}, {"nan", *nanDesc}};
+  const NamelistGroup group{"nml", 2, items};
+  static char t1[]{"&nml x=1 2 nan(q) 4 nan(1)=5 nan(q)/"};
+  StaticDescriptor<1, true> statDesc;
+  Descriptor &internalDesc{statDesc.descriptor()};
+  internalDesc.Establish(TypeCode{CFI_type_char},
+      /*elementBytes=*/std::strlen(t1), t1, 0, nullptr, CFI_attribute_pointer);
+  auto inCookie{IONAME(BeginInternalArrayListInput)(
+      internalDesc, nullptr, 0, __FILE__, __LINE__)};
+  ASSERT_TRUE(IONAME(InputNamelist)(inCookie, group));
+  ASSERT_EQ(IONAME(EndIoStatement)(inCookie), IostatOk)
+      << "namelist real input for nans";
+  char out[40];
+  internalDesc.Establish(TypeCode{CFI_type_char}, /*elementBytes=*/sizeof out,
+      out, 0, nullptr, CFI_attribute_pointer);
+  auto outCookie{IONAME(BeginInternalArrayListOutput)(
+      internalDesc, nullptr, 0, __FILE__, __LINE__)};
+  ASSERT_TRUE(IONAME(OutputNamelist)(outCookie, group));
+  ASSERT_EQ(IONAME(EndIoStatement)(outCookie), IostatOk) << "namelist output";
+  std::string got{out, sizeof out};
+  static const std::string expect{" &NML X= 1. 2. NaN 4. 0.,NAN= 5. NaN/   "};
+  EXPECT_EQ(got, expect);
+}
+
 // TODO: Internal NAMELIST error tests