Skip to content

Commit fc17d43

Browse files
author
blake2-ppc
committed
vec: Add .shift_opt() -> Option<T>
Add a function to safely retrieve the first element of a ~[T], as Option<T>. Implement shift() using shift_opt(). Add tests for both .shift() and .shift_opt()
1 parent d805b83 commit fc17d43

File tree

1 file changed

+40
-12
lines changed

1 file changed

+40
-12
lines changed

src/libstd/vec.rs

Lines changed: 40 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1093,6 +1093,7 @@ pub trait OwnedVector<T> {
10931093
fn pop(&mut self) -> T;
10941094
fn pop_opt(&mut self) -> Option<T>;
10951095
fn shift(&mut self) -> T;
1096+
fn shift_opt(&mut self) -> Option<T>;
10961097
fn unshift(&mut self, x: T);
10971098
fn insert(&mut self, i: uint, x:T);
10981099
fn remove(&mut self, i: uint) -> T;
@@ -1305,20 +1306,26 @@ impl<T> OwnedVector<T> for ~[T] {
13051306
}
13061307

13071308
/// Removes the first element from a vector and return it
1309+
#[inline]
13081310
fn shift(&mut self) -> T {
1309-
unsafe {
1310-
assert!(!self.is_empty());
1311-
1312-
if self.len() == 1 { return self.pop() }
1311+
self.shift_opt().expect("shift: empty vector")
1312+
}
13131313

1314-
if self.len() == 2 {
1315-
let last = self.pop();
1316-
let first = self.pop();
1317-
self.push(last);
1318-
return first;
1319-
}
1314+
/// Removes the first element from a vector and return it, or `None` if it is empty
1315+
fn shift_opt(&mut self) -> Option<T> {
1316+
unsafe {
1317+
let ln = match self.len() {
1318+
0 => return None,
1319+
1 => return self.pop_opt(),
1320+
2 => {
1321+
let last = self.pop();
1322+
let first = self.pop_opt();
1323+
self.push(last);
1324+
return first;
1325+
}
1326+
x => x
1327+
};
13201328

1321-
let ln = self.len();
13221329
let next_ln = self.len() - 1;
13231330

13241331
// Save the last element. We're going to overwrite its position
@@ -1354,7 +1361,7 @@ impl<T> OwnedVector<T> for ~[T] {
13541361
let vp = raw::to_mut_ptr(*self);
13551362
let vp = ptr::mut_offset(vp, next_ln - 1);
13561363

1357-
ptr::replace_ptr(vp, work_elt)
1364+
Some(ptr::replace_ptr(vp, work_elt))
13581365
}
13591366
}
13601367

@@ -2763,6 +2770,27 @@ mod tests {
27632770
assert_eq!([&[1], &[2], &[3]].connect_vec(&0), ~[1, 0, 2, 0, 3]);
27642771
}
27652772

2773+
#[test]
2774+
fn test_shift() {
2775+
let mut x = ~[1, 2, 3];
2776+
assert_eq!(x.shift(), 1);
2777+
assert_eq!(&x, &~[2, 3]);
2778+
assert_eq!(x.shift(), 2);
2779+
assert_eq!(x.shift(), 3);
2780+
assert_eq!(x.len(), 0);
2781+
}
2782+
2783+
#[test]
2784+
fn test_shift_opt() {
2785+
let mut x = ~[1, 2, 3];
2786+
assert_eq!(x.shift_opt(), Some(1));
2787+
assert_eq!(&x, &~[2, 3]);
2788+
assert_eq!(x.shift_opt(), Some(2));
2789+
assert_eq!(x.shift_opt(), Some(3));
2790+
assert_eq!(x.shift_opt(), None);
2791+
assert_eq!(x.len(), 0);
2792+
}
2793+
27662794
#[test]
27672795
fn test_unshift() {
27682796
let mut x = ~[1, 2, 3];

0 commit comments

Comments
 (0)