From 07f4e5e51239f6c4d3364be8bdc7ce711a9b6879 Mon Sep 17 00:00:00 2001
From: Shahms King <shahms@google.com>
Date: Fri, 28 Jan 2022 13:51:27 -0800
Subject: [PATCH] fix(rust_extractor): set the working_directory (#5200)

---
 kythe/rust/extractor/src/bin/extractor.rs      | 7 +++++++
 kythe/rust/extractor/tests/integration_test.rs | 2 ++
 2 files changed, 9 insertions(+)

diff --git a/kythe/rust/extractor/src/bin/extractor.rs b/kythe/rust/extractor/src/bin/extractor.rs
index e197732efa..86198ee0e9 100644
--- a/kythe/rust/extractor/src/bin/extractor.rs
+++ b/kythe/rust/extractor/src/bin/extractor.rs
@@ -23,6 +23,7 @@ use extra_actions_base_rust_proto::*;
 use kythe_rust_extractor::vname_util::VNameRule;
 use protobuf::Message;
 use sha2::{Digest, Sha256};
+use std::env;
 use std::fs::File;
 use std::io::{Read, Write};
 use std::path::{Path, PathBuf};
@@ -109,6 +110,9 @@ fn main() -> Result<()> {
         build_output_path,
         required_input,
         unit_vname,
+        env::current_dir()?
+            .to_str()
+            .ok_or_else(|| anyhow!("working directory is invalid UTF-8"))?,
     );
     let mut indexed_compilation_bytes: Vec<u8> = Vec::new();
     indexed_compilation
@@ -156,12 +160,14 @@ fn get_spawn_info(file_path: impl AsRef<Path>) -> Result<SpawnInfo> {
 /// * `build_output_path` - The output path of the build target
 /// * `required_input` - The generated data for the CompilationUnit
 ///   `required_input` field
+/// * `working_directory` - The working_directory field for the CompilationUnit
 fn create_indexed_compilation(
     source_files: Vec<String>,
     arguments: Vec<String>,
     build_output_path: &str,
     required_input: Vec<CompilationUnit_FileInput>,
     mut unit_vname: VName,
+    working_directory: &str,
 ) -> IndexedCompilation {
     let mut compilation_unit = CompilationUnit::new();
 
@@ -171,6 +177,7 @@ fn create_indexed_compilation(
     compilation_unit.set_argument(protobuf::RepeatedField::from_vec(arguments));
     compilation_unit.set_required_input(protobuf::RepeatedField::from_vec(required_input));
     compilation_unit.set_output_key(build_output_path.to_string());
+    compilation_unit.set_working_directory(working_directory.to_string());
 
     let mut indexed_compilation = IndexedCompilation::new();
     indexed_compilation.set_unit(compilation_unit);
diff --git a/kythe/rust/extractor/tests/integration_test.rs b/kythe/rust/extractor/tests/integration_test.rs
index 224eab74a1..f9b73d083b 100644
--- a/kythe/rust/extractor/tests/integration_test.rs
+++ b/kythe/rust/extractor/tests/integration_test.rs
@@ -193,6 +193,8 @@ fn correct_arguments_succeed(
     let required_inputs = compilation_unit.get_required_input().to_vec();
     let source_input = required_inputs.get(0).expect("Failed to get first required input");
 
+    assert!(!compilation_unit.get_working_directory().is_empty());
+
     let source_file_path = source_input.get_info().get_path();
     assert!(
         source_file_path.contains("kythe/rust/extractor/main.rs"),