Skip to content

Commit 4f184eb

Browse files
authored
Rollup merge of rust-lang#48012 - scottmcm:faster-rangeinclusive-fold, r=alexcrichton
Override try_[r]fold for RangeInclusive Because the last item needs special handling, it seems that LLVM has trouble canonicalizing the loops in external iteration. With the override, it becomes obvious that the start==end case exits the loop (as opposed to the one *after* that exiting the loop in external iteration). Demo adapted from rust-lang#45222 ```rust #[no_mangle] pub fn foo3r(n: u64) -> u64 { let mut count = 0; (0..n).for_each(|_| { (0 ..= n).rev().for_each(|j| { count += j; }) }); count } ``` <details> <summary>Current nightly ASM, 100 lines (https://play.rust-lang.org/?gist=f5674c702c6e2045c3aab5d03763e5f6&version=nightly&mode=release)</summary> ```asm foo3r: pushq %rbx .Lcfi0: .Lcfi1: testq %rdi, %rdi je .LBB0_1 testb $1, %dil jne .LBB0_4 xorl %eax, %eax xorl %r8d, %r8d cmpq $1, %rdi jne .LBB0_11 jmp .LBB0_23 .LBB0_1: xorl %eax, %eax popq %rbx retq .LBB0_4: xorl %r8d, %r8d movq $-1, %r9 xorl %eax, %eax movq %rdi, %r11 xorl %r10d, %r10d jmp .LBB0_5 .LBB0_8: addq %r11, %rax movq %rsi, %r11 movq %rdx, %r10 .LBB0_5: cmpq %r11, %r10 movl $1, %ecx cmovbq %r9, %rcx cmoveq %r8, %rcx testq %rcx, %rcx movl $0, %esi movl $1, %edx je .LBB0_8 cmpq $-1, %rcx jne .LBB0_9 leaq -1(%r11), %rsi movq %r10, %rdx jmp .LBB0_8 .LBB0_9: movl $1, %r8d cmpq $1, %rdi je .LBB0_23 .LBB0_11: xorl %r9d, %r9d movq $-1, %r10 .LBB0_12: movq %rdi, %rsi xorl %r11d, %r11d jmp .LBB0_13 .LBB0_16: addq %rsi, %rax movq %rcx, %rsi movq %rbx, %r11 .LBB0_13: cmpq %rsi, %r11 movl $1, %edx cmovbq %r10, %rdx cmoveq %r9, %rdx testq %rdx, %rdx movl $0, %ecx movl $1, %ebx je .LBB0_16 cmpq $-1, %rdx jne .LBB0_17 leaq -1(%rsi), %rcx movq %r11, %rbx jmp .LBB0_16 .LBB0_17: movq %rdi, %rcx xorl %r11d, %r11d jmp .LBB0_18 .LBB0_21: addq %rcx, %rax movq %rsi, %rcx movq %rbx, %r11 .LBB0_18: cmpq %rcx, %r11 movl $1, %edx cmovbq %r10, %rdx cmoveq %r9, %rdx testq %rdx, %rdx movl $0, %esi movl $1, %ebx je .LBB0_21 cmpq $-1, %rdx jne .LBB0_22 leaq -1(%rcx), %rsi movq %r11, %rbx jmp .LBB0_21 .LBB0_22: addq $2, %r8 cmpq %rdi, %r8 jne .LBB0_12 .LBB0_23: popq %rbx retq .Lfunc_end0: ``` </details><br> With this PR: ```asm foo3r: test rcx, rcx je .LBB3_1 lea r8, [rcx - 1] lea rdx, [rcx - 2] mov rax, r8 mul rdx shld rdx, rax, 63 imul r8, r8 add r8, rcx sub r8, rdx imul r8, rcx mov rax, r8 ret .LBB3_1: xor r8d, r8d mov rax, r8 ret ```
2 parents a026e8a + 1b1e887 commit 4f184eb

File tree

2 files changed

+65
-1
lines changed

2 files changed

+65
-1
lines changed

src/libcore/iter/range.rs

+45-1
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
use convert::TryFrom;
1212
use mem;
13-
use ops::{self, Add, Sub};
13+
use ops::{self, Add, Sub, Try};
1414
use usize;
1515

1616
use super::{FusedIterator, TrustedLen};
@@ -397,6 +397,28 @@ impl<A: Step> Iterator for ops::RangeInclusive<A> {
397397
fn max(mut self) -> Option<A> {
398398
self.next_back()
399399
}
400+
401+
#[inline]
402+
fn try_fold<B, F, R>(&mut self, init: B, mut f: F) -> R where
403+
Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try<Ok=B>
404+
{
405+
let mut accum = init;
406+
if self.start <= self.end {
407+
loop {
408+
let (x, done) =
409+
if self.start < self.end {
410+
let n = self.start.add_one();
411+
(mem::replace(&mut self.start, n), false)
412+
} else {
413+
self.end.replace_zero();
414+
(self.start.replace_one(), true)
415+
};
416+
accum = f(accum, x)?;
417+
if done { break }
418+
}
419+
}
420+
Try::from_ok(accum)
421+
}
400422
}
401423

402424
#[unstable(feature = "inclusive_range", reason = "recently added, follows RFC", issue = "28237")]
@@ -418,6 +440,28 @@ impl<A: Step> DoubleEndedIterator for ops::RangeInclusive<A> {
418440
_ => None,
419441
}
420442
}
443+
444+
#[inline]
445+
fn try_rfold<B, F, R>(&mut self, init: B, mut f: F) -> R where
446+
Self: Sized, F: FnMut(B, Self::Item) -> R, R: Try<Ok=B>
447+
{
448+
let mut accum = init;
449+
if self.start <= self.end {
450+
loop {
451+
let (x, done) =
452+
if self.start < self.end {
453+
let n = self.end.sub_one();
454+
(mem::replace(&mut self.end, n), false)
455+
} else {
456+
self.start.replace_one();
457+
(self.end.replace_zero(), true)
458+
};
459+
accum = f(accum, x)?;
460+
if done { break }
461+
}
462+
}
463+
Try::from_ok(accum)
464+
}
421465
}
422466

423467
#[unstable(feature = "fused", issue = "35602")]

src/libcore/tests/iter.rs

+20
Original file line numberDiff line numberDiff line change
@@ -1459,6 +1459,26 @@ fn test_range_inclusive_min() {
14591459
assert_eq!(r.min(), None);
14601460
}
14611461

1462+
#[test]
1463+
fn test_range_inclusive_folds() {
1464+
assert_eq!((1..=10).sum::<i32>(), 55);
1465+
assert_eq!((1..=10).rev().sum::<i32>(), 55);
1466+
1467+
let mut it = 40..=50;
1468+
assert_eq!(it.try_fold(0, i8::checked_add), None);
1469+
assert_eq!(it, 44..=50);
1470+
assert_eq!(it.try_rfold(0, i8::checked_add), None);
1471+
assert_eq!(it, 44..=47);
1472+
1473+
let mut it = 10..=20;
1474+
assert_eq!(it.try_fold(0, |a,b| Some(a+b)), Some(165));
1475+
assert_eq!(it, 1..=0);
1476+
1477+
let mut it = 10..=20;
1478+
assert_eq!(it.try_rfold(0, |a,b| Some(a+b)), Some(165));
1479+
assert_eq!(it, 1..=0);
1480+
}
1481+
14621482
#[test]
14631483
fn test_repeat() {
14641484
let mut it = repeat(42);

0 commit comments

Comments
 (0)