[flang] Preserve compiler directives in -E output

flang/lib/Parser/parsing.cpp

@@ -79,16 +79,24 @@ const SourceFile *Parsing::Prescan(const std::string &path, Options options) {
       .set_expandIncludeLines(!options.prescanAndReformat ||
           options.expandIncludeLinesInPreprocessedOutput)
       .AddCompilerDirectiveSentinel("dir$");
-  if (options.features.IsEnabled(LanguageFeature::OpenACC)) {
+  bool noneOfTheAbove{!options.features.IsEnabled(LanguageFeature::OpenACC) &&
+      !options.features.IsEnabled(LanguageFeature::OpenMP) &&
+      !options.features.IsEnabled(LanguageFeature::CUDA)};
+  if (options.features.IsEnabled(LanguageFeature::OpenACC) ||
+      (options.prescanAndReformat && noneOfTheAbove)) {
     prescanner.AddCompilerDirectiveSentinel("$acc");
   }
-  if (options.features.IsEnabled(LanguageFeature::OpenMP)) {
+  if (options.features.IsEnabled(LanguageFeature::OpenMP) ||
+      (options.prescanAndReformat && noneOfTheAbove)) {
     prescanner.AddCompilerDirectiveSentinel("$omp");
     prescanner.AddCompilerDirectiveSentinel("$"); // OMP conditional line
   }
-  if (options.features.IsEnabled(LanguageFeature::CUDA)) {
+  if (options.features.IsEnabled(LanguageFeature::CUDA) ||
+      (options.prescanAndReformat && noneOfTheAbove)) {
     prescanner.AddCompilerDirectiveSentinel("$cuf");
     prescanner.AddCompilerDirectiveSentinel("@cuf");
+  }
+  if (options.features.IsEnabled(LanguageFeature::CUDA)) {
     preprocessor_.Define("_CUDA", "1");
   }
   ProvenanceRange range{allSources.AddIncludedFile(
@@ -119,11 +127,13 @@ void Parsing::EmitPreprocessedSource(
   int sourceLine{0};
   int column{1};
   bool inDirective{false};
+  bool ompConditionalLine{false};
   bool inContinuation{false};
   bool lineWasBlankBefore{true};
   const AllSources &allSources{allCooked().allSources()};
-  // All directives that flang support are known to have a length of 3 chars
-  constexpr int directiveNameLength{3};
+  // All directives that flang supports are known to have a length of 4 chars,
+  // except for OpenMP conditional compilation lines (!$).
+  constexpr int directiveNameLength{4};
   // We need to know the current directive in order to provide correct
   // continuation for the directive
   std::string directive;
@@ -133,6 +143,7 @@ void Parsing::EmitPreprocessedSource(
       out << '\n'; // TODO: DOS CR-LF line ending if necessary
       column = 1;
       inDirective = false;
+      ompConditionalLine = false;
       inContinuation = false;
       lineWasBlankBefore = true;
       ++sourceLine;
@@ -153,16 +164,21 @@ void Parsing::EmitPreprocessedSource(
         return ch;
       }};
 
+      bool inDirectiveSentinel{false};
       if (ch == '!' && lineWasBlankBefore) {
         // Other comment markers (C, *, D) in original fixed form source
         // input card column 1 will have been deleted or normalized to !,
         // which signifies a comment (directive) in both source forms.
         inDirective = true;
-      }
-      bool inDirectiveSentinel{
-          inDirective && directive.size() < directiveNameLength};
-      if (inDirectiveSentinel && IsLetter(ch)) {
-        directive += getOriginalChar(ch);
+        inDirectiveSentinel = true;
+      } else if (inDirective && !ompConditionalLine &&
+          directive.size() < directiveNameLength) {
+        if (IsLetter(ch) || ch == '$' || ch == '@') {
+          directive += getOriginalChar(ch);
+          inDirectiveSentinel = true;
+        } else if (directive == "$"s) {
+          ompConditionalLine = true;
+        }
       }
 
       std::optional<SourcePosition> position{provenance
@@ -199,9 +215,16 @@ void Parsing::EmitPreprocessedSource(
         // column limit override option.
         // OpenMP and OpenACC directives' continuations should have the
         // corresponding sentinel at the next line.
-        const auto continuation{
-            inDirective ? "&\n!$" + directive + "&" : "&\n     &"s};
-        out << continuation;
+        out << "&\n";
+        if (inDirective) {
+          if (ompConditionalLine) {
+            out << "!$   &";
+          } else {
+            out << '!' << directive << '&';
+          }
+        } else {
+          out << "     &";
+        }
         column = 7; // start of fixed form source field
         ++sourceLine;
         inContinuation = true;
@@ -212,11 +235,20 @@ void Parsing::EmitPreprocessedSource(
           out << ' ';
         }
       }
-      if (!inContinuation && !inDirectiveSentinel && position &&
-          position->column <= 72 && ch != ' ') {
-        // Preserve original indentation
-        for (; column < position->column; ++column) {
-          out << ' ';
+      if (ch != ' ') {
+        if (ompConditionalLine) {
+          // Only digits can stay in the label field
+          if (!(ch >= '0' && ch <= '9')) {
+            for (; column < 7; ++column) {
+              out << ' ';
+            }
+          }
+        } else if (!inContinuation && !inDirectiveSentinel && position &&
+            position->column <= 72) {
+          // Preserve original indentation
+          for (; column < position->column; ++column) {
+            out << ' ';
+          }
         }
       }
       out << getOriginalChar(ch);