1+ use crate :: iter:: { DoubleEndedIterator , FusedIterator , Iterator , TrustedLen } ;
12use crate :: ops:: Try ;
23use crate :: usize;
34
4- use super :: super :: { DoubleEndedIterator , FusedIterator , Iterator , TrustedLen } ;
5-
65/// An iterator that links two iterators together, in a chain.
76///
87/// This `struct` is created by the [`chain`] method on [`Iterator`]. See its
@@ -14,37 +13,34 @@ use super::super::{DoubleEndedIterator, FusedIterator, Iterator, TrustedLen};
1413#[ must_use = "iterators are lazy and do nothing unless consumed" ]
1514#[ stable( feature = "rust1" , since = "1.0.0" ) ]
1615pub struct Chain < A , B > {
17- a : A ,
18- b : B ,
19- state : ChainState ,
16+ // These are "fused" with `Option` so we don't need separate state to track which part is
17+ // already exhausted, and we may also get niche layout for `None`. We don't use the real `Fuse`
18+ // adapter because its specialization for `FusedIterator` unconditionally descends into the
19+ // iterator, and that could be expensive to keep revisiting stuff like nested chains. It also
20+ // hurts compiler performance to add more iterator layers to `Chain`.
21+ a : Option < A > ,
22+ b : Option < B > ,
2023}
2124impl < A , B > Chain < A , B > {
2225 pub ( in super :: super ) fn new ( a : A , b : B ) -> Chain < A , B > {
23- Chain { a, b , state : ChainState :: Both }
26+ Chain { a : Some ( a ) , b : Some ( b ) }
2427 }
2528}
2629
27- // The iterator protocol specifies that iteration ends with the return value
28- // `None` from `.next()` (or `.next_back()`) and it is unspecified what
29- // further calls return. The chain adaptor must account for this since it uses
30- // two subiterators.
31- //
32- // It uses three states:
33- //
34- // - Both: `a` and `b` are remaining
35- // - Front: `a` remaining
36- // - Back: `b` remaining
37- //
38- // The fourth state (neither iterator is remaining) only occurs after Chain has
39- // returned None once, so we don't need to store this state.
40- #[ derive( Clone , Debug ) ]
41- enum ChainState {
42- // both front and back iterator are remaining
43- Both ,
44- // only front is remaining
45- Front ,
46- // only back is remaining
47- Back ,
30+ /// Fuse the iterator if the expression is `None`.
31+ macro_rules! fuse {
32+ ( $self: ident . $iter: ident . $( $call: tt) +) => {
33+ match $self. $iter {
34+ Some ( ref mut iter) => match iter. $( $call) + {
35+ None => {
36+ $self. $iter = None ;
37+ None
38+ }
39+ item => item,
40+ } ,
41+ None => None ,
42+ }
43+ } ;
4844}
4945
5046#[ stable( feature = "rust1" , since = "1.0.0" ) ]
@@ -57,128 +53,101 @@ where
5753
5854 #[ inline]
5955 fn next ( & mut self ) -> Option < A :: Item > {
60- match self . state {
61- ChainState :: Both => match self . a . next ( ) {
62- elt @ Some ( ..) => elt,
63- None => {
64- self . state = ChainState :: Back ;
65- self . b . next ( )
66- }
67- } ,
68- ChainState :: Front => self . a . next ( ) ,
69- ChainState :: Back => self . b . next ( ) ,
56+ match fuse ! ( self . a. next( ) ) {
57+ None => fuse ! ( self . b. next( ) ) ,
58+ item => item,
7059 }
7160 }
7261
7362 #[ inline]
7463 #[ rustc_inherit_overflow_checks]
7564 fn count ( self ) -> usize {
76- match self . state {
77- ChainState :: Both => self . a . count ( ) + self . b . count ( ) ,
78- ChainState :: Front => self . a . count ( ) ,
79- ChainState :: Back => self . b . count ( ) ,
80- }
65+ let a_count = match self . a {
66+ Some ( a) => a. count ( ) ,
67+ None => 0 ,
68+ } ;
69+ let b_count = match self . b {
70+ Some ( b) => b. count ( ) ,
71+ None => 0 ,
72+ } ;
73+ a_count + b_count
8174 }
8275
83- fn try_fold < Acc , F , R > ( & mut self , init : Acc , mut f : F ) -> R
76+ fn try_fold < Acc , F , R > ( & mut self , mut acc : Acc , mut f : F ) -> R
8477 where
8578 Self : Sized ,
8679 F : FnMut ( Acc , Self :: Item ) -> R ,
8780 R : Try < Ok = Acc > ,
8881 {
89- let mut accum = init;
90- match self . state {
91- ChainState :: Both | ChainState :: Front => {
92- accum = self . a . try_fold ( accum, & mut f) ?;
93- if let ChainState :: Both = self . state {
94- self . state = ChainState :: Back ;
95- }
96- }
97- _ => { }
82+ if let Some ( ref mut a) = self . a {
83+ acc = a. try_fold ( acc, & mut f) ?;
84+ self . a = None ;
9885 }
99- if let ChainState :: Back = self . state {
100- accum = self . b . try_fold ( accum, & mut f) ?;
86+ if let Some ( ref mut b) = self . b {
87+ acc = b. try_fold ( acc, f) ?;
88+ self . b = None ;
10189 }
102- Try :: from_ok ( accum )
90+ Try :: from_ok ( acc )
10391 }
10492
105- fn fold < Acc , F > ( self , init : Acc , mut f : F ) -> Acc
93+ fn fold < Acc , F > ( self , mut acc : Acc , mut f : F ) -> Acc
10694 where
10795 F : FnMut ( Acc , Self :: Item ) -> Acc ,
10896 {
109- let mut accum = init;
110- match self . state {
111- ChainState :: Both | ChainState :: Front => {
112- accum = self . a . fold ( accum, & mut f) ;
113- }
114- _ => { }
97+ if let Some ( a) = self . a {
98+ acc = a. fold ( acc, & mut f) ;
11599 }
116- match self . state {
117- ChainState :: Both | ChainState :: Back => {
118- accum = self . b . fold ( accum, & mut f) ;
119- }
120- _ => { }
100+ if let Some ( b) = self . b {
101+ acc = b. fold ( acc, f) ;
121102 }
122- accum
103+ acc
123104 }
124105
125106 #[ inline]
126107 fn nth ( & mut self , mut n : usize ) -> Option < A :: Item > {
127- match self . state {
128- ChainState :: Both | ChainState :: Front => {
129- for x in self . a . by_ref ( ) {
130- if n == 0 {
131- return Some ( x) ;
132- }
133- n -= 1 ;
134- }
135- if let ChainState :: Both = self . state {
136- self . state = ChainState :: Back ;
108+ if let Some ( ref mut a) = self . a {
109+ while let Some ( x) = a. next ( ) {
110+ if n == 0 {
111+ return Some ( x) ;
137112 }
113+ n -= 1 ;
138114 }
139- ChainState :: Back => { }
115+ self . a = None ;
140116 }
141- if let ChainState :: Back = self . state { self . b . nth ( n) } else { None }
117+ fuse ! ( self . b. nth( n) )
142118 }
143119
144120 #[ inline]
145121 fn find < P > ( & mut self , mut predicate : P ) -> Option < Self :: Item >
146122 where
147123 P : FnMut ( & Self :: Item ) -> bool ,
148124 {
149- match self . state {
150- ChainState :: Both => match self . a . find ( & mut predicate) {
151- None => {
152- self . state = ChainState :: Back ;
153- self . b . find ( predicate)
154- }
155- v => v,
156- } ,
157- ChainState :: Front => self . a . find ( predicate) ,
158- ChainState :: Back => self . b . find ( predicate) ,
125+ match fuse ! ( self . a. find( & mut predicate) ) {
126+ None => fuse ! ( self . b. find( predicate) ) ,
127+ item => item,
159128 }
160129 }
161130
162131 #[ inline]
163132 fn last ( self ) -> Option < A :: Item > {
164- match self . state {
165- ChainState :: Both => {
166- // Must exhaust a before b.
167- let a_last = self . a . last ( ) ;
168- let b_last = self . b . last ( ) ;
169- b_last . or ( a_last )
170- }
171- ChainState :: Front => self . a . last ( ) ,
172- ChainState :: Back => self . b . last ( ) ,
173- }
133+ // Must exhaust a before b.
134+ let a_last = match self . a {
135+ Some ( a ) => a . last ( ) ,
136+ None => None ,
137+ } ;
138+ let b_last = match self . b {
139+ Some ( b ) => b . last ( ) ,
140+ None => None ,
141+ } ;
142+ b_last . or ( a_last )
174143 }
175144
176145 #[ inline]
177146 fn size_hint ( & self ) -> ( usize , Option < usize > ) {
178- match self . state {
179- ChainState :: Both => {
180- let ( a_lower, a_upper) = self . a . size_hint ( ) ;
181- let ( b_lower, b_upper) = self . b . size_hint ( ) ;
147+ match self {
148+ Chain { a : Some ( a ) , b : Some ( b ) } => {
149+ let ( a_lower, a_upper) = a. size_hint ( ) ;
150+ let ( b_lower, b_upper) = b. size_hint ( ) ;
182151
183152 let lower = a_lower. saturating_add ( b_lower) ;
184153
@@ -189,8 +158,9 @@ where
189158
190159 ( lower, upper)
191160 }
192- ChainState :: Front => self . a . size_hint ( ) ,
193- ChainState :: Back => self . b . size_hint ( ) ,
161+ Chain { a : Some ( a) , b : None } => a. size_hint ( ) ,
162+ Chain { a : None , b : Some ( b) } => b. size_hint ( ) ,
163+ Chain { a : None , b : None } => ( 0 , Some ( 0 ) ) ,
194164 }
195165 }
196166}
@@ -203,82 +173,71 @@ where
203173{
204174 #[ inline]
205175 fn next_back ( & mut self ) -> Option < A :: Item > {
206- match self . state {
207- ChainState :: Both => match self . b . next_back ( ) {
208- elt @ Some ( ..) => elt,
209- None => {
210- self . state = ChainState :: Front ;
211- self . a . next_back ( )
212- }
213- } ,
214- ChainState :: Front => self . a . next_back ( ) ,
215- ChainState :: Back => self . b . next_back ( ) ,
176+ match fuse ! ( self . b. next_back( ) ) {
177+ None => fuse ! ( self . a. next_back( ) ) ,
178+ item => item,
216179 }
217180 }
218181
219182 #[ inline]
220183 fn nth_back ( & mut self , mut n : usize ) -> Option < A :: Item > {
221- match self . state {
222- ChainState :: Both | ChainState :: Back => {
223- for x in self . b . by_ref ( ) . rev ( ) {
224- if n == 0 {
225- return Some ( x) ;
226- }
227- n -= 1 ;
228- }
229- if let ChainState :: Both = self . state {
230- self . state = ChainState :: Front ;
184+ if let Some ( ref mut b) = self . b {
185+ while let Some ( x) = b. next_back ( ) {
186+ if n == 0 {
187+ return Some ( x) ;
231188 }
189+ n -= 1 ;
232190 }
233- ChainState :: Front => { }
191+ self . b = None ;
234192 }
235- if let ChainState :: Front = self . state { self . a . nth_back ( n) } else { None }
193+ fuse ! ( self . a. nth_back( n) )
236194 }
237195
238- fn try_rfold < Acc , F , R > ( & mut self , init : Acc , mut f : F ) -> R
196+ #[ inline]
197+ fn rfind < P > ( & mut self , mut predicate : P ) -> Option < Self :: Item >
198+ where
199+ P : FnMut ( & Self :: Item ) -> bool ,
200+ {
201+ match fuse ! ( self . b. rfind( & mut predicate) ) {
202+ None => fuse ! ( self . a. rfind( predicate) ) ,
203+ item => item,
204+ }
205+ }
206+
207+ fn try_rfold < Acc , F , R > ( & mut self , mut acc : Acc , mut f : F ) -> R
239208 where
240209 Self : Sized ,
241210 F : FnMut ( Acc , Self :: Item ) -> R ,
242211 R : Try < Ok = Acc > ,
243212 {
244- let mut accum = init;
245- match self . state {
246- ChainState :: Both | ChainState :: Back => {
247- accum = self . b . try_rfold ( accum, & mut f) ?;
248- if let ChainState :: Both = self . state {
249- self . state = ChainState :: Front ;
250- }
251- }
252- _ => { }
213+ if let Some ( ref mut b) = self . b {
214+ acc = b. try_rfold ( acc, & mut f) ?;
215+ self . b = None ;
253216 }
254- if let ChainState :: Front = self . state {
255- accum = self . a . try_rfold ( accum, & mut f) ?;
217+ if let Some ( ref mut a) = self . a {
218+ acc = a. try_rfold ( acc, f) ?;
219+ self . a = None ;
256220 }
257- Try :: from_ok ( accum )
221+ Try :: from_ok ( acc )
258222 }
259223
260- fn rfold < Acc , F > ( self , init : Acc , mut f : F ) -> Acc
224+ fn rfold < Acc , F > ( self , mut acc : Acc , mut f : F ) -> Acc
261225 where
262226 F : FnMut ( Acc , Self :: Item ) -> Acc ,
263227 {
264- let mut accum = init;
265- match self . state {
266- ChainState :: Both | ChainState :: Back => {
267- accum = self . b . rfold ( accum, & mut f) ;
268- }
269- _ => { }
228+ if let Some ( b) = self . b {
229+ acc = b. rfold ( acc, & mut f) ;
270230 }
271- match self . state {
272- ChainState :: Both | ChainState :: Front => {
273- accum = self . a . rfold ( accum, & mut f) ;
274- }
275- _ => { }
231+ if let Some ( a) = self . a {
232+ acc = a. rfold ( acc, f) ;
276233 }
277- accum
234+ acc
278235 }
279236}
280237
281238// Note: *both* must be fused to handle double-ended iterators.
239+ // Now that we "fuse" both sides, we *could* implement this unconditionally,
240+ // but we should be cautious about committing to that in the public API.
282241#[ stable( feature = "fused" , since = "1.26.0" ) ]
283242impl < A , B > FusedIterator for Chain < A , B >
284243where
0 commit comments