From 9d3fa18ec8eb8614efbeb72834940b7b012897b5 Mon Sep 17 00:00:00 2001 From: coolreader18 <33094578+coolreader18@users.noreply.github.com> Date: Sat, 6 Jul 2019 23:05:27 -0500 Subject: [PATCH] Add __mul__ to deque --- tests/snippets/test_collections.py | 4 ++ vm/src/obj/objlist.rs | 8 ++- vm/src/obj/objsequence.rs | 89 ++++++++++++++++++++---------- vm/src/obj/objtuple.rs | 7 ++- vm/src/stdlib/collections.rs | 28 ++++++++-- 5 files changed, 96 insertions(+), 40 deletions(-) diff --git a/tests/snippets/test_collections.py b/tests/snippets/test_collections.py index 4dd771a0c..820123c43 100644 --- a/tests/snippets/test_collections.py +++ b/tests/snippets/test_collections.py @@ -27,3 +27,7 @@ d.clear() assert d == deque() assert d == deque([], 4) + +assert deque([1, 2, 3]) * 2 == deque([1, 2, 3, 1, 2, 3]) + +assert deque([1, 2, 3], 4) * 2 == deque([3, 1, 2, 3]) diff --git a/vm/src/obj/objlist.rs b/vm/src/obj/objlist.rs index 92991bb21..b4d3a3f0f 100644 --- a/vm/src/obj/objlist.rs +++ b/vm/src/obj/objlist.rs @@ -392,12 +392,16 @@ impl PyListRef { } fn mul(self, counter: isize, vm: &VirtualMachine) -> PyObjectRef { - let new_elements = seq_mul(&self.elements.borrow(), counter); + let new_elements = seq_mul(&self.elements.borrow().as_slice(), counter) + .cloned() + .collect(); vm.ctx.new_list(new_elements) } fn imul(self, counter: isize, _vm: &VirtualMachine) -> Self { - let new_elements = seq_mul(&self.elements.borrow(), counter); + let new_elements = seq_mul(&self.elements.borrow().as_slice(), counter) + .cloned() + .collect(); self.elements.replace(new_elements); self } diff --git a/vm/src/obj/objsequence.rs b/vm/src/obj/objsequence.rs index a141ca80f..b360a5d51 100644 --- a/vm/src/obj/objsequence.rs +++ b/vm/src/obj/objsequence.rs @@ -216,35 +216,37 @@ pub fn get_item( } } -pub trait SimpleSeq<'a> { +type DynPyIter<'a> = Box + 'a>; + +pub trait SimpleSeq { fn len(&self) -> usize; - fn iter(&'a self) -> Box + 'a>; + fn iter(&self) -> DynPyIter; } -impl<'a> SimpleSeq<'a> for &'a [PyObjectRef] { +impl SimpleSeq for &[PyObjectRef] { fn len(&self) -> usize { (&**self).len() } - fn iter(&'a self) -> Box + 'a> { + fn iter(&self) -> DynPyIter { Box::new((&**self).iter()) } } -impl<'a> SimpleSeq<'a> for std::collections::VecDeque { +impl SimpleSeq for std::collections::VecDeque { fn len(&self) -> usize { self.len() } - fn iter(&'a self) -> Box + 'a> { + fn iter(&self) -> DynPyIter { Box::new(self.iter()) } } // impl<'a, I> -pub fn seq_equal<'a>( +pub fn seq_equal( vm: &VirtualMachine, - zelf: &'a dyn SimpleSeq<'a>, - other: &'a dyn SimpleSeq<'a>, + zelf: &dyn SimpleSeq, + other: &dyn SimpleSeq, ) -> Result { if zelf.len() == other.len() { for (a, b) in Iterator::zip(zelf.iter(), other.iter()) { @@ -262,10 +264,10 @@ pub fn seq_equal<'a>( } } -pub fn seq_lt<'a>( +pub fn seq_lt( vm: &VirtualMachine, - zelf: &'a dyn SimpleSeq<'a>, - other: &'a dyn SimpleSeq<'a>, + zelf: &dyn SimpleSeq, + other: &dyn SimpleSeq, ) -> Result { if zelf.len() == other.len() { for (a, b) in Iterator::zip(zelf.iter(), other.iter()) { @@ -302,10 +304,10 @@ pub fn seq_lt<'a>( } } -pub fn seq_gt<'a>( +pub fn seq_gt( vm: &VirtualMachine, - zelf: &'a dyn SimpleSeq<'a>, - other: &'a dyn SimpleSeq<'a>, + zelf: &dyn SimpleSeq, + other: &dyn SimpleSeq, ) -> Result { if zelf.len() == other.len() { for (a, b) in Iterator::zip(zelf.iter(), other.iter()) { @@ -341,32 +343,61 @@ pub fn seq_gt<'a>( } } -pub fn seq_ge<'a>( +pub fn seq_ge( vm: &VirtualMachine, - zelf: &'a dyn SimpleSeq<'a>, - other: &'a dyn SimpleSeq<'a>, + zelf: &dyn SimpleSeq, + other: &dyn SimpleSeq, ) -> Result { Ok(seq_gt(vm, zelf, other)? || seq_equal(vm, zelf, other)?) } -pub fn seq_le<'a>( +pub fn seq_le( vm: &VirtualMachine, - zelf: &'a dyn SimpleSeq<'a>, - other: &'a dyn SimpleSeq<'a>, + zelf: &dyn SimpleSeq, + other: &dyn SimpleSeq, ) -> Result { Ok(seq_lt(vm, zelf, other)? || seq_equal(vm, zelf, other)?) } -pub fn seq_mul(elements: &[PyObjectRef], counter: isize) -> Vec { - let current_len = elements.len(); - let new_len = counter.max(0) as usize * current_len; - let mut new_elements = Vec::with_capacity(new_len); - - for _ in 0..counter { - new_elements.extend(elements.to_owned()); +pub struct SeqMul<'a> { + seq: &'a dyn SimpleSeq, + repetitions: usize, + iter: DynPyIter<'a>, +} +impl SeqMul<'_> { + fn next_iter(&mut self) { + self.repetitions -= 1; + self.iter = self.seq.iter(); } +} +impl<'a> Iterator for SeqMul<'a> { + type Item = &'a PyObjectRef; + fn next(&mut self) -> Option { + match self.iter.next() { + Some(item) => Some(item), + None => { + if self.repetitions == 0 { + None + } else { + self.next_iter(); + self.next() + } + } + } + } + fn size_hint(&self) -> (usize, Option) { + let size = self.iter.len() + (self.repetitions * self.seq.len()); + (size, Some(size)) + } +} +impl ExactSizeIterator for SeqMul<'_> {} - new_elements +pub fn seq_mul(seq: &dyn SimpleSeq, repetitions: isize) -> SeqMul { + SeqMul { + seq, + repetitions: repetitions.max(0) as usize - 1, + iter: seq.iter(), + } } pub fn get_elements_cell<'a>(obj: &'a PyObjectRef) -> &'a RefCell> { diff --git a/vm/src/obj/objtuple.rs b/vm/src/obj/objtuple.rs index a4773e4e6..e13d3160f 100644 --- a/vm/src/obj/objtuple.rs +++ b/vm/src/obj/objtuple.rs @@ -164,13 +164,14 @@ impl PyTupleRef { } fn mul(self, counter: isize, vm: &VirtualMachine) -> PyObjectRef { - let new_elements = seq_mul(&self.elements, counter); + let new_elements = seq_mul(&self.elements.as_slice(), counter) + .cloned() + .collect(); vm.ctx.new_tuple(new_elements) } fn rmul(self, counter: isize, vm: &VirtualMachine) -> PyObjectRef { - let new_elements = seq_mul(&self.elements, counter); - vm.ctx.new_tuple(new_elements) + self.mul(counter, vm) } fn getitem(self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult { diff --git a/vm/src/stdlib/collections.rs b/vm/src/stdlib/collections.rs index 262aea94f..587d5ba1c 100644 --- a/vm/src/stdlib/collections.rs +++ b/vm/src/stdlib/collections.rs @@ -45,8 +45,8 @@ impl PyDeque { fn append(&self, obj: PyObjectRef, _vm: &VirtualMachine) { let mut deque = self.deque.borrow_mut(); if self.maxlen.get() == Some(deque.len()) { - deque.pop_front(); - } + deque.pop_front(); + } deque.push_back(obj); } @@ -54,8 +54,8 @@ impl PyDeque { fn appendleft(&self, obj: PyObjectRef, _vm: &VirtualMachine) { let mut deque = self.deque.borrow_mut(); if self.maxlen.get() == Some(deque.len()) { - deque.pop_back(); - } + deque.pop_back(); + } deque.push_front(obj); } @@ -125,8 +125,8 @@ impl PyDeque { let mut deque = self.deque.borrow_mut(); if self.maxlen.get() == Some(deque.len()) { - return Err(vm.new_index_error("deque already at its maximum size".to_string())); - } + return Err(vm.new_index_error("deque already at its maximum size".to_string())); + } let idx = if idx < 0 { if -idx as usize > deque.len() { @@ -321,6 +321,22 @@ impl PyDeque { let eq = objsequence::seq_ge(vm, lhs, rhs)?; Ok(vm.new_bool(eq)) } + + #[pymethod(name = "__mul__")] + fn mul(&self, n: isize, _vm: &VirtualMachine) -> Self { + let deque: &VecDeque<_> = &self.deque.borrow(); + let mul = objsequence::seq_mul(deque, n); + let skipped = if let Some(maxlen) = self.maxlen.get() { + mul.len() - maxlen + } else { + 0 + }; + let deque = mul.skip(skipped).cloned().collect(); + PyDeque { + deque: RefCell::new(deque), + maxlen: self.maxlen.clone(), + } + } } pub fn make_module(vm: &VirtualMachine) -> PyObjectRef {