From 63538a087926f9f9dc992918b06be4720f4c59fb Mon Sep 17 00:00:00 2001
From: cynecx <me@cynecx.net>
Date: Wed, 21 Jun 2023 10:39:56 -0700
Subject: [PATCH] [MC] Add .pushsection/.popsection support to COFFAsmParser

The COFFAsmParser (to my surprise) didn't support the .pushsection and
.popsection directives. These directives aren't directly useful, however for
frontends that have inline asm support this is really useful. Rust in
particular, has support for inline asm, which can be used together with these
directives to "emulate" features like static generics. This patch adds support
for the two mentioned directives.

Reviewed By: MaskRay

Differential Revision: https://reviews.llvm.org/D152085
---
 llvm/lib/MC/MCParser/COFFAsmParser.cpp | 31 ++++++++++++++-
 llvm/test/MC/COFF/section.s            | 72 ++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/MC/MCParser/COFFAsmParser.cpp b/llvm/lib/MC/MCParser/COFFAsmParser.cpp
index 3f88257..bfded36 100644
--- a/llvm/lib/MC/MCParser/COFFAsmParser.cpp
+++ b/llvm/lib/MC/MCParser/COFFAsmParser.cpp
@@ -56,6 +56,10 @@ class COFFAsmParser : public MCAsmParserExtension {
     addDirectiveHandler<&COFFAsmParser::ParseSectionDirectiveData>(".data");
     addDirectiveHandler<&COFFAsmParser::ParseSectionDirectiveBSS>(".bss");
     addDirectiveHandler<&COFFAsmParser::ParseDirectiveSection>(".section");
+    addDirectiveHandler<&COFFAsmParser::ParseDirectivePushSection>(
+        ".pushsection");
+    addDirectiveHandler<&COFFAsmParser::ParseDirectivePopSection>(
+        ".popsection");
     addDirectiveHandler<&COFFAsmParser::ParseDirectiveDef>(".def");
     addDirectiveHandler<&COFFAsmParser::ParseDirectiveScl>(".scl");
     addDirectiveHandler<&COFFAsmParser::ParseDirectiveType>(".type");
@@ -115,6 +119,9 @@ class COFFAsmParser : public MCAsmParserExtension {
   }
 
   bool ParseDirectiveSection(StringRef, SMLoc);
+  bool parseSectionArguments(StringRef, SMLoc);
+  bool ParseDirectivePushSection(StringRef, SMLoc);
+  bool ParseDirectivePopSection(StringRef, SMLoc);
   bool ParseDirectiveDef(StringRef, SMLoc);
   bool ParseDirectiveScl(StringRef, SMLoc);
   bool ParseDirectiveType(StringRef, SMLoc);
@@ -343,7 +350,12 @@ bool COFFAsmParser::ParseSectionName(StringRef &SectionName) {
   return false;
 }
 
+bool COFFAsmParser::ParseDirectiveSection(StringRef directive, SMLoc loc) {
+  return parseSectionArguments(directive, loc);
+}
+
 // .section name [, "flags"] [, identifier [ identifier ], identifier]
+// .pushsection <same as above>
 //
 // Supported flags:
 //   a: Ignored.
@@ -358,7 +370,7 @@ bool COFFAsmParser::ParseSectionName(StringRef &SectionName) {
 //   y: Not-readable section (clears 'r')
 //
 // Subsections are not supported.
-bool COFFAsmParser::ParseDirectiveSection(StringRef, SMLoc) {
+bool COFFAsmParser::parseSectionArguments(StringRef, SMLoc) {
   StringRef SectionName;
 
   if (ParseSectionName(SectionName))
@@ -417,6 +429,23 @@ bool COFFAsmParser::ParseDirectiveSection(StringRef, SMLoc) {
   return false;
 }
 
+bool COFFAsmParser::ParseDirectivePushSection(StringRef directive, SMLoc loc) {
+  getStreamer().pushSection();
+
+  if (parseSectionArguments(directive, loc)) {
+    getStreamer().popSection();
+    return true;
+  }
+
+  return false;
+}
+
+bool COFFAsmParser::ParseDirectivePopSection(StringRef, SMLoc) {
+  if (!getStreamer().popSection())
+    return TokError(".popsection without corresponding .pushsection");
+  return false;
+}
+
 bool COFFAsmParser::ParseDirectiveDef(StringRef, SMLoc) {
   StringRef SymbolName;
 
diff --git a/llvm/test/MC/COFF/section.s b/llvm/test/MC/COFF/section.s
index 05a9155..9c1a11e 100644
--- a/llvm/test/MC/COFF/section.s
+++ b/llvm/test/MC/COFF/section.s
@@ -1,5 +1,6 @@
 // RUN: llvm-mc -triple i386-pc-win32 -filetype=obj %s | llvm-readobj -S - | FileCheck %s
 // RUN: llvm-mc -triple x86_64-pc-win32 -filetype=obj %s | llvm-readobj -S - | FileCheck %s
+// RUN: not llvm-mc -triple x86_64-pc-win32 -filetype=obj --defsym ERR=1 %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ERR
 
 .section .foo$bar; .long 1
 .section .foo@bar; .long 1
@@ -208,3 +209,74 @@
 
 // CHECK-NOT:    Section {
 // CHECK:      ]
+
+.section data1; .quad 0
+
+.pushsection data2; .quad 0
+.popsection
+
+// Back to section data1
+.quad 2
+
+// CHECK:       Section {
+// CHECK-NEXT:    Number:
+// CHECK-NEXT:    Name: data1
+// CHECK:         RawDataSize: 16
+
+// CHECK:       Section {
+// CHECK-NEXT:    Number:
+// CHECK-NEXT:    Name: data2
+// CHECK:         RawDataSize: 8
+
+.section .data3,"dw"; .quad 1
+
+.pushsection .data4,"dw"; .quad 1
+.popsection
+
+.pushsection .data5,"dr"; .quad 1
+.popsection
+
+// in section .data3
+.quad 4
+
+// Notice the different section flags here.
+// This shouldn't overwrite the intial section flags.
+.pushsection .data4,"dr"; .quad 1
+.popsection
+
+// CHECK:       Section {
+// CHECK-NEXT:    Number:
+// CHECK-NEXT:    Name: .data3
+// CHECK:         RawDataSize: 16
+// CHECK:         Characteristics [
+// CHECK-NEXT:      IMAGE_SCN_ALIGN_1BYTES
+// CHECK-NEXT:      IMAGE_SCN_CNT_INITIALIZED_DATA
+// CHECK-NEXT:      IMAGE_SCN_MEM_READ
+// CHECK-NEXT:      IMAGE_SCN_MEM_WRITE
+// CHECK-NEXT:    ]
+
+// CHECK:       Section {
+// CHECK-NEXT:    Number:
+// CHECK-NEXT:    Name: .data4
+// CHECK:         RawDataSize: 16
+// CHECK:         Characteristics [
+// CHECK-NEXT:      IMAGE_SCN_ALIGN_1BYTES
+// CHECK-NEXT:      IMAGE_SCN_CNT_INITIALIZED_DATA
+// CHECK-NEXT:      IMAGE_SCN_MEM_READ
+// CHECK-NEXT:      IMAGE_SCN_MEM_WRITE
+// CHECK-NEXT:    ]
+
+// CHECK:       Section {
+// CHECK-NEXT:    Number:
+// CHECK-NEXT:    Name: .data5
+// CHECK:         RawDataSize: 8
+// CHECK:         Characteristics [
+// CHECK-NEXT:      IMAGE_SCN_ALIGN_1BYTES
+// CHECK-NEXT:      IMAGE_SCN_CNT_INITIALIZED_DATA
+// CHECK-NEXT:      IMAGE_SCN_MEM_READ
+// CHECK-NEXT:    ]
+
+.ifdef ERR
+// ERR: :[[#@LINE+1]]:12: error: .popsection without corresponding .pushsection
+.popsection
+.endif
-- 
2.7.4