From b8219ac7f74d06c4983e96a1b9d0dbe00d7a52c7 Mon Sep 17 00:00:00 2001
From: Marcel Hlopko <hlopko@google.com>
Date: Wed, 2 Jun 2021 09:09:26 +0200
Subject: [PATCH] Add support for manifest based runfiles. (#760)

* Add support for manifest based runfiles.

These are typically used on Windows when using symlinks which are needed
for directory based runfiles is not possible.
---
 .bazelci/presubmit.yml     |  2 -
 tools/runfiles/runfiles.rs | 99 +++++++++++++++++++++++++++++++++-----
 2 files changed, 88 insertions(+), 13 deletions(-)

diff --git a/.bazelci/presubmit.yml b/.bazelci/presubmit.yml
index 9881deb325..478cfdf7d0 100644
--- a/.bazelci/presubmit.yml
+++ b/.bazelci/presubmit.yml
@@ -53,7 +53,6 @@ tasks:
       - "-//test/chained_direct_deps:mod3_doc_test"
       - "-@examples//fibonacci:fibonacci_doc_test"
       - "-@examples//hello_lib:hello_lib_doc_test"
-      - "-//tools/runfiles:runfiles_doc_test"
       - "-@examples//ffi/rust_calling_c/simple/..."
       # See https://github.com/bazelbuild/bazel/issues/9987
       - "-@examples//ffi/rust_calling_c:matrix_dylib_test"
@@ -71,7 +70,6 @@ tasks:
       - "-//test/test_env/..."
       - "-//test/proto/..."
       - "-//tools/rust_analyzer/..."
-      - "-//tools/runfiles/..."
       - "-//test/rustfmt/..."
       - "@examples//..."
       - "-@examples//ffi/rust_calling_c:matrix_dylib_test"
diff --git a/tools/runfiles/runfiles.rs b/tools/runfiles/runfiles.rs
index 8f3bf5d92b..6486c2e0e3 100644
--- a/tools/runfiles/runfiles.rs
+++ b/tools/runfiles/runfiles.rs
@@ -13,7 +13,7 @@
 //!     ```
 //!
 //! 2.  Import the runfiles library.
-//!     ```
+//!     ```ignore
 //!     extern crate runfiles;
 //!
 //!     use runfiles::Runfiles;
@@ -31,23 +31,55 @@
 //!     // ...
 //!     ```
 
+use std::collections::HashMap;
 use std::env;
+use std::ffi::OsString;
 use std::fs;
 use std::io;
 use std::path::Path;
 use std::path::PathBuf;
 
+enum Mode {
+    DirectoryBased(PathBuf),
+    ManifestBased(HashMap<PathBuf, PathBuf>),
+}
+
 pub struct Runfiles {
-    runfiles_dir: PathBuf,
+    mode: Mode,
 }
 
 impl Runfiles {
-    /// Creates a directory based Runfiles object.
-    ///
-    /// Manifest based creation is not currently supported.
+    /// Creates a manifest based Runfiles object when
+    /// RUNFILES_MANIFEST_ONLY environment variable is present,
+    /// or a directory based Runfiles object otherwise.
     pub fn create() -> io::Result<Self> {
+        if is_manifest_only() {
+            Self::create_manifest_based()
+        } else {
+            Self::create_directory_based()
+        }
+    }
+
+    fn create_directory_based() -> io::Result<Self> {
+        Ok(Runfiles {
+            mode: Mode::DirectoryBased(find_runfiles_dir()?),
+        })
+    }
+
+    fn create_manifest_based() -> io::Result<Self> {
+        let manifest_path = find_manifest_path()?;
+        let manifest_content = std::fs::read_to_string(manifest_path)?;
+        let path_mapping = manifest_content
+            .lines()
+            .map(|line| {
+                let pair = line
+                    .split_once(" ")
+                    .expect("manifest file contained unexpected content");
+                (pair.0.into(), pair.1.into())
+            })
+            .collect::<HashMap<_, _>>();
         Ok(Runfiles {
-            runfiles_dir: find_runfiles_dir()?,
+            mode: Mode::ManifestBased(path_mapping),
         })
     }
 
@@ -61,12 +93,25 @@ impl Runfiles {
         if path.is_absolute() {
             return path.to_path_buf();
         }
-        self.runfiles_dir.join(path)
+        match &self.mode {
+            Mode::DirectoryBased(runfiles_dir) => runfiles_dir.join(path),
+            Mode::ManifestBased(path_mapping) => path_mapping
+                .get(path)
+                .expect(&format!(
+                    "Path {} not found among runfiles.",
+                    path.to_string_lossy()
+                ))
+                .clone(),
+        }
     }
 }
 
 /// Returns the .runfiles directory for the currently executing binary.
 pub fn find_runfiles_dir() -> io::Result<PathBuf> {
+    assert_ne!(
+        std::env::var_os("RUNFILES_MANIFEST_ONLY").unwrap_or(OsString::from("0")),
+        "1"
+    );
     let exec_path = std::env::args().nth(0).expect("arg 0 was not set");
 
     let mut binary_path = PathBuf::from(&exec_path);
@@ -108,10 +153,31 @@ pub fn find_runfiles_dir() -> io::Result<PathBuf> {
         }
     }
 
-    Err(io::Error::new(
-        io::ErrorKind::Other,
-        "Failed to find .runfiles directory.",
-    ))
+    Err(make_io_error("failed to find .runfiles directory"))
+}
+
+fn make_io_error(msg: &str) -> io::Error {
+    io::Error::new(io::ErrorKind::Other, msg)
+}
+
+fn is_manifest_only() -> bool {
+    match std::env::var("RUNFILES_MANIFEST_ONLY") {
+        Ok(val) => val == "1",
+        Err(_) => false,
+    }
+}
+
+fn find_manifest_path() -> io::Result<PathBuf> {
+    assert_eq!(
+        std::env::var_os("RUNFILES_MANIFEST_ONLY").expect("RUNFILES_MANIFEST_ONLY was not set"),
+        OsString::from("1")
+    );
+    match std::env::var_os("RUNFILES_MANIFEST_FILE") {
+        Some(path) => Ok(path.into()),
+        None => Err(
+            make_io_error(
+                "RUNFILES_MANIFEST_ONLY was set to '1', but RUNFILES_MANIFEST_FILE was not set. Did Bazel change?"))
+    }
 }
 
 #[cfg(test)]
@@ -132,4 +198,15 @@ mod test {
 
         assert_eq!("Example Text!", buffer);
     }
+
+    #[test]
+    fn test_manifest_based_can_read_data_from_runfiles() {
+        let mut path_mapping = HashMap::new();
+        path_mapping.insert("a/b".into(), "c/d".into());
+        let r = Runfiles {
+            mode: Mode::ManifestBased(path_mapping),
+        };
+
+        assert_eq!(r.rlocation("a/b"), PathBuf::from("c/d"));
+    }
 }