Environment variable iterator overrides

src/libstd/env.rs

@@ -142,13 +142,29 @@ impl Iterator for Vars {
         })
     }
     fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
+    fn count(self) -> usize { self.inner.count() }
+    fn nth(&mut self, n: usize) -> Option<(String, String)> {
+        self.inner.nth(n).map(|(a, b)| {
+            (a.into_string().unwrap(), b.into_string().unwrap())
+        })
+    }
+    fn last(self) -> Option<(String, String)> {
+        self.inner.last().map(|(a, b)| {
+            (a.into_string().unwrap(), b.into_string().unwrap())
+        })
+    }
 }
 
 #[stable(feature = "env", since = "1.0.0")]
 impl Iterator for VarsOs {
     type Item = (OsString, OsString);
     fn next(&mut self) -> Option<(OsString, OsString)> { self.inner.next() }
     fn size_hint(&self) -> (usize, Option<usize>) { self.inner.size_hint() }
+    fn count(self) -> usize { self.inner.count() }
+    fn nth(&mut self, n: usize) -> Option<(OsString, OsString)> {
+        self.inner.nth(n)
+    }
+    fn last(self) -> Option<(OsString, OsString)> { self.inner.last() }
 }
 
 /// Fetches the environment variable `key` from the current process.

src/libstd/sys/unix/os.rs

@@ -407,6 +407,11 @@ impl Iterator for Env {
     type Item = (OsString, OsString);
     fn next(&mut self) -> Option<(OsString, OsString)> { self.iter.next() }
     fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
+    fn count(self) -> usize { self.iter.count() }
+    fn nth(&mut self, n: usize) -> Option<(OsString, OsString)> {
+        self.iter.nth(n)
+    }
+    fn last(self) -> Option<(OsString, OsString)> { self.iter.last() }
 }
 
 #[cfg(target_os = "macos")]

src/libstd/sys/windows/os.rs

@@ -72,14 +72,23 @@ pub fn error_string(errnum: i32) -> String {
 }
 
 pub struct Env {
+    inner: EnvSlices,
+}
+
+// This is a sub-iterator for `Env` to avoid allocating `OsString`s when not
+// needed, by separating the iteration over the `LPWCH` from constructing the
+// result strings this allows methods on `Env` such as `nth` to avoid allocations.
+struct EnvSlices {
     base: c::LPWCH,
     cur: c::LPWCH,
 }
 
-impl Iterator for Env {
-    type Item = (OsString, OsString);
+impl Iterator for EnvSlices {
+    // These aren't really 'static, but that's required to avoid some unwanted
+    // lifetime definitions. The slices have the same lifetime as this iterator.
+    type Item = (&'static [u16], &'static [u16]);
 
-    fn next(&mut self) -> Option<(OsString, OsString)> {
+    fn next(&mut self) -> Option<(&'static [u16], &'static [u16])> {
         loop {
             unsafe {
                 if *self.cur == 0 { return None }
@@ -100,16 +109,37 @@ impl Iterator for Env {
                     Some(p) => p,
                     None => continue,
                 };
-                return Some((
-                    OsStringExt::from_wide(&s[..pos]),
-                    OsStringExt::from_wide(&s[pos+1..]),
-                ))
+                return Some((&s[..pos], &s[pos+1..]))
             }
         }
     }
 }
 
-impl Drop for Env {
+impl Env {
+    fn convert((a, b): (&'static [u16], &'static [u16])) -> (OsString, OsString) {
+        (OsStringExt::from_wide(a), OsStringExt::from_wide(b))
+    }
+}
+
+impl Iterator for Env {
+    type Item = (OsString, OsString);
+
+    fn next(&mut self) -> Option<(OsString, OsString)> {
+        self.inner.next().map(Self::convert)
+    }
+
+    fn count(self) -> usize { self.inner.count() }
+
+    fn nth(&mut self, n: usize) -> Option<(OsString, OsString)> {
+        self.inner.nth(n).map(Self::convert)
+    }
+
+    fn last(self) -> Option<(OsString, OsString)> {
+        self.inner.last().map(Self::convert)
+    }
+}
+
+impl Drop for EnvSlices {
     fn drop(&mut self) {
         unsafe { c::FreeEnvironmentStringsW(self.base); }
     }
@@ -122,7 +152,7 @@ pub fn env() -> Env {
             panic!("failure getting env string from OS: {}",
                    io::Error::last_os_error());
         }
-        Env { base: ch, cur: ch }
+        Env { inner: EnvSlices { base: ch, cur: ch } }
     }
 }
 

src/test/run-pass/env-vars-iter.rs

@@ -0,0 +1,68 @@
+// Copyright 2016 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.
+
+use std::process::Command;
+use std::env;
+
+fn check_var(&(ref key, ref val): &(String, String)) -> bool {
+    match &**key {
+        "FOO" => { assert_eq!(val, "BAR"); true },
+        "BAR" => { assert_eq!(val, "BAZ"); true },
+        "BAZ" => { assert_eq!(val, "FOO"); true },
+        _ => false,
+    }
+}
+
+fn main() {
+    if let Some(arg) = env::args().nth(1) {
+        match &*arg {
+            "empty" => {
+                assert_eq!(env::vars().count(), 0);
+                assert_eq!(env::vars().next(), None);
+                assert_eq!(env::vars().nth(1), None);
+                assert_eq!(env::vars().last(), None);
+            },
+            "many" => {
+                assert!(env::vars().count() >= 3);
+                assert!(env::vars().last().is_some());
+                assert_eq!(env::vars().filter(check_var).count(), 3);
+                assert_eq!(
+                    (0..env::vars().count())
+                        .map(|i| env::vars().nth(i).unwrap())
+                        .filter(check_var)
+                        .count(),
+                    3);
+            },
+            arg => {
+                panic!("Unexpected arg {}", arg);
+            },
+        }
+    } else {
+        // Command::env_clear does not work on Windows.
+        // https://github.com/rust-lang/rust/issues/31259
+        if !cfg!(windows) {
+            let status = Command::new(env::current_exe().unwrap())
+                                 .arg("empty")
+                                 .env_clear()
+                                 .status()
+                                 .unwrap();
+            assert_eq!(status.code(), Some(0));
+        }
+
+        let status = Command::new(env::current_exe().unwrap())
+                             .arg("many")
+                             .env("FOO", "BAR")
+                             .env("BAR", "BAZ")
+                             .env("BAZ", "FOO")
+                             .status()
+                             .unwrap();
+        assert_eq!(status.code(), Some(0));
+    }
+}