auto merge of #15353 : aturon/rust/env-hashmap, r=alexcrichton

This commit adds `env_insert` and `env_remove` methods to the `Command`
builder, easing updates to the environment variables for the child
process. The existing method, `env`, is still available for overriding
the entire environment in one shot (after which the `env_insert` and
`env_remove` methods can be used to make further adjustments).

To support these new methods, the internal `env` representation for
`Command` has been changed to an optional `HashMap` holding owned
`CString`s (to support non-utf8 data). The `HashMap` is only
materialized if the environment is updated. The implementation does not
try hard to avoid allocation, since the cost of launching a process will
dwarf any allocation cost.

This patch also adds `PartialOrd`, `Eq`, and `Hash` implementations for
`CString`.

Author: bors

src/compiletest/procsrv.rs

@@ -8,12 +8,11 @@
 // option. This file may not be copied, modified, or distributed
 // except according to those terms.
 
-use std::os;
 use std::str;
 use std::io::process::{ProcessExit, Command, Process, ProcessOutput};
 use std::dynamic_lib::DynamicLibrary;
 
-fn target_env(lib_path: &str, aux_path: Option<&str>) -> Vec<(String, String)> {
+fn add_target_env(cmd: &mut Command, lib_path: &str, aux_path: Option<&str>) {
     // Need to be sure to put both the lib_path and the aux path in the dylib
     // search path for the child.
     let mut path = DynamicLibrary::search_path();
@@ -23,19 +22,11 @@ fn target_env(lib_path: &str, aux_path: Option<&str>) -> Vec<(String, String)> {
     }
     path.insert(0, Path::new(lib_path));
 
-    // Remove the previous dylib search path var
-    let var = DynamicLibrary::envvar();
-    let mut env: Vec<(String,String)> = os::env();
-    match env.iter().position(|&(ref k, _)| k.as_slice() == var) {
-        Some(i) => { env.remove(i); }
-        None => {}
-    }
-
     // Add the new dylib search path var
+    let var = DynamicLibrary::envvar();
     let newpath = DynamicLibrary::create_path(path.as_slice());
     let newpath = str::from_utf8(newpath.as_slice()).unwrap().to_string();
-    env.push((var.to_string(), newpath));
-    return env;
+    cmd.env(var.to_string(), newpath);
 }
 
 pub struct Result {pub status: ProcessExit, pub out: String, pub err: String}
@@ -47,8 +38,14 @@ pub fn run(lib_path: &str,
            env: Vec<(String, String)> ,
            input: Option<String>) -> Option<Result> {
 
-    let env = env.clone().append(target_env(lib_path, aux_path).as_slice());
-    match Command::new(prog).args(args).env(env.as_slice()).spawn() {
+    let mut cmd = Command::new(prog);
+    cmd.args(args);
+    add_target_env(&mut cmd, lib_path, aux_path);
+    for (key, val) in env.move_iter() {
+        cmd.env(key, val);
+    }
+
+    match cmd.spawn() {
         Ok(mut process) => {
             for input in input.iter() {
                 process.stdin.get_mut_ref().write(input.as_bytes()).unwrap();
@@ -73,8 +70,14 @@ pub fn run_background(lib_path: &str,
            env: Vec<(String, String)> ,
            input: Option<String>) -> Option<Process> {
 
-    let env = env.clone().append(target_env(lib_path, aux_path).as_slice());
-    match Command::new(prog).args(args).env(env.as_slice()).spawn() {
+    let mut cmd = Command::new(prog);
+    cmd.args(args);
+    add_target_env(&mut cmd, lib_path, aux_path);
+    for (key, val) in env.move_iter() {
+        cmd.env(key, val);
+    }
+
+    match cmd.spawn() {
         Ok(mut process) => {
             for input in input.iter() {
                 process.stdin.get_mut_ref().write(input.as_bytes()).unwrap();

src/compiletest/runtest.rs

@@ -574,7 +574,7 @@ fn run_debuginfo_lldb_test(config: &Config, props: &TestProps, testfile: &Path)
         cmd.arg("./src/etc/lldb_batchmode.py")
            .arg(test_executable)
            .arg(debugger_script)
-           .env([("PYTHONPATH", config.lldb_python_dir.clone().unwrap().as_slice())]);
+           .env_set_all([("PYTHONPATH", config.lldb_python_dir.clone().unwrap().as_slice())]);
 
         let (status, out, err) = match cmd.spawn() {
             Ok(process) => {

src/libnative/io/process.rs

@@ -729,7 +729,7 @@ fn with_argv<T>(prog: &CString, args: &[CString],
 }
 
 #[cfg(unix)]
-fn with_envp<T>(env: Option<&[(CString, CString)]>,
+fn with_envp<T>(env: Option<&[(&CString, &CString)]>,
                 cb: proc(*const c_void) -> T) -> T {
     // On posixy systems we can pass a char** for envp, which is a
     // null-terminated array of "k=v\0" strings. Since we must create
@@ -762,7 +762,7 @@ fn with_envp<T>(env: Option<&[(CString, CString)]>,
 }
 
 #[cfg(windows)]
-fn with_envp<T>(env: Option<&[(CString, CString)]>, cb: |*mut c_void| -> T) -> T {
+fn with_envp<T>(env: Option<&[(&CString, &CString)]>, cb: |*mut c_void| -> T) -> T {
     // On win32 we pass an "environment block" which is not a char**, but
     // rather a concatenation of null-terminated k=v\0 sequences, with a final
     // \0 to terminate.

src/librustdoc/test.rs

@@ -176,26 +176,15 @@ fn runtest(test: &str, cratename: &str, libs: HashSet<Path>, should_fail: bool,
     // environment to ensure that the target loads the right libraries at
     // runtime. It would be a sad day if the *host* libraries were loaded as a
     // mistake.
-    let exe = outdir.path().join("rust_out");
-    let env = {
+    let mut cmd = Command::new(outdir.path().join("rust_out"));
+    let newpath = {
         let mut path = DynamicLibrary::search_path();
         path.insert(0, libdir.clone());
-
-        // Remove the previous dylib search path var
-        let var = DynamicLibrary::envvar();
-        let mut env: Vec<(String,String)> = os::env().move_iter().collect();
-        match env.iter().position(|&(ref k, _)| k.as_slice() == var) {
-            Some(i) => { env.remove(i); }
-            None => {}
-        };
-
-        // Add the new dylib search path var
-        let newpath = DynamicLibrary::create_path(path.as_slice());
-        env.push((var.to_string(),
-                  str::from_utf8(newpath.as_slice()).unwrap().to_string()));
-        env
+        DynamicLibrary::create_path(path.as_slice())
     };
-    match Command::new(exe).env(env.as_slice()).output() {
+    cmd.env(DynamicLibrary::envvar(), newpath.as_slice());
+
+    match cmd.output() {
         Err(e) => fail!("couldn't run the test: {}{}", e,
                         if e.kind == io::PermissionDenied {
                             " - maybe your tempdir is mounted with noexec?"

src/librustrt/c_str.rs

@@ -69,6 +69,7 @@ use core::prelude::*;
 
 use alloc::libc_heap::malloc_raw;
 use collections::string::String;
+use collections::hash;
 use core::kinds::marker;
 use core::mem;
 use core::ptr;
@@ -116,6 +117,22 @@ impl PartialEq for CString {
     }
 }
 
+impl PartialOrd for CString {
+    #[inline]
+    fn partial_cmp(&self, other: &CString) -> Option<Ordering> {
+        self.as_bytes().partial_cmp(&other.as_bytes())
+    }
+}
+
+impl Eq for CString {}
+
+impl<S: hash::Writer> hash::Hash<S> for CString {
+    #[inline]
+    fn hash(&self, state: &mut S) {
+        self.as_bytes().hash(state)
+    }
+}
+
 impl CString {
     /// Create a C String from a pointer.
     pub unsafe fn new(buf: *const libc::c_char, owns_buffer: bool) -> CString {

src/librustrt/lib.rs

@@ -17,7 +17,7 @@
        html_root_url = "http://doc.rust-lang.org/0.11.0/")]
 
 #![feature(macro_rules, phase, globs, thread_local, managed_boxes, asm)]
-#![feature(linkage, lang_items, unsafe_destructor)]
+#![feature(linkage, lang_items, unsafe_destructor, default_type_params)]
 #![no_std]
 #![experimental]
 

src/librustrt/rtio.rs

@@ -75,7 +75,7 @@ pub struct ProcessConfig<'a> {
 
     /// Optional environment to specify for the program. If this is None, then
     /// it will inherit the current process's environment.
-    pub env: Option<&'a [(CString, CString)]>,
+    pub env: Option<&'a [(&'a CString, &'a CString)]>,
 
     /// Optional working directory for the new process. If this is None, then
     /// the current directory of the running process is inherited.

src/librustuv/process.rs

@@ -193,7 +193,7 @@ fn with_argv<T>(prog: &CString, args: &[CString],
 }
 
 /// Converts the environment to the env array expected by libuv
-fn with_env<T>(env: Option<&[(CString, CString)]>,
+fn with_env<T>(env: Option<&[(&CString, &CString)]>,
                cb: |*const *const libc::c_char| -> T) -> T {
     // We can pass a char** for envp, which is a null-terminated array
     // of "k=v\0" strings. Since we must create these strings locally,

src/libstd/io/process.rs

@@ -16,6 +16,7 @@ use prelude::*;
 
 use str;
 use fmt;
+use os;
 use io::{IoResult, IoError};
 use io;
 use libc;
@@ -24,6 +25,7 @@ use owned::Box;
 use rt::rtio::{RtioProcess, ProcessConfig, IoFactory, LocalIo};
 use rt::rtio;
 use c_str::CString;
+use collections::HashMap;
 
 /// Signal a process to exit, without forcibly killing it. Corresponds to
 /// SIGTERM on unix platforms.
@@ -78,6 +80,9 @@ pub struct Process {
     pub extra_io: Vec<Option<io::PipeStream>>,
 }
 
+/// A HashMap representation of environment variables.
+pub type EnvMap = HashMap<CString, CString>;
+
 /// The `Command` type acts as a process builder, providing fine-grained control
 /// over how a new process should be spawned. A default configuration can be
 /// generated using `Command::new(program)`, where `program` gives a path to the
@@ -100,7 +105,7 @@ pub struct Command {
     // methods below, and serialized into rt::rtio::ProcessConfig.
     program: CString,
     args: Vec<CString>,
-    env: Option<Vec<(CString, CString)>>,
+    env: Option<EnvMap>,
     cwd: Option<CString>,
     stdin: StdioContainer,
     stdout: StdioContainer,
@@ -147,31 +152,53 @@ impl Command {
     }
 
     /// Add an argument to pass to the program.
-    pub fn arg<'a, T:ToCStr>(&'a mut self, arg: T) -> &'a mut Command {
+    pub fn arg<'a, T: ToCStr>(&'a mut self, arg: T) -> &'a mut Command {
         self.args.push(arg.to_c_str());
         self
     }
 
     /// Add multiple arguments to pass to the program.
-    pub fn args<'a, T:ToCStr>(&'a mut self, args: &[T]) -> &'a mut Command {
+    pub fn args<'a, T: ToCStr>(&'a mut self, args: &[T]) -> &'a mut Command {
         self.args.extend(args.iter().map(|arg| arg.to_c_str()));;
         self
     }
+    // Get a mutable borrow of the environment variable map for this `Command`.
+    fn get_env_map<'a>(&'a mut self) -> &'a mut EnvMap {
+        match self.env {
+            Some(ref mut map) => map,
+            None => {
+                // if the env is currently just inheriting from the parent's,
+                // materialize the parent's env into a hashtable.
+                self.env = Some(os::env_as_bytes().move_iter()
+                                   .map(|(k, v)| (k.as_slice().to_c_str(),
+                                                  v.as_slice().to_c_str()))
+                                   .collect());
+                self.env.as_mut().unwrap()
+            }
+        }
+    }
 
-    /// Sets the environment for the child process (rather than inheriting it
-    /// from the current process).
-
-    // FIXME (#13851): We should change this interface to allow clients to (1)
-    // build up the env vector incrementally and (2) allow both inheriting the
-    // current process's environment AND overriding/adding additional
-    // environment variables. The underlying syscalls assume that the
-    // environment has no duplicate names, so we really want to use a hashtable
-    // to compute the environment to pass down to the syscall; resolving issue
-    // #13851 will make it possible to use the standard hashtable.
-    pub fn env<'a, T:ToCStr>(&'a mut self, env: &[(T,T)]) -> &'a mut Command {
-        self.env = Some(env.iter().map(|&(ref name, ref val)| {
-            (name.to_c_str(), val.to_c_str())
-        }).collect());
+    /// Inserts or updates an environment variable mapping.
+    pub fn env<'a, T: ToCStr, U: ToCStr>(&'a mut self, key: T, val: U)
+                                         -> &'a mut Command {
+        self.get_env_map().insert(key.to_c_str(), val.to_c_str());
+        self
+    }
+
+    /// Removes an environment variable mapping.
+    pub fn env_remove<'a, T: ToCStr>(&'a mut self, key: T) -> &'a mut Command {
+        self.get_env_map().remove(&key.to_c_str());
+        self
+    }
+
+    /// Sets the entire environment map for the child process.
+    ///
+    /// If the given slice contains multiple instances of an environment
+    /// variable, the *rightmost* instance will determine the value.
+    pub fn env_set_all<'a, T: ToCStr, U: ToCStr>(&'a mut self, env: &[(T,U)])
+                                                 -> &'a mut Command {
+        self.env = Some(env.iter().map(|&(ref k, ref v)| (k.to_c_str(), v.to_c_str()))
+                                  .collect());
         self
     }
 
@@ -245,10 +272,15 @@ impl Command {
         let extra_io: Vec<rtio::StdioContainer> =
             self.extra_io.iter().map(|x| to_rtio(*x)).collect();
         LocalIo::maybe_raise(|io| {
+            let env = match self.env {
+                None => None,
+                Some(ref env_map) =>
+                    Some(env_map.iter().collect::<Vec<_>>())
+            };
             let cfg = ProcessConfig {
                 program: &self.program,
                 args: self.args.as_slice(),
-                env: self.env.as_ref().map(|env| env.as_slice()),
+                env: env.as_ref().map(|e| e.as_slice()),
                 cwd: self.cwd.as_ref(),
                 stdin: to_rtio(self.stdin),
                 stdout: to_rtio(self.stdout),
@@ -872,16 +904,50 @@ mod tests {
         }
     })
 
-    iotest!(fn test_add_to_env() {
+    iotest!(fn test_override_env() {
         let new_env = vec![("RUN_TEST_NEW_ENV", "123")];
-        let prog = env_cmd().env(new_env.as_slice()).spawn().unwrap();
+        let prog = env_cmd().env_set_all(new_env.as_slice()).spawn().unwrap();
         let result = prog.wait_with_output().unwrap();
         let output = str::from_utf8_lossy(result.output.as_slice()).into_string();
 
         assert!(output.as_slice().contains("RUN_TEST_NEW_ENV=123"),
                 "didn't find RUN_TEST_NEW_ENV inside of:\n\n{}", output);
     })
 
+    iotest!(fn test_add_to_env() {
+        let prog = env_cmd().env("RUN_TEST_NEW_ENV", "123").spawn().unwrap();
+        let result = prog.wait_with_output().unwrap();
+        let output = str::from_utf8_lossy(result.output.as_slice()).into_string();
+
+        assert!(output.as_slice().contains("RUN_TEST_NEW_ENV=123"),
+                "didn't find RUN_TEST_NEW_ENV inside of:\n\n{}", output);
+    })
+
+    iotest!(fn test_remove_from_env() {
+        use os;
+
+        // save original environment
+        let old_env = os::getenv("RUN_TEST_NEW_ENV");
+
+        os::setenv("RUN_TEST_NEW_ENV", "123");
+        let prog = env_cmd().env_remove("RUN_TEST_NEW_ENV").spawn().unwrap();
+        let result = prog.wait_with_output().unwrap();
+        let output = str::from_utf8_lossy(result.output.as_slice()).into_string();
+
+        // restore original environment
+        match old_env {
+            None => {
+                os::unsetenv("RUN_TEST_NEW_ENV");
+            }
+            Some(val) => {
+                os::setenv("RUN_TEST_NEW_ENV", val.as_slice());
+            }
+        }
+
+        assert!(!output.as_slice().contains("RUN_TEST_NEW_ENV"),
+                "found RUN_TEST_NEW_ENV inside of:\n\n{}", output);
+    })
+
     #[cfg(unix)]
     pub fn sleeper() -> Process {
         Command::new("sleep").arg("1000").spawn().unwrap()

src/test/run-pass/backtrace.rs

@@ -37,19 +37,8 @@ fn double() {
 }
 
 fn runtest(me: &str) {
-    let mut env = os::env().move_iter()
-                           .map(|(ref k, ref v)| {
-                               (k.to_string(), v.to_string())
-                           }).collect::<Vec<(String,String)>>();
-    match env.iter()
-             .position(|&(ref s, _)| "RUST_BACKTRACE" == s.as_slice()) {
-        Some(i) => { env.remove(i); }
-        None => {}
-    }
-    env.push(("RUST_BACKTRACE".to_string(), "1".to_string()));
-
     // Make sure that the stack trace is printed
-    let mut p = Command::new(me).arg("fail").env(env.as_slice()).spawn().unwrap();
+    let mut p = Command::new(me).arg("fail").env("RUST_BACKTRACE", "1").spawn().unwrap();
     let out = p.wait_with_output().unwrap();
     assert!(!out.status.success());
     let s = str::from_utf8(out.error.as_slice()).unwrap();
@@ -73,7 +62,8 @@ fn runtest(me: &str) {
             "bad output3: {}", s);
 
     // Make sure a stack trace isn't printed too many times
-    let mut p = Command::new(me).arg("double-fail").env(env.as_slice()).spawn().unwrap();
+    let mut p = Command::new(me).arg("double-fail")
+                                .env("RUST_BACKTRACE", "1").spawn().unwrap();
     let out = p.wait_with_output().unwrap();
     assert!(!out.status.success());
     let s = str::from_utf8(out.error.as_slice()).unwrap();