[clang] Add support for optional func-name argument in #pragma init_seg

clang/include/clang/Sema/Sema.h

@@ -2075,6 +2075,10 @@ class Sema final : public SemaBase {
   /// Last section used with #pragma init_seg.
   StringLiteral *CurInitSeg;
   SourceLocation CurInitSegLoc;
+  /// Optional function name provided to #pragma init_seg([seg][, func-name]).
+  /// When present, stores the IdentifierInfo and its location.
+  IdentifierInfo *CurInitSegFn;
+  SourceLocation CurInitSegFnLoc;
 
   /// Sections used with #pragma alloc_text.
   llvm::StringMap<std::tuple<StringRef, SourceLocation>> FunctionToSectionMap;
@@ -2245,8 +2249,12 @@ class Sema final : public SemaBase {
                             StringLiteral *SegmentName);
 
   /// Called on well-formed \#pragma init_seg().
+  /// If a function name is provided (", func-name"), it will be passed
+  /// via Func and FuncLoc.
   void ActOnPragmaMSInitSeg(SourceLocation PragmaLocation,
-                            StringLiteral *SegmentName);
+                            StringLiteral *SegmentName,
+                            IdentifierInfo *Func = nullptr,
+                            SourceLocation FuncLoc = SourceLocation());
 
   /// Called on well-formed \#pragma alloc_text().
   void ActOnPragmaMSAllocText(

clang/lib/Parse/ParsePragma.cpp

@@ -1229,6 +1229,11 @@ bool Parser::HandlePragmaMSInitSeg(StringRef PragmaName,
 
   // Parse either the known section names or the string section name.
   StringLiteral *SegmentName = nullptr;
+  // Optional function identifier provided as the second argument to
+  // #pragma init_seg([segment][, func-name]). Declared here so it's in scope
+  // for the call to ActOnPragmaMSInitSeg below.
+  IdentifierInfo *FuncII = nullptr;
+  SourceLocation FuncLoc;
   if (Tok.isAnyIdentifier()) {
     auto *II = Tok.getIdentifierInfo();
     StringRef Section = llvm::StringSwitch<StringRef>(II->getName())
@@ -1259,7 +1264,31 @@ bool Parser::HandlePragmaMSInitSeg(StringRef PragmaName,
           << PragmaName;
       return false;
     }
-    // FIXME: Add support for the '[, func-name]' part of the pragma.
+    // Nothing else here; the optional ', func-name' (if present) will be
+    // parsed after the segment parsing so it works whether the segment was
+    // specified as a known identifier or as a string literal.
+  }
+
+  // After parsing the segment name (either via identifier mapping or a
+  // string literal), optionally parse a comma and an identifier naming the
+  // helper function to be used in place of atexit.
+  // Only consume the comma if an identifier follows; otherwise leave it
+  // for ExpectAndConsume(tok::r_paren) to report the error naturally (e.g.,
+  // for malformed pragmas like #pragma init_seg("a", "b")).
+  if (Tok.is(tok::comma)) {
+    // Save current token state for potential backtrack.
+    Token SavedTok = Tok;
+    PP.Lex(Tok); // tentatively consume comma
+    if (Tok.is(tok::identifier)) {
+      // The comma is followed by an identifier; keep both consumed.
+      FuncII = Tok.getIdentifierInfo();
+      FuncLoc = Tok.getLocation();
+      PP.Lex(Tok); // consume identifier and move to next token
+    } else {
+      // Not an identifier after comma; restore and don't consume the comma.
+      // The r_paren check below will produce the expected diagnostics.
+      Tok = SavedTok;
+    }
   }
 
   if (!SegmentName) {
@@ -1273,7 +1302,8 @@ bool Parser::HandlePragmaMSInitSeg(StringRef PragmaName,
                        PragmaName))
     return false;
 
-  Actions.ActOnPragmaMSInitSeg(PragmaLocation, SegmentName);
+  Actions.ActOnPragmaMSInitSeg(PragmaLocation, SegmentName, FuncII,
+                               FuncLoc);
   return true;
 }
 

clang/lib/Sema/Sema.cpp

@@ -313,7 +313,7 @@ Sema::Sema(Preprocessor &pp, ASTContext &ctxt, ASTConsumer &consumer,
       DataSegStack(nullptr), BSSSegStack(nullptr), ConstSegStack(nullptr),
       CodeSegStack(nullptr), StrictGuardStackCheckStack(false),
       FpPragmaStack(FPOptionsOverride()), CurInitSeg(nullptr),
-      VisContext(nullptr), PragmaAttributeCurrentTargetDecl(nullptr),
+  CurInitSegFn(nullptr), CurInitSegFnLoc(), VisContext(nullptr), PragmaAttributeCurrentTargetDecl(nullptr),
       StdCoroutineTraitsCache(nullptr), IdResolver(pp),
       OriginalLexicalContext(nullptr), StdInitializerList(nullptr),
       StdTypeIdentity(nullptr),

clang/lib/Sema/SemaAttr.cpp

@@ -888,12 +888,17 @@ void Sema::ActOnPragmaMSSection(SourceLocation PragmaLocation,
 }
 
 void Sema::ActOnPragmaMSInitSeg(SourceLocation PragmaLocation,
-                                StringLiteral *SegmentName) {
-  // There's no stack to maintain, so we just have a current section.  When we
-  // see the default section, reset our current section back to null so we stop
-  // tacking on unnecessary attributes.
-  CurInitSeg = SegmentName->getString() == ".CRT$XCU" ? nullptr : SegmentName;
+                                StringLiteral *SegmentName,
+                                IdentifierInfo *Func,
+                                SourceLocation FuncLoc) {
+
+  // When we see the default section, reset back to null.
+  CurInitSeg = (SegmentName->getString() == ".CRT$XCU") ? nullptr : SegmentName;
   CurInitSegLoc = PragmaLocation;
+
+  // Store optional function identifier (may be null).
+  CurInitSegFn = Func;
+  CurInitSegFnLoc = Func ? FuncLoc : SourceLocation();
 }
 
 void Sema::ActOnPragmaMSAllocText(

clang/test/SemaCXX/pragma-init_seg.cpp

@@ -17,6 +17,9 @@
 #pragma init_seg("a" L"b") // expected-warning {{expected non-wide string literal in '#pragma init_seg'}}
 
 #pragma init_seg(compiler)
+// Test optional func-name parsing: declare an identifier and use it.
+int myexit();
+#pragma init_seg(compiler, myexit)
 #else
 #pragma init_seg(compiler) // expected-warning {{'#pragma init_seg' is only supported when targeting a Microsoft environment}}
 #endif