From ba2cbe1e5570a9b10f33cd6a0e57c0759c9d00d7 Mon Sep 17 00:00:00 2001
From: Brian Atkinson <bca@google.com>
Date: Wed, 19 Feb 2020 16:25:30 -0800
Subject: [PATCH] Avoid direct access to the env var
 TEST_UNDECLARED_OUTPUTS_DIR.

On Windows, Bazel populates this path with `/`s only making proper path
management impossible without sanitizing the path up front. This changes to
accessing the env var through an indirection layer which will fix path problems
on Windows when the codebase is ready to switch over.

PiperOrigin-RevId: 296083765
Change-Id: I26bbaf83ba5e3fafd3ab0a0de08f6cb597b94477
---
 .../mlir/tensorflow/utils/dump_mlir_util.cc   | 30 +++++++++----------
 .../mlir/tensorflow/utils/dump_mlir_util.h    |  2 +-
 .../compiler/tf2xla/mlir_bridge_pass.cc       |  5 ++--
 3 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc b/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc
index ead26c8f17d..f06734a26bd 100644
--- a/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc
+++ b/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.cc
@@ -27,6 +27,7 @@ limitations under the License.
 #include "mlir/IR/Operation.h"  // TF:llvm-project
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/path.h"
 
 namespace tensorflow {
 
@@ -97,18 +98,18 @@ struct WritableFileRawStream : public llvm::raw_ostream {
 Status CreateFileForDumping(llvm::StringRef name,
                             std::unique_ptr<llvm::raw_ostream>* os,
                             std::string* filepath, llvm::StringRef dirname) {
-  const char* dir = nullptr;
+  std::string dir;
   if (!dirname.empty())
-    dir = dirname.data();
+    dir = std::string(dirname);
   else
     dir = GetDumpDirFromEnvVar();
 
-  if (!dir) {
+  if (dir.empty()) {
     return Status(error::Code::INVALID_ARGUMENT,
                   "(TF_DUMP_GRAPH_PREFIX not specified)");
   }
 
-  if (std::strncmp(dir, "-", 2) == 0) {
+  if (dir == "-") {
     *os = std::make_unique<LogInfoRawStream>();
     *filepath = "LOG(INFO)";
     return Status();
@@ -151,25 +152,24 @@ std::string DumpMlirOpToFile(llvm::StringRef name, mlir::Operation* op,
   return filepath;
 }
 
-const char* GetDumpDirFromEnvVar() {
+std::string GetDumpDirFromEnvVar() {
   const char* prefix_env = getenv("TF_DUMP_GRAPH_PREFIX");
   if (!prefix_env) {
     LOG(WARNING)
         << "Failed to dump MLIR module because dump location is not "
         << " specified through TF_DUMP_GRAPH_PREFIX environment variable.";
-    return nullptr;
+    return "";
   }
 
-  if (absl::EqualsIgnoreCase(prefix_env, "sponge")) {
-    const char* tmp_dir = getenv("TEST_UNDECLARED_OUTPUTS_DIR");
-    if (!tmp_dir) {
-      LOG(WARNING) << "TF_DUMP_GRAPH_PREFIX=sponge but "
-                      "TEST_UNDECLARED_OUTPUT_DIRS is not set";
-      return nullptr;
-    }
-    return tmp_dir;
+  std::string result = prefix_env;
+
+  if (absl::EqualsIgnoreCase(result, "sponge") &&
+      !io::GetTestUndeclaredOutputsDir(&result)) {
+    LOG(WARNING) << "TF_DUMP_GRAPH_PREFIX=sponge but "
+                    "TEST_UNDECLARED_OUTPUT_DIRS is not set";
+    return "";
   }
-  return prefix_env;
+  return result;
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h b/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h
index 7c25a809089..14c0d1f0b6e 100644
--- a/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h
+++ b/tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h
@@ -54,7 +54,7 @@ std::string DumpMlirOpToFile(llvm::StringRef name, mlir::Operation* op,
 // Default is reading from TF_DUMP_GRAPH_PREFIX, and if the string is 'sponge'
 // read from TEST_UNDECLARED_OUTPUTS_DIR. Returns nullptr if the directory
 // cannot be determined and generates a warning message.
-const char* GetDumpDirFromEnvVar();
+std::string GetDumpDirFromEnvVar();
 
 }  // namespace tensorflow
 
diff --git a/tensorflow/compiler/tf2xla/mlir_bridge_pass.cc b/tensorflow/compiler/tf2xla/mlir_bridge_pass.cc
index a0ffd1908c5..7ac4cb8fb06 100644
--- a/tensorflow/compiler/tf2xla/mlir_bridge_pass.cc
+++ b/tensorflow/compiler/tf2xla/mlir_bridge_pass.cc
@@ -35,11 +35,10 @@ namespace tensorflow {
 // This require the TF_DUMP_GRAPH_PREFIX to be set to a path that exist (or can
 // be created).
 static void DumpModule(mlir::ModuleOp module, llvm::StringRef file_prefix) {
-  const char* prefix_env = GetDumpDirFromEnvVar();
-  if (!prefix_env) {
+  std::string prefix = GetDumpDirFromEnvVar();
+  if (prefix.empty()) {
     return;
   }
-  std::string prefix = prefix_env;
 
   auto* env = tensorflow::Env::Default();
   auto status = env->RecursivelyCreateDir(prefix);