[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/3] [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/3] 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");

>From f46aa7b78a8ac0f4bf95af1260183877339cdcd5 Mon Sep 17 00:00:00 2001
From: Eric Astor <epastor at google.com>
Date: Fri, 17 Jan 2025 20:21:48 +0000
Subject: [PATCH 3/3] Actually add tests this time

---
 llvm/test/tools/llvm-ml/bare_proc_error.asm  | 7 +++++++
 llvm/test/tools/llvm-ml/no_section_error.asm | 4 ++++
 2 files changed, 11 insertions(+)
 create mode 100644 llvm/test/tools/llvm-ml/bare_proc_error.asm
 create mode 100644 llvm/test/tools/llvm-ml/no_section_error.asm

diff --git a/llvm/test/tools/llvm-ml/bare_proc_error.asm b/llvm/test/tools/llvm-ml/bare_proc_error.asm
new file mode 100644
index 00000000000000..582abeaab76b8f
--- /dev/null
+++ b/llvm/test/tools/llvm-ml/bare_proc_error.asm
@@ -0,0 +1,7 @@
+; RUN: not llvm-ml -filetype=s %s /Fo /dev/null 2>&1 | FileCheck %s
+
+; CHECK: :[[# @LINE + 1]]:1: error: expected section directive before 'PROC' directive
+foo PROC
+; CHECK: :[[# @LINE + 1]]:6: error: expected section directive before assembly directive
+  ret
+foo ENDP
diff --git a/llvm/test/tools/llvm-ml/no_section_error.asm b/llvm/test/tools/llvm-ml/no_section_error.asm
new file mode 100644
index 00000000000000..65c111908b81aa
--- /dev/null
+++ b/llvm/test/tools/llvm-ml/no_section_error.asm
@@ -0,0 +1,4 @@
+; RUN: not llvm-ml -filetype=s %s /Fo /dev/null 2>&1 | FileCheck %s
+
+; CHECK: :[[# @LINE + 1]]:6: error: expected section directive before assembly directive in 'BYTE' directive
+BYTE 2, 3, 4