[rust-project] Record env vars

Some Rust procedural macros depend on the values of environment
variables available at compile time (typically for use inside an
include! directive to pull in generated code).

These environment variables must be made avaiable not just for the
procedural macro build, but for each invocation of it - so they need to
be available in all crates.

Informing rust-analyzer of these environment variables allows it to
expand more procedural macros and results in better source code
expansion.

Bug: chromium/1293933
Change-Id: Ic635d962ccd0782289e7f593118422086cbe8c49
Reviewed-on: https://gn-review.googlesource.com/c/gn/+/13360
Reviewed-by: Brett Wilson <brettw@chromium.org>
Commit-Queue: Brett Wilson <brettw@chromium.org>
diff --git a/src/gn/rust_project_writer.cc b/src/gn/rust_project_writer.cc
index 4e66dcc..eea527c 100644
--- a/src/gn/rust_project_writer.cc
+++ b/src/gn/rust_project_writer.cc
@@ -10,6 +10,7 @@
 #include <tuple>
 
 #include "base/json/string_escape.h"
+#include "base/strings/string_split.h"
 #include "gn/builder.h"
 #include "gn/deps_iterator.h"
 #include "gn/ninja_target_command_util.h"
@@ -329,6 +330,17 @@
     }
   }
 
+  // Note any environment variables. These may be used by proc macros
+  // invoked by the current crate (so we want to record these for all crates,
+  // not just proc macro crates)
+  for (const auto& env_var : target->config_values().rustenv()) {
+    std::vector<std::string> parts = base::SplitString(
+        env_var, "=", base::TRIM_WHITESPACE, base::SPLIT_WANT_ALL);
+    if (parts.size() >= 2) {
+      crate.AddRustenv(parts[0], parts[1]);
+    }
+  }
+
   // Add the rest of the crate dependencies.
   for (const auto& dep : crate_deps) {
     auto idx = lookup[dep];
@@ -436,8 +448,29 @@
       rust_project << "        \"" << escaped_config << "\"";
     }
     rust_project << NEWLINE;
-    rust_project << "      ]" NEWLINE;  // end cfgs
+    rust_project << "      ]";  // end cfgs
 
+    if (!crate.rustenv().empty()) {
+      rust_project << "," NEWLINE;
+      rust_project << "      \"env\": {";
+      bool first_env = true;
+      for (const auto& env : crate.rustenv()) {
+        if (!first_env)
+          rust_project << ",";
+        first_env = false;
+        std::string escaped_key, escaped_val;
+        base::EscapeJSONString(env.first, false, &escaped_key);
+        base::EscapeJSONString(env.second, false, &escaped_val);
+        rust_project << NEWLINE;
+        rust_project << "        \"" << escaped_key << "\": \"" << escaped_val
+                     << "\"";
+      }
+
+      rust_project << NEWLINE;
+      rust_project << "      }" NEWLINE;  // end env vars
+    } else {
+      rust_project << NEWLINE;
+    }
     rust_project << "    }";  // end crate
   }
   rust_project << NEWLINE "  ]" NEWLINE;  // end crate list
diff --git a/src/gn/rust_project_writer_helpers.h b/src/gn/rust_project_writer_helpers.h
index d53ff07..24b13d0 100644
--- a/src/gn/rust_project_writer_helpers.h
+++ b/src/gn/rust_project_writer_helpers.h
@@ -14,6 +14,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include "base/containers/flat_map.h"
 #include "build_settings.h"
 #include "gn/source_file.h"
 #include "gn/target.h"
@@ -44,6 +45,11 @@
   // Add a config item to the crate.
   void AddConfigItem(std::string cfg_item) { configs_.push_back(cfg_item); }
 
+  // Add a key-value environment variable pair used when building this crate.
+  void AddRustenv(std::string key, std::string value) {
+    rustenv_.emplace(key, value);
+  }
+
   // Add another crate as a dependency of this one.
   void AddDependency(CrateIndex index, std::string name) {
     deps_.push_back(std::make_pair(index, name));
@@ -91,6 +97,10 @@
     return proc_macro_dynamic_library_;
   }
 
+  // Returns environment variables applied to this, which may be necessary
+  // for correct functioning of environment variables
+  const base::flat_map<std::string, std::string>& rustenv() { return rustenv_; }
+
  private:
   SourceFile root_;
   CrateIndex index_;
@@ -101,6 +111,7 @@
   std::optional<std::string> compiler_target_;
   std::vector<std::string> compiler_args_;
   std::optional<OutputFile> proc_macro_dynamic_library_;
+  base::flat_map<std::string, std::string> rustenv_;
 };
 
 using CrateList = std::vector<Crate>;
diff --git a/src/gn/rust_project_writer_unittest.cc b/src/gn/rust_project_writer_unittest.cc
index 77a9d34..2d4f0dc 100644
--- a/src/gn/rust_project_writer_unittest.cc
+++ b/src/gn/rust_project_writer_unittest.cc
@@ -535,6 +535,8 @@
   target.sources().push_back(lib);
   target.source_types_used().Set(SourceFile::SOURCE_RS);
   target.rust_values().set_crate_root(lib);
+  target.config_values().rustenv().push_back("TEST_ENV_VAR=baz");
+  target.config_values().rustenv().push_back("TEST_ENV_VAR2=baz2");
   target.rust_values().crate_name() = "foo";
   target.config_values().rustflags().push_back("--cfg=feature=\"foo_enabled\"");
   target.SetToolchain(setup.toolchain());
@@ -568,7 +570,11 @@
       "        \"test\",\n"
       "        \"debug_assertions\",\n"
       "        \"feature=\\\"foo_enabled\\\"\"\n"
-      "      ]\n"
+      "      ],\n"
+      "      \"env\": {\n"
+      "        \"TEST_ENV_VAR\": \"baz\",\n"
+      "        \"TEST_ENV_VAR2\": \"baz2\"\n"
+      "      }\n"
       "    }\n"
       "  ]\n"
       "}\n";