[Flang] Fix parsing error on loop count compiler directive

[mnadeem]

Summary

Fixes:
https://github.com/llvm/llvm-project/issues/58731
https://github.com/llvm/llvm-project/issues/56678

Only handles this form which was reported: !DIR$ LOOP COUNT (n1[, n2]...)

[klausler]

flang/lib/Parser/Fortran-parsers.cpp

1219

Two minor comments: I used "ignore_tkr" for the IGNORE_TKR directive parser since it actually has an underscore in its name. LOOP COUNT does not, so "loopCount" would be a better name.

Second, be sure to run your modified code through clang-format 16.x before merging.

[sscalpone]

@mnadeem Please update the documentation llvm-project/flang/docs/Directives.md to include dir$ loop count. Thank you!

[mnadeem]

In D141976#4061262, @sscalpone wrote:

@mnadeem Please update the documentation llvm-project/flang/docs/Directives.md to include dir$ loop count. Thank you!

This directive is parsed but ignored, should I still add it to Directives.md?

flang/include/flang/Parser/dump-parse-tree.h

#include "llvm/Frontend/OpenACC/ACC.inc"
   NODE(parser, CodimensionDecl)
   NODE(parser, CodimensionStmt)
   NODE(parser, CoindexedNamedObject)
   NODE(parser, CommonBlockObject)
   NODE(parser, CommonStmt)
   NODE(CommonStmt, Block)
   NODE(parser, CompilerDirective)
   NODE(CompilerDirective, IgnoreTKR)
+  NODE(CompilerDirective, LoopCount)
   NODE(CompilerDirective, NameValue)
   NODE(parser, ComplexLiteralConstant)
   NODE(parser, ComplexPart)
   NODE(parser, ComponentArraySpec)
   NODE(parser, ComponentAttrSpec)
   NODE(parser, ComponentDataSource)
   NODE(parser, ComponentDecl)
   NODE(parser, FillDecl)

flang/include/flang/Parser/parse-tree.h

 struct StmtFunctionStmt {
   TUPLE_CLASS_BOILERPLATE(StmtFunctionStmt);
   std::tuple<Name, std::list<Name>, Scalar<Expr>> t;
   Statement<ActionStmt> ConvertToAssignment();
 };
 
 // Compiler directives
 // !DIR$ IGNORE_TKR [ [(tkr...)] name ]...
+// !DIR$ LOOP COUNT (n1[, n2]...)
 // !DIR$ name...
 struct CompilerDirective {
   UNION_CLASS_BOILERPLATE(CompilerDirective);
   struct IgnoreTKR {
     TUPLE_CLASS_BOILERPLATE(IgnoreTKR);
     std::tuple<std::list<const char *>, Name> t;
   };
+  struct LoopCount {
+    WRAPPER_CLASS_BOILERPLATE(LoopCount, std::list<std::uint64_t>);
+  };
   struct NameValue {
     TUPLE_CLASS_BOILERPLATE(NameValue);
     std::tuple<Name, std::optional<std::uint64_t>> t;
   };
   CharBlock source;
-  std::variant<std::list<IgnoreTKR>, std::list<NameValue>> u;
+  std::variant<std::list<IgnoreTKR>, LoopCount, std::list<NameValue>> u;
 };
 
 // Legacy extensions
 struct BasedPointer {
   TUPLE_CLASS_BOILERPLATE(BasedPointer);
   std::tuple<ObjectName, ObjectName, std::optional<ArraySpec>> t;
 };
 WRAPPER_CLASS(BasedPointerStmt, std::list<BasedPointer>);

flang/lib/Parser/Fortran-parsers.cpp

 
 // R942 dealloc-opt -> STAT = stat-variable | ERRMSG = errmsg-variable
 // R1165 sync-stat -> STAT = stat-variable | ERRMSG = errmsg-variable
 TYPE_PARSER(construct<StatOrErrmsg>("STAT =" >> statVariable) ||
     construct<StatOrErrmsg>("ERRMSG =" >> msgVariable))
 
 // Directives, extensions, and deprecated statements
 // !DIR$ IGNORE_TKR [ [(tkr...)] name ]...
+// !DIR$ LOOP COUNT (n1[, n2]...)
 // !DIR$ name...
 constexpr auto beginDirective{skipStuffBeforeStatement >> "!"_ch};
 constexpr auto endDirective{space >> endOfLine};
 constexpr auto ignore_tkr{
     "DIR$ IGNORE_TKR" >> optionalList(construct<CompilerDirective::IgnoreTKR>(
                              defaulted(parenthesized(some("tkr"_ch))), name))};
+constexpr auto loopCount{
+    "DIR$ LOOP COUNT" >> construct<CompilerDirective::LoopCount>(
+                             parenthesized(nonemptyList(digitString64)))};
+
 TYPE_PARSER(beginDirective >>
     sourced(construct<CompilerDirective>(ignore_tkr) ||
+        construct<CompilerDirective>(loopCount) ||
         construct<CompilerDirective>(
             "DIR$" >> many(construct<CompilerDirective::NameValue>(name,
                           maybe(("="_tok || ":"_tok) >> digitString64))))) /
         endDirective)
 
 TYPE_PARSER(extension<LanguageFeature::CrayPointer>(
     "nonstandard usage: based POINTER"_port_en_US,
     construct<BasedPointerStmt>(

flang/lib/Parser/unparse.cpp

#undef FMT
   // Directives, extensions, and deprecated constructs
   void Unparse(const CompilerDirective &x) {
     common::visit(
         common::visitors{
             [&](const std::list<CompilerDirective::IgnoreTKR> &tkr) {
               Word("!DIR$ IGNORE_TKR"); // emitted even if tkr list is empty
               Walk(" ", tkr, ", ");
             },
+            [&](const CompilerDirective::LoopCount &lcount) {
+              Walk("!DIR$ LOOP COUNT (", lcount.v, ", ", ")");
+            },
             [&](const std::list<CompilerDirective::NameValue> &names) {
               Walk("!DIR$ ", names, " ");
             },
         },
         x.u);
     Put('\n');
   }
   void Unparse(const CompilerDirective::IgnoreTKR &x) {

flang/test/Parser/compiler-directives.f90

 ! RUN: %flang_fc1 -fdebug-unparse %s 2>&1
 
 ! Test that compiler directives can appear in various places.
 
 !dir$ integer
 module m
   !dir$ integer
   use iso_fortran_env
   !dir$ integer
   implicit integer(a-z)
   !dir$ integer
   !dir$ integer=64
   !dir$ integer = 64
   !dir$ optimize:1
   !dir$ optimize : 1
+  !dir$ loop count (10000)
+  !dir$ loop count (1, 500, 5000, 10000)
 end