[ms] [llvm-ml] Support command-line defines

Enable command-line defines as textmacros

Reviewed By: thakis

Differential Revision: https://reviews.llvm.org/D90059

include/llvm/MC/MCParser/MCAsmParser.h

@@ -184,6 +184,8 @@ public:
 
   virtual bool isParsingMasm() const { return false; }
 
+  virtual bool defineMacro(StringRef Name, StringRef Value) { return true; }
+
   virtual bool lookUpField(StringRef Name, AsmFieldInfo &Info) const {
     return true;
   }

lib/MC/MCParser/MasmParser.cpp

@@ -501,6 +501,8 @@ public:
 
   bool isParsingMasm() const override { return true; }
 
+  bool defineMacro(StringRef Name, StringRef Value) override;
+
   bool lookUpField(StringRef Name, AsmFieldInfo &Info) const override;
   bool lookUpField(StringRef Base, StringRef Member,
                    AsmFieldInfo &Info) const override;
@@ -6905,6 +6907,19 @@ static int rewritesSort(const AsmRewrite *AsmRewriteA,
   llvm_unreachable("Unstable rewrite sort.");
 }
 
+bool MasmParser::defineMacro(StringRef Name, StringRef Value) {
+  Variable &Var = Variables[Name.lower()];
+  if (Var.Name.empty()) {
+    Var.Name = Name;
+  } else if (!Var.Redefinable) {
+    return TokError("invalid variable redefinition");
+  }
+  Var.Redefinable = true;
+  Var.IsText = true;
+  Var.TextValue = Value.str();
+  return false;
+}
+
 bool MasmParser::lookUpField(StringRef Name, AsmFieldInfo &Info) const {
   const std::pair<StringRef, StringRef> BaseMember = Name.split('.');
   const StringRef Base = BaseMember.first, Member = BaseMember.second;

test/tools/llvm-ml/command_line_defines.asm

@@ -0,0 +1,38 @@
+; RUN: llvm-ml -filetype=s %s /Fo - /DT1=test1 /D T2=test2 | FileCheck %s
+
+.code
+
+t1:
+  ret
+; CHECK-NOT: t1:
+; CHECK-LABEL: test1:
+; CHECK-NOT: t1:
+
+t2:
+  ret
+; CHECK-NOT: t2:
+; CHECK-LABEL: test2:
+; CHECK-NOT: t2:
+
+t3:
+ifdef t1
+  xor eax, eax
+endif
+  ret
+; CHECK-LABEL: t3:
+; CHECK: xor eax, eax
+; CHECK: ret
+
+t4:
+ifdef undefined
+  xor eax, eax
+elseifdef t2
+  xor ebx, ebx
+endif
+  ret
+; CHECK-LABEL: t4:
+; CHECK-NOT: xor eax, eax
+; CHECK: xor ebx, ebx
+; CHECK: ret
+
+end

tools/llvm-ml/Opts.td

@@ -31,6 +31,9 @@ def help : MLFlag<"?">,
            HelpText<"Display available options">;
 def help_long : MLFlag<"help">, Alias<help>;
 def assemble_only : MLFlag<"c">, HelpText<"Assemble only; do not link">;
+def define : MLJoinedOrSeparate<"D">, MetaVarName<"<macro>=<value>">,
+             HelpText<"Define <macro> to <value> (or blank if <value> "
+                      "omitted)">;
 def output_file : MLJoinedOrSeparate<"Fo">, HelpText<"Names the output file">;
 def include_path : MLJoinedOrSeparate<"I">,
                    HelpText<"Sets path for include files">;
@@ -72,7 +75,6 @@ def coff_object_file : UnsupportedFlag<"coff">, HelpText<"">;
 def preserve_identifier_case : UnsupportedFlag<"Cp">, HelpText<"">;
 def uppercase_identifiers : UnsupportedFlag<"Cu">, HelpText<"">;
 def preserve_extern_case : UnsupportedFlag<"Cx">, HelpText<"">;
-def define : UnsupportedJoinedOrSeparate<"D">, HelpText<"">;
 def output_preprocessed : UnsupportedFlag<"EP">, HelpText<"">;
 def errorreport : UnsupportedJoined<"ERRORREPORT">, HelpText<"">;
 def stacksize : UnsupportedSeparate<"F">, HelpText<"">;

tools/llvm-ml/llvm-ml.cpp

@@ -147,6 +147,17 @@ static int AssembleInput(StringRef ProgName, const Target *TheTarget,
   Parser->getLexer().setLexMasmHexFloats(true);
   Parser->getLexer().setLexMasmStrings(true);
 
+  auto Defines = InputArgs.getAllArgValues(OPT_define);
+  for (StringRef Define : Defines) {
+    const auto NameValue = Define.split('=');
+    StringRef Name = NameValue.first, Value = NameValue.second;
+    if (Parser->defineMacro(Name, Value)) {
+      WithColor::error(errs(), ProgName)
+          << "can't define macro '" << Name << "' = '" << Value << "'\n";
+      return 1;
+    }
+  }
+
   int Res = Parser->Run(/*NoInitialTextSection=*/true);
 
   return Res;