From 4b5317e937ac5c87c30a67bc76a4e253c2c3516c Mon Sep 17 00:00:00 2001
From: Eric Astor <epastor@google.com>
Date: Wed, 9 Jun 2021 16:05:12 -0400
Subject: [PATCH] [ms] [llvm-ml] Add support for INCLUDE environment variable

Also adds support for the ML.exe command-line flag /X, which ignores the INCLUDE environment variable.

This relands commit c43f413b01b021a8f7b6fce013296114fa92a245 using lit's cross-platform `env` support.

Differential Revision: https://reviews.llvm.org/D103989
---
 .../test/tools/llvm-ml/include_by_env_var.asm | 20 +++++++++++++++++++
 .../llvm-ml/include_by_env_var_errors.asm     | 16 +++++++++++++++
 llvm/tools/llvm-ml/Opts.td                    |  3 ++-
 llvm/tools/llvm-ml/llvm-ml.cpp                | 18 +++++++++++++++--
 4 files changed, 54 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/tools/llvm-ml/include_by_env_var.asm
 create mode 100644 llvm/test/tools/llvm-ml/include_by_env_var_errors.asm

diff --git a/llvm/test/tools/llvm-ml/include_by_env_var.asm b/llvm/test/tools/llvm-ml/include_by_env_var.asm
new file mode 100644
index 000000000000..6f4b269d5807
--- /dev/null
+++ b/llvm/test/tools/llvm-ml/include_by_env_var.asm
@@ -0,0 +1,20 @@
+; RUN: env INCLUDE=%S llvm-ml -filetype=s %s /Fo - | FileCheck %s
+
+include included.inc
+
+.code
+
+t1:
+mov eax, Const
+
+; CHECK-LABEL: t1:
+; CHECK-NEXT: mov eax, 8
+
+t2:
+push_pop ebx
+
+; CHECK-LABEL: t2:
+; CHECK-NEXT: push ebx
+; CHECK-NEXT: pop ebx
+
+end
diff --git a/llvm/test/tools/llvm-ml/include_by_env_var_errors.asm b/llvm/test/tools/llvm-ml/include_by_env_var_errors.asm
new file mode 100644
index 000000000000..7c2b5e8356b9
--- /dev/null
+++ b/llvm/test/tools/llvm-ml/include_by_env_var_errors.asm
@@ -0,0 +1,16 @@
+; RUN: not llvm-ml -filetype=s %s /Fo - 2>&1 | FileCheck %s --implicit-check-not=error:
+; RUN: env INCLUDE=%S not llvm-ml -filetype=s %s /X /Fo - 2>&1 | FileCheck %s --implicit-check-not=error:
+
+; CHECK: :[[# @LINE + 1]]:9: error: Could not find include file 'included.inc'
+include included.inc
+
+.code
+
+t1:
+mov eax, Const
+
+t2:
+; CHECK: :[[# @LINE + 1]]:1: error: invalid instruction mnemonic 'push_pop'
+push_pop ebx
+
+end
diff --git a/llvm/tools/llvm-ml/Opts.td b/llvm/tools/llvm-ml/Opts.td
index b5ba920c8769..06871b49ad0d 100644
--- a/llvm/tools/llvm-ml/Opts.td
+++ b/llvm/tools/llvm-ml/Opts.td
@@ -73,6 +73,8 @@ def assembly_file : MLJoinedOrSeparate<"Ta">,
 def error_on_warning : MLFlag<"WX">, Alias<fatal_warnings>;
 def parse_only : MLFlag<"Zs">, HelpText<"Run a syntax-check only">,
                  Alias<filetype>, AliasArgs<["null"]>;
+def ignore_include_envvar : MLFlag<"X">,
+                            HelpText<"Ignore the INCLUDE environment variable">;
 
 def tiny_model_support : UnsupportedFlag<"AT">, HelpText<"">;
 def alternate_linker : UnsupportedJoined<"Bl">, HelpText<"">;
@@ -105,7 +107,6 @@ def listing_title : UnsupportedSeparate<"St">, HelpText<"">;
 def listing_false_conditionals : UnsupportedFlag<"Sx">, HelpText<"">;
 def extra_warnings : UnsupportedFlag<"w">, HelpText<"">;
 def warning_level : UnsupportedJoined<"W">, HelpText<"">;
-def ignore_include_envvar : UnsupportedFlag<"X">, HelpText<"">;
 def line_number_info : UnsupportedFlag<"Zd">, HelpText<"">;
 def export_all_symbols : UnsupportedFlag<"Zf">, HelpText<"">;
 def codeview_info : UnsupportedFlag<"Zi">, HelpText<"">;
diff --git a/llvm/tools/llvm-ml/llvm-ml.cpp b/llvm/tools/llvm-ml/llvm-ml.cpp
index 14f75d0d814e..1a6686253ff2 100644
--- a/llvm/tools/llvm-ml/llvm-ml.cpp
+++ b/llvm/tools/llvm-ml/llvm-ml.cpp
@@ -36,6 +36,7 @@
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
+#include "llvm/Support/Process.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/TargetRegistry.h"
 #include "llvm/Support/TargetSelect.h"
@@ -263,8 +264,21 @@ int main(int Argc, char **Argv) {
   SrcMgr.AddNewSourceBuffer(std::move(*BufferPtr), SMLoc());
 
   // Record the location of the include directories so that the lexer can find
-  // it later.
-  SrcMgr.setIncludeDirs(InputArgs.getAllArgValues(OPT_include_path));
+  // included files later.
+  std::vector<std::string> IncludeDirs =
+      InputArgs.getAllArgValues(OPT_include_path);
+  if (!InputArgs.hasArg(OPT_ignore_include_envvar)) {
+    if (llvm::Optional<std::string> IncludeEnvVar =
+            llvm::sys::Process::GetEnv("INCLUDE")) {
+      SmallVector<StringRef, 8> Dirs;
+      StringRef(*IncludeEnvVar)
+          .split(Dirs, ";", /*MaxSplit=*/-1, /*KeepEmpty=*/false);
+      IncludeDirs.reserve(IncludeDirs.size() + Dirs.size());
+      for (StringRef Dir : Dirs)
+        IncludeDirs.push_back(Dir.str());
+    }
+  }
+  SrcMgr.setIncludeDirs(IncludeDirs);
 
   std::unique_ptr<MCRegisterInfo> MRI(TheTarget->createMCRegInfo(TripleName));
   assert(MRI && "Unable to create target register info!");