[llvm] [llvm-ml] Remove unsafe getCurrentSegmentOnly() call (PR #123355)

[Eric Astor via llvm-commits]

https://github.com/ericastor updated https://github.com/llvm/llvm-project/pull/123355

>From e8b956df20b850985f8a55d0a3408d632a35f92d Mon Sep 17 00:00:00 2001
From: Eric Astor <epastor at google.com>
Date: Fri, 17 Jan 2025 15:27:52 +0000
Subject: [PATCH 1/2] [llvm-ml] Remove unsafe getCurrentSegmentOnly() call

---
 llvm/lib/MC/MCParser/MasmParser.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/MC/MCParser/MasmParser.cpp b/llvm/lib/MC/MCParser/MasmParser.cpp
index 78261c1f9fedb2..6c49a17180f27b 100644
--- a/llvm/lib/MC/MCParser/MasmParser.cpp
+++ b/llvm/lib/MC/MCParser/MasmParser.cpp
@@ -1454,7 +1454,7 @@ bool MasmParser::Run(bool NoInitialTextSection, bool NoFinalize) {
 }
 
 bool MasmParser::checkForValidSection() {
-  if (!ParsingMSInlineAsm && !getStreamer().getCurrentSectionOnly()) {
+  if (!ParsingMSInlineAsm && !getStreamer().getCurrentFragment()) {
     Out.initSections(false, getTargetParser().getSTI());
     return Error(getTok().getLoc(),
                  "expected section directive before assembly directive");

>From 70d25e7872d9a9a4e676b17f688ca803691853a4 Mon Sep 17 00:00:00 2001
From: Eric Astor <epastor at google.com>
Date: Fri, 17 Jan 2025 20:19:56 +0000
Subject: [PATCH 2/2] Add tests, and more fixes

---
 llvm/lib/MC/MCParser/COFFMasmParser.cpp | 5 +++++
 llvm/lib/MC/MCParser/MasmParser.cpp     | 3 ++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/MC/MCParser/COFFMasmParser.cpp b/llvm/lib/MC/MCParser/COFFMasmParser.cpp
index c323e64a40aee6..c10751e514ced6 100644
--- a/llvm/lib/MC/MCParser/COFFMasmParser.cpp
+++ b/llvm/lib/MC/MCParser/COFFMasmParser.cpp
@@ -441,6 +441,11 @@ bool COFFMasmParser::parseDirectiveOption(StringRef Directive, SMLoc Loc) {
 ///          statements
 ///      label "endproc"
 bool COFFMasmParser::parseDirectiveProc(StringRef Directive, SMLoc Loc) {
+  if (!getStreamer().getCurrentFragment()) {
+    return Error(getTok().getLoc(), "expected section directive before '" +
+                                        Directive + "' directive");
+  }
+
   StringRef Label;
   if (getParser().parseIdentifier(Label))
     return Error(Loc, "expected identifier for procedure");
diff --git a/llvm/lib/MC/MCParser/MasmParser.cpp b/llvm/lib/MC/MCParser/MasmParser.cpp
index 6c49a17180f27b..b2c956e0a45981 100644
--- a/llvm/lib/MC/MCParser/MasmParser.cpp
+++ b/llvm/lib/MC/MCParser/MasmParser.cpp
@@ -1454,7 +1454,8 @@ bool MasmParser::Run(bool NoInitialTextSection, bool NoFinalize) {
 }
 
 bool MasmParser::checkForValidSection() {
-  if (!ParsingMSInlineAsm && !getStreamer().getCurrentFragment()) {
+  if (!ParsingMSInlineAsm && !(getStreamer().getCurrentFragment() &&
+                               getStreamer().getCurrentSectionOnly())) {
     Out.initSections(false, getTargetParser().getSTI());
     return Error(getTok().getLoc(),
                  "expected section directive before assembly directive");