diff --git a/src/libstd/env.rs b/src/libstd/env.rs index f7e13a35e9e50..895733d761499 100644 --- a/src/libstd/env.rs +++ b/src/libstd/env.rs @@ -142,6 +142,17 @@ impl Iterator for Vars { }) } fn size_hint(&self) -> (usize, Option) { 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")] @@ -149,6 +160,11 @@ impl Iterator for VarsOs { type Item = (OsString, OsString); fn next(&mut self) -> Option<(OsString, OsString)> { self.inner.next() } fn size_hint(&self) -> (usize, Option) { 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. diff --git a/src/libstd/sys/unix/os.rs b/src/libstd/sys/unix/os.rs index b6a0bd844094b..6052cbb381bd2 100644 --- a/src/libstd/sys/unix/os.rs +++ b/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) { 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")] diff --git a/src/libstd/sys/windows/os.rs b/src/libstd/sys/windows/os.rs index 688475a756574..0cfb02c0b57bf 100644 --- a/src/libstd/sys/windows/os.rs +++ b/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 } } } } diff --git a/src/test/run-pass/env-vars-iter.rs b/src/test/run-pass/env-vars-iter.rs new file mode 100644 index 0000000000000..7afd052d5af9f --- /dev/null +++ b/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 or the MIT license +// , 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)); + } +}