Auto merge of #46789 - Diggsey:command-env-capture, r=dtolnay

Capture `Command` environment at spawn

Fixes #28975

This tracks a set of changes to the environment and then replays them at spawn time.

3 files changed, 145 insertions, 0 deletions

diff --git a/src/libstd/sys_common/mod.rs b/src/libstd/sys_common/mod.rs
index 5c4d7b52754..b16299a1d63 100644
--- a/src/libstd/sys_common/mod.rs
+++ b/src/libstd/sys_common/mod.rs
@@ -44,6 +44,7 @@ pub mod thread_local;
 pub mod util;
 pub mod wtf8;
 pub mod bytestring;
+pub mod process;
 
 cfg_if! {
     if #[cfg(any(target_os = "redox", target_os = "l4re"))] {
diff --git a/src/libstd/sys_common/process.rs b/src/libstd/sys_common/process.rs
new file mode 100644
index 00000000000..fd1a5fdb410
--- /dev/null
+++ b/src/libstd/sys_common/process.rs
@@ -0,0 +1,124 @@
+// Copyright 2014 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#![allow(dead_code)]
+#![unstable(feature = "process_internals", issue = "0")]
+
+use ffi::{OsStr, OsString};
+use env;
+use collections::BTreeMap;
+use alloc::borrow::Borrow;
+
+pub trait EnvKey:
+    From<OsString> + Into<OsString> +
+    Borrow<OsStr> + Borrow<Self> + AsRef<OsStr> +
+    Ord + Clone {}
+
+// Implement a case-sensitive environment variable key
+#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
+pub struct DefaultEnvKey(OsString);
+
+impl From<OsString> for DefaultEnvKey {
+    fn from(k: OsString) -> Self { DefaultEnvKey(k) }
+}
+
+impl From<DefaultEnvKey> for OsString {
+    fn from(k: DefaultEnvKey) -> Self { k.0 }
+}
+
+impl Borrow<OsStr> for DefaultEnvKey {
+    fn borrow(&self) -> &OsStr { &self.0 }
+}
+
+impl AsRef<OsStr> for DefaultEnvKey {
+    fn as_ref(&self) -> &OsStr { &self.0 }
+}
+
+impl EnvKey for DefaultEnvKey {}
+
+// Stores a set of changes to an environment
+#[derive(Clone, Debug)]
+pub struct CommandEnv<K> {
+    clear: bool,
+    vars: BTreeMap<K, Option<OsString>>
+}
+
+impl<K: EnvKey> Default for CommandEnv<K> {
+    fn default() -> Self {
+        CommandEnv {
+            clear: false,
+            vars: Default::default()
+        }
+    }
+}
+
+impl<K: EnvKey> CommandEnv<K> {
+    // Capture the current environment with these changes applied
+    pub fn capture(&self) -> BTreeMap<K, OsString> {
+        let mut result = BTreeMap::<K, OsString>::new();
+        if !self.clear {
+            for (k, v) in env::vars_os() {
+                result.insert(k.into(), v);
+            }
+        }
+        for (k, maybe_v) in &self.vars {
+            if let &Some(ref v) = maybe_v {
+                result.insert(k.clone(), v.clone());
+            } else {
+                result.remove(k);
+            }
+        }
+        result
+    }
+
+    // Apply these changes directly to the current environment
+    pub fn apply(&self) {
+        if self.clear {
+            for (k, _) in env::vars_os() {
+                env::remove_var(k);
+            }
+        }
+        for (key, maybe_val) in self.vars.iter() {
+            if let &Some(ref val) = maybe_val {
+                env::set_var(key, val);
+            } else {
+                env::remove_var(key);
+            }
+        }
+    }
+
+    pub fn is_unchanged(&self) -> bool {
+        !self.clear && self.vars.is_empty()
+    }
+
+    pub fn capture_if_changed(&self) -> Option<BTreeMap<K, OsString>> {
+        if self.is_unchanged() {
+            None
+        } else {
+            Some(self.capture())
+        }
+    }
+
+    // The following functions build up changes
+    pub fn set(&mut self, key: &OsStr, value: &OsStr) {
+        self.vars.insert(key.to_owned().into(), Some(value.to_owned()));
+    }
+    pub fn remove(&mut self, key: &OsStr) {
+        if self.clear {
+            self.vars.remove(key);
+        } else {
+            self.vars.insert(key.to_owned().into(), None);
+        }
+    }
+    pub fn clear(&mut self) {
+        self.clear = true;
+        self.vars.clear();
+    }
+}
diff --git a/src/libstd/sys_common/wtf8.rs b/src/libstd/sys_common/wtf8.rs
index b2fc559bb37..46d554d6411 100644
--- a/src/libstd/sys_common/wtf8.rs
+++ b/src/libstd/sys_common/wtf8.rs
@@ -134,6 +134,12 @@ impl ops::Deref for Wtf8Buf {
     }
 }
 
+impl ops::DerefMut for Wtf8Buf {
+    fn deref_mut(&mut self) -> &mut Wtf8 {
+        self.as_mut_slice()
+    }
+}
+
 /// Format the string with double quotes,
 /// and surrogates as `\u` followed by four hexadecimal digits.
 /// Example: `"a\u{D800}"` for a string with code points [U+0061, U+D800]
@@ -221,6 +227,11 @@ impl Wtf8Buf {
         unsafe { Wtf8::from_bytes_unchecked(&self.bytes) }
     }
 
+    #[inline]
+    pub fn as_mut_slice(&mut self) -> &mut Wtf8 {
+        unsafe { Wtf8::from_mut_bytes_unchecked(&mut self.bytes) }
+    }
+
     /// Reserves capacity for at least `additional` more bytes to be inserted
     /// in the given `Wtf8Buf`.
     /// The collection may reserve more space to avoid frequent reallocations.
@@ -486,6 +497,15 @@ impl Wtf8 {
         mem::transmute(value)
     }
 
+    /// Creates a mutable WTF-8 slice from a mutable WTF-8 byte slice.
+    ///
+    /// Since the byte slice is not checked for valid WTF-8, this functions is
+    /// marked unsafe.
+    #[inline]
+    unsafe fn from_mut_bytes_unchecked(value: &mut [u8]) -> &mut Wtf8 {
+        mem::transmute(value)
+    }
+
     /// Returns the length, in WTF-8 bytes.
     #[inline]
     pub fn len(&self) -> usize {