|
1 | 1 | use crate::convert::TryFrom; |
2 | 2 | use crate::{ |
3 | 3 | intrinsics, |
4 | | - iter::{from_fn, TrustedLen}, |
5 | | - ops::{Range, Try}, |
| 4 | + iter::{from_fn, range::RangeIter, TrustedLen}, |
| 5 | + ops::{range::legacy::Range, Try}, |
6 | 6 | }; |
7 | 7 |
|
8 | 8 | /// An iterator for stepping iterators by a custom amount. |
@@ -486,6 +486,89 @@ macro_rules! spec_int_ranges { |
486 | 486 | /// the outer length calculation won't encounter clamped values |
487 | 487 | #[unstable(feature = "trusted_len", issue = "37572")] |
488 | 488 | unsafe impl TrustedLen for StepBy<Range<$t>> {} |
| 489 | + |
| 490 | + impl SpecRangeSetup<RangeIter<$t>> for RangeIter<$t> { |
| 491 | + #[inline] |
| 492 | + fn setup(mut r: RangeIter<$t>, step: usize) -> RangeIter<$t> { |
| 493 | + r.inner = <Range<$t> as SpecRangeSetup::<Range<$t>>>::setup(r.inner.clone(), step); |
| 494 | + r |
| 495 | + } |
| 496 | + } |
| 497 | + |
| 498 | + unsafe impl StepByImpl<RangeIter<$t>> for StepBy<RangeIter<$t>> { |
| 499 | + #[inline] |
| 500 | + fn spec_next(&mut self) -> Option<$t> { |
| 501 | + // if a step size larger than the type has been specified fall back to |
| 502 | + // t::MAX, in which case remaining will be at most 1. |
| 503 | + // The `+ 1` can't overflow since the constructor substracted 1 from the original value. |
| 504 | + let step = <$t>::try_from(self.step + 1).unwrap_or(<$t>::MAX); |
| 505 | + let remaining = self.iter.inner.end; |
| 506 | + if remaining > 0 { |
| 507 | + let val = self.iter.inner.start; |
| 508 | + // this can only overflow during the last step, after which the value |
| 509 | + // will not be used |
| 510 | + self.iter.inner.start = val.wrapping_add(step); |
| 511 | + self.iter.inner.end = remaining - 1; |
| 512 | + Some(val) |
| 513 | + } else { |
| 514 | + None |
| 515 | + } |
| 516 | + } |
| 517 | + |
| 518 | + #[inline] |
| 519 | + fn spec_size_hint(&self) -> (usize, Option<usize>) { |
| 520 | + let remaining = self.iter.inner.end as usize; |
| 521 | + (remaining, Some(remaining)) |
| 522 | + } |
| 523 | + |
| 524 | + // The methods below are all copied from the Iterator trait default impls. |
| 525 | + // We have to repeat them here so that the specialization overrides the StepByImpl defaults |
| 526 | + |
| 527 | + #[inline] |
| 528 | + fn spec_nth(&mut self, n: usize) -> Option<Self::Item> { |
| 529 | + self.advance_by(n).ok()?; |
| 530 | + self.next() |
| 531 | + } |
| 532 | + |
| 533 | + #[inline] |
| 534 | + fn spec_try_fold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R |
| 535 | + where |
| 536 | + F: FnMut(Acc, Self::Item) -> R, |
| 537 | + R: Try<Output = Acc> |
| 538 | + { |
| 539 | + let mut accum = init; |
| 540 | + while let Some(x) = self.next() { |
| 541 | + accum = f(accum, x)?; |
| 542 | + } |
| 543 | + try { accum } |
| 544 | + } |
| 545 | + |
| 546 | + #[inline] |
| 547 | + fn spec_fold<Acc, F>(self, init: Acc, mut f: F) -> Acc |
| 548 | + where |
| 549 | + F: FnMut(Acc, Self::Item) -> Acc |
| 550 | + { |
| 551 | + // if a step size larger than the type has been specified fall back to |
| 552 | + // t::MAX, in which case remaining will be at most 1. |
| 553 | + let step = <$t>::try_from(self.step + 1).unwrap_or(<$t>::MAX); |
| 554 | + let remaining = self.iter.inner.end; |
| 555 | + let mut acc = init; |
| 556 | + let mut val = self.iter.inner.start; |
| 557 | + for _ in 0..remaining { |
| 558 | + acc = f(acc, val); |
| 559 | + // this can only overflow during the last step, after which the value |
| 560 | + // will no longer be used |
| 561 | + val = val.wrapping_add(step); |
| 562 | + } |
| 563 | + acc |
| 564 | + } |
| 565 | + } |
| 566 | + |
| 567 | + /// Safety: This macro is only applied to ranges over types <= usize |
| 568 | + /// which means the inner length is guaranteed to fit into a usize and so |
| 569 | + /// the outer length calculation won't encounter clamped values |
| 570 | + #[unstable(feature = "trusted_len", issue = "37572")] |
| 571 | + unsafe impl TrustedLen for StepBy<RangeIter<$t>> {} |
489 | 572 | )*) |
490 | 573 | } |
491 | 574 |
|
@@ -550,6 +633,64 @@ macro_rules! spec_int_ranges_r { |
550 | 633 | accum |
551 | 634 | } |
552 | 635 | } |
| 636 | + |
| 637 | + unsafe impl StepByBackImpl<RangeIter<$t>> for StepBy<RangeIter<$t>> { |
| 638 | + |
| 639 | + #[inline] |
| 640 | + fn spec_next_back(&mut self) -> Option<Self::Item> |
| 641 | + where RangeIter<$t>: DoubleEndedIterator + ExactSizeIterator, |
| 642 | + { |
| 643 | + let step = (self.step + 1) as $t; |
| 644 | + let remaining = self.iter.inner.end; |
| 645 | + if remaining > 0 { |
| 646 | + let start = self.iter.inner.start; |
| 647 | + self.iter.inner.end = remaining - 1; |
| 648 | + Some(start + step * (remaining - 1)) |
| 649 | + } else { |
| 650 | + None |
| 651 | + } |
| 652 | + } |
| 653 | + |
| 654 | + // The methods below are all copied from the Iterator trait default impls. |
| 655 | + // We have to repeat them here so that the specialization overrides the StepByImplBack defaults |
| 656 | + |
| 657 | + #[inline] |
| 658 | + fn spec_nth_back(&mut self, n: usize) -> Option<Self::Item> |
| 659 | + where Self: DoubleEndedIterator, |
| 660 | + { |
| 661 | + if self.advance_back_by(n).is_err() { |
| 662 | + return None; |
| 663 | + } |
| 664 | + self.next_back() |
| 665 | + } |
| 666 | + |
| 667 | + #[inline] |
| 668 | + fn spec_try_rfold<Acc, F, R>(&mut self, init: Acc, mut f: F) -> R |
| 669 | + where |
| 670 | + Self: DoubleEndedIterator, |
| 671 | + F: FnMut(Acc, Self::Item) -> R, |
| 672 | + R: Try<Output = Acc> |
| 673 | + { |
| 674 | + let mut accum = init; |
| 675 | + while let Some(x) = self.next_back() { |
| 676 | + accum = f(accum, x)?; |
| 677 | + } |
| 678 | + try { accum } |
| 679 | + } |
| 680 | + |
| 681 | + #[inline] |
| 682 | + fn spec_rfold<Acc, F>(mut self, init: Acc, mut f: F) -> Acc |
| 683 | + where |
| 684 | + Self: DoubleEndedIterator, |
| 685 | + F: FnMut(Acc, Self::Item) -> Acc |
| 686 | + { |
| 687 | + let mut accum = init; |
| 688 | + while let Some(x) = self.next_back() { |
| 689 | + accum = f(accum, x); |
| 690 | + } |
| 691 | + accum |
| 692 | + } |
| 693 | + } |
553 | 694 | )*) |
554 | 695 | } |
555 | 696 |
|
|
0 commit comments