Skip to content

Commit

Permalink
Add Record::drain to take out elements by range
Browse files Browse the repository at this point in the history
Matches the general behavior of `Vec::drain` or
`indexmap::IndexMap::drain`:
- Drop the remaining elements (implementing the unstable `keep_rest()`
would not be compatible with something like `indexmap`)
- No `AsRef<[T]>` or `Drain::as_slice()` behavior as this would make
layout assumptions.
- `Drain: DoubleEndedIterator`

Found useful in nushell#10903
  • Loading branch information
sholderbach committed Nov 8, 2023
1 parent 86cd387 commit 05723a8
Showing 1 changed file with 67 additions and 0 deletions.
67 changes: 67 additions & 0 deletions crates/nu-protocol/src/value/record.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
use std::ops::RangeBounds;

use crate::Value;

use serde::{Deserialize, Serialize};
Expand Down Expand Up @@ -218,6 +220,42 @@ impl Record {
iter: self.vals.into_iter(),
}
}

/// Obtain an iterator to remove elements in `range`
///
/// Elements not consumed from the iterator will be dropped
///
/// ```rust
/// use nu_protocol::{record, Value};
///
/// let mut rec = record!(
/// "a" => Value::test_nothing(),
/// "b" => Value::test_int(42),
/// "c" => Value::test_string("foo"),
/// );
/// {
/// let mut drainer = rec.drain(1..);
/// assert_eq!(drainer.next(), Some(("b".into(), Value::test_int(42))));
/// // Dropping the `Drain`
/// }
/// let mut rec_iter = rec.into_iter();
/// assert_eq!(rec_iter.next(), Some(("a".into(), Value::test_nothing())));
/// assert_eq!(rec_iter.next(), None);
/// ```
pub fn drain<R>(&mut self, range: R) -> Drain
where
R: RangeBounds<usize> + Clone,
{
assert_eq!(
self.cols.len(),
self.vals.len(),
"Length of cols and vals must be equal for sane `Record::drain`"
);
Drain {
keys: self.cols.drain(range.clone()),
values: self.vals.drain(range),
}
}
}

impl FromIterator<(String, Value)> for Record {
Expand Down Expand Up @@ -357,6 +395,35 @@ impl ExactSizeIterator for IntoValues {
}
}

pub struct Drain<'a> {
keys: std::vec::Drain<'a, String>,
values: std::vec::Drain<'a, Value>,
}

impl Iterator for Drain<'_> {
type Item = (String, Value);

fn next(&mut self) -> Option<Self::Item> {
Some((self.keys.next()?, self.values.next()?))
}

fn size_hint(&self) -> (usize, Option<usize>) {
self.keys.size_hint()
}
}

impl DoubleEndedIterator for Drain<'_> {
fn next_back(&mut self) -> Option<Self::Item> {
Some((self.keys.next_back()?, self.values.next_back()?))
}
}

impl ExactSizeIterator for Drain<'_> {
fn len(&self) -> usize {
self.keys.len()
}
}

#[macro_export]
macro_rules! record {
{$($col:expr => $val:expr),+ $(,)?} => {
Expand Down

0 comments on commit 05723a8

Please sign in to comment.