[yaml2obj] Add -D k=v to preprocess the input YAML

Examples:

```
yaml2obj -D MACHINE=EM_386 a.yaml -o a.o
yaml2obj -D MACHINE=0x1234 a.yaml -o a.o
```

where a.yaml contains:

```
--- !ELF
FileHeader:
  Class:   ELFCLASS64
  Data:    ELFDATA2MSB
  Type:    ET_REL
  Machine: [[MACHINE]]
```

Reviewed By: grimar, jhenderson

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

2 files changed, 107 insertions, 1 deletions

diff --git a/llvm/test/tools/yaml2obj/macro.yaml b/llvm/test/tools/yaml2obj/macro.yaml
new file mode 100644
index 0000000..d7834fa
--- /dev/null
+++ b/llvm/test/tools/yaml2obj/macro.yaml
@@ -0,0 +1,60 @@
+# RUN: yaml2obj -D NAME= %s | llvm-nm - | FileCheck --check-prefix=FBAR %s
+# FBAR: U fbar
+
+# RUN: yaml2obj -D NAME=o %s | llvm-nm - | FileCheck --check-prefix=FOOBAR %s
+# RUN: yaml2obj -DNAME=o %s | llvm-nm - | FileCheck --check-prefix=FOOBAR %s
+# FOOBAR: U foobar
+
+# RUN: not yaml2obj -D NAME %s 2>&1 | FileCheck --check-prefix=ERR1 %s
+# ERR1: error: invalid syntax for -D: NAME
+
+# RUN: not yaml2obj -D =value %s 2>&1 | FileCheck --check-prefix=ERR2 %s
+# ERR2: error: invalid syntax for -D: =value
+
+# RUN: not yaml2obj -D NAME=a -D NAME=b %s 2>&1 | FileCheck --check-prefix=ERR3 %s
+# ERR3: error: 'NAME' redefined
+
+--- !ELF
+FileHeader:
+  Class:   ELFCLASS64
+  Data:    ELFDATA2LSB
+  Type:    ET_REL
+  Machine: EM_X86_64
+Symbols:
+  - Name: f[[NAME]][[NAME]]bar
+
+## Digits can appear in macro names. Test that we don't expand macros recursively.
+# RUN: yaml2obj --docnum=2 -D a0='[[a1]]' -D a1='[[a0]]' %s | llvm-nm --just-symbol-name --no-sort - > %t.recur
+# RUN: echo -e '[[a1]]\n[[a0]]' > %t0.recur
+# RUN: diff -u %t0.recur %t.recur
+
+--- !ELF
+FileHeader:
+  Class:   ELFCLASS64
+  Data:    ELFDATA2LSB
+  Type:    ET_REL
+  Machine: EM_X86_64
+Symbols:
+  - Name: "[[a0]]"
+  - Name: "[[a1]]"
+
+## Test unterminated [[.
+# RUN: yaml2obj --docnum=3 %s | llvm-nm --just-symbol-name --no-sort - > %t.nosubst
+# RUN: echo -e 'a[[\n[[a]\n[[a[[a]]\n[[a][[a]][[b]]' > %t0.nosubst
+# RUN: diff -u %t0.nosubst %t.nosubst
+
+# RUN: yaml2obj --docnum=3 -D a=b -D b=c %s | llvm-nm --just-symbol-name --no-sort - > %t.subst
+# RUN: echo -e 'a[[\n[[a]\n[[ab\n[[a]bc' > %t0.subst
+# RUN: diff -u %t0.subst %t.subst
+
+--- !ELF
+FileHeader:
+  Class:   ELFCLASS64
+  Data:    ELFDATA2LSB
+  Type:    ET_REL
+  Machine: EM_X86_64
+Symbols:
+  - Name: "a[["
+  - Name: "[[a]"
+  - Name: "[[a[[a]]"
+  - Name: "[[a][[a]][[b]]"
diff --git a/llvm/tools/yaml2obj/yaml2obj.cpp b/llvm/tools/yaml2obj/yaml2obj.cpp
index dd92e01..8a23b22 100644
--- a/llvm/tools/yaml2obj/yaml2obj.cpp
+++ b/llvm/tools/yaml2obj/yaml2obj.cpp
@@ -34,6 +34,11 @@ cl::OptionCategory Cat("yaml2obj Options");
 cl::opt<std::string> Input(cl::Positional, cl::desc("<input file>"),
                            cl::init("-"), cl::cat(Cat));
 
+cl::list<std::string>
+    D("D", cl::Prefix,
+      cl::desc("Defined the specified macros to their specified "
+               "definition. The syntax is <macro>=<definition>"));
+
 cl::opt<unsigned>
     DocNum("docnum", cl::init(1),
            cl::desc("Read specified document from input (default = 1)"),
@@ -44,6 +49,44 @@ cl::opt<std::string> OutputFilename("o", cl::desc("Output filename"),
                                     cl::Prefix, cl::cat(Cat));
 } // namespace
 
+static Optional<std::string> preprocess(StringRef Buf,
+                                        yaml::ErrorHandler ErrHandler) {
+  DenseMap<StringRef, StringRef> Defines;
+  for (StringRef Define : D) {
+    StringRef Macro, Definition;
+    std::tie(Macro, Definition) = Define.split('=');
+    if (!Define.count('=') || Macro.empty()) {
+      ErrHandler("invalid syntax for -D: " + Define);
+      return {};
+    }
+    if (!Defines.try_emplace(Macro, Definition).second) {
+      ErrHandler("'" + Macro + "'" + " redefined");
+      return {};
+    }
+  }
+
+  std::string Preprocessed;
+  while (!Buf.empty()) {
+    if (Buf.startswith("[[")) {
+      size_t I = Buf.find_first_of("[]", 2);
+      if (Buf.substr(I).startswith("]]")) {
+        StringRef Macro = Buf.substr(2, I - 2);
+        auto It = Defines.find(Macro);
+        if (It != Defines.end()) {
+          Preprocessed += It->second;
+          Buf = Buf.substr(I + 2);
+          continue;
+        }
+      }
+    }
+
+    Preprocessed += Buf[0];
+    Buf = Buf.substr(1);
+  }
+
+  return Preprocessed;
+}
+
 int main(int argc, char **argv) {
   InitLLVM X(argc, argv);
   cl::HideUnrelatedOptions(Cat);
@@ -68,7 +111,10 @@ int main(int argc, char **argv) {
   if (!Buf)
     return 1;
 
-  yaml::Input YIn(Buf.get()->getBuffer());
+  Optional<std::string> Buffer = preprocess(Buf.get()->getBuffer(), ErrHandler);
+  if (!Buffer)
+    return 1;
+  yaml::Input YIn(*Buffer);
   if (!convertYAML(YIn, Out->os(), ErrHandler, DocNum))
     return 1;