From 94fae7a581d6772e9f4ec0270485a4308eb8d2dc Mon Sep 17 00:00:00 2001
From: Eric Astor <epastor@google.com>
Date: Thu, 11 Aug 2022 14:18:32 +0000
Subject: [PATCH] [ms] [llvm-ml] Add support for nested PROC/ENDP pairs

This is believed to match behavior by ML.EXE and ML64.EXE.

Reviewed By: thakis

Differential Revision: https://reviews.llvm.org/D131522
---
 llvm/lib/MC/MCParser/COFFMasmParser.cpp       | 21 ++++++-------
 llvm/test/tools/llvm-ml/nested_proc.asm       | 30 +++++++++++++++++++
 .../test/tools/llvm-ml/nested_proc_errors.asm | 13 ++++++++
 3 files changed, 54 insertions(+), 10 deletions(-)
 create mode 100644 llvm/test/tools/llvm-ml/nested_proc.asm
 create mode 100644 llvm/test/tools/llvm-ml/nested_proc_errors.asm

diff --git a/llvm/lib/MC/MCParser/COFFMasmParser.cpp b/llvm/lib/MC/MCParser/COFFMasmParser.cpp
index f2165b486449..34aa5bf2ae39 100644
--- a/llvm/lib/MC/MCParser/COFFMasmParser.cpp
+++ b/llvm/lib/MC/MCParser/COFFMasmParser.cpp
@@ -207,8 +207,9 @@ class COFFMasmParser : public MCAsmParserExtension {
                               SectionKind::getBSS());
   }
 
-  StringRef CurrentProcedure;
-  bool CurrentProcedureFramed;
+  /// Stack of active procedure definitions.
+  SmallVector<StringRef, 1> CurrentProcedures;
+  SmallVector<bool, 1> CurrentProceduresFramed;
 
 public:
   COFFMasmParser() = default;
@@ -478,8 +479,8 @@ bool COFFMasmParser::ParseDirectiveProc(StringRef Directive, SMLoc Loc) {
   }
   getStreamer().emitLabel(Sym, Loc);
 
-  CurrentProcedure = Label;
-  CurrentProcedureFramed = Framed;
+  CurrentProcedures.push_back(Label);
+  CurrentProceduresFramed.push_back(Framed);
   return false;
 }
 bool COFFMasmParser::ParseDirectiveEndProc(StringRef Directive, SMLoc Loc) {
@@ -488,17 +489,17 @@ bool COFFMasmParser::ParseDirectiveEndProc(StringRef Directive, SMLoc Loc) {
   if (getParser().parseIdentifier(Label))
     return Error(LabelLoc, "expected identifier for procedure end");
 
-  if (CurrentProcedure.empty())
+  if (CurrentProcedures.empty())
     return Error(Loc, "endp outside of procedure block");
-  else if (CurrentProcedure != Label)
+  else if (!CurrentProcedures.back().equals_insensitive(Label))
     return Error(LabelLoc, "endp does not match current procedure '" +
-                               CurrentProcedure + "'");
+                               CurrentProcedures.back() + "'");
 
-  if (CurrentProcedureFramed) {
+  if (CurrentProceduresFramed.back()) {
     getStreamer().emitWinCFIEndProc(Loc);
   }
-  CurrentProcedure = "";
-  CurrentProcedureFramed = false;
+  CurrentProcedures.pop_back();
+  CurrentProceduresFramed.pop_back();
   return false;
 }
 
diff --git a/llvm/test/tools/llvm-ml/nested_proc.asm b/llvm/test/tools/llvm-ml/nested_proc.asm
new file mode 100644
index 000000000000..540541f597fc
--- /dev/null
+++ b/llvm/test/tools/llvm-ml/nested_proc.asm
@@ -0,0 +1,30 @@
+; RUN: llvm-ml -m32 -filetype=s %s /Fo - | FileCheck %s
+; RUN: llvm-ml -m64 -filetype=s %s /Fo - | FileCheck %s
+
+.code
+
+t1 PROC
+  xor eax, eax
+t1_nested PROC
+  ret
+t1_nested ENDP
+t1 ENDP
+
+; CHECK-LABEL: t1:
+; CHECK: xor eax, eax
+; CHECK: t1_nested:
+; CHECK: ret
+
+t2 PROC
+  xor eax, eax
+t2_nested PROC
+  ret
+T2_nEsTeD ENDP
+t2 ENDP
+
+; CHECK-LABEL: t2:
+; CHECK: xor eax, eax
+; CHECK: t2_nested:
+; CHECK: ret
+
+END
diff --git a/llvm/test/tools/llvm-ml/nested_proc_errors.asm b/llvm/test/tools/llvm-ml/nested_proc_errors.asm
new file mode 100644
index 000000000000..32356e9cf553
--- /dev/null
+++ b/llvm/test/tools/llvm-ml/nested_proc_errors.asm
@@ -0,0 +1,13 @@
+; RUN: not llvm-ml -filetype=s %s /Fo - 2>&1 | FileCheck %s --implicit-check-not=error:
+
+.code
+
+t1 PROC
+  xor eax, eax
+t1_nested PROC
+  ret
+t1 ENDP
+t1_nested ENDP
+; CHECK: :[[# @LINE - 2]]:1: error: endp does not match current procedure 't1_nested'
+
+END