1414
1515
1616use std:: util:: { swap, replace} ;
17- use std:: iterator:: { FromIterator , Extendable } ;
17+ use std:: iterator:: { FromIterator , Extendable , Peekable } ;
18+ use std:: cmp:: Ordering ;
1819
1920// This is implemented as an AA tree, which is a simplified variation of
2021// a red-black tree where red (horizontal) nodes can only be added
@@ -529,24 +530,24 @@ impl<T: TotalOrd> TreeSet<T> {
529530
530531 /// Visit the values (in-order) representing the difference
531532 pub fn difference < ' a > ( & ' a self , other : & ' a TreeSet < T > ) -> Difference < ' a , T > {
532- Difference { a : Focus :: new ( self . iter ( ) ) , b : Focus :: new ( other. iter ( ) ) }
533+ Difference { a : self . iter ( ) . peekable ( ) , b : other. iter ( ) . peekable ( ) }
533534 }
534535
535536 /// Visit the values (in-order) representing the symmetric difference
536537 pub fn symmetric_difference < ' a > ( & ' a self , other : & ' a TreeSet < T > )
537538 -> SymDifference < ' a , T > {
538- SymDifference { a : Focus :: new ( self . iter ( ) ) , b : Focus :: new ( other. iter ( ) ) }
539+ SymDifference { a : self . iter ( ) . peekable ( ) , b : other. iter ( ) . peekable ( ) }
539540 }
540541
541542 /// Visit the values (in-order) representing the intersection
542543 pub fn intersection < ' a > ( & ' a self , other : & ' a TreeSet < T > )
543544 -> Intersection < ' a , T > {
544- Intersection { a : Focus :: new ( self . iter ( ) ) , b : Focus :: new ( other. iter ( ) ) }
545+ Intersection { a : self . iter ( ) . peekable ( ) , b : other. iter ( ) . peekable ( ) }
545546 }
546547
547548 /// Visit the values (in-order) representing the union
548549 pub fn union < ' a > ( & ' a self , other : & ' a TreeSet < T > ) -> Union < ' a , T > {
549- Union { a : Focus :: new ( self . iter ( ) ) , b : Focus :: new ( other. iter ( ) ) }
550+ Union { a : self . iter ( ) . peekable ( ) , b : other. iter ( ) . peekable ( ) }
550551 }
551552}
552553
@@ -560,61 +561,47 @@ pub struct TreeSetRevIterator<'self, T> {
560561 priv iter : TreeMapRevIterator < ' self , T , ( ) >
561562}
562563
563- // Encapsulate an iterator and hold its latest value until stepped forward
564- struct Focus < A , T > {
565- priv iter : T ,
566- priv focus : Option < A > ,
567- }
568-
569- impl < A , T : Iterator < A > > Focus < A , T > {
570- fn new ( mut it : T ) -> Focus < A , T > {
571- Focus { focus : it. next ( ) , iter : it}
572- }
573- fn step ( & mut self ) {
574- self . focus = self . iter . next ( )
575- }
576- }
577-
578564/// Lazy iterator producing elements in the set difference (in-order)
579565pub struct Difference < ' self , T > {
580- priv a: Focus < & ' self T , TreeSetIterator < ' self , T > > ,
581- priv b: Focus < & ' self T , TreeSetIterator < ' self , T > > ,
566+ priv a: Peekable < & ' self T , TreeSetIterator < ' self , T > > ,
567+ priv b: Peekable < & ' self T , TreeSetIterator < ' self , T > > ,
582568}
583569
584570/// Lazy iterator producing elements in the set symmetric difference (in-order)
585571pub struct SymDifference < ' self , T > {
586- priv a: Focus < & ' self T , TreeSetIterator < ' self , T > > ,
587- priv b: Focus < & ' self T , TreeSetIterator < ' self , T > > ,
572+ priv a: Peekable < & ' self T , TreeSetIterator < ' self , T > > ,
573+ priv b: Peekable < & ' self T , TreeSetIterator < ' self , T > > ,
588574}
589575
590576/// Lazy iterator producing elements in the set intersection (in-order)
591577pub struct Intersection < ' self , T > {
592- priv a: Focus < & ' self T , TreeSetIterator < ' self , T > > ,
593- priv b: Focus < & ' self T , TreeSetIterator < ' self , T > > ,
578+ priv a: Peekable < & ' self T , TreeSetIterator < ' self , T > > ,
579+ priv b: Peekable < & ' self T , TreeSetIterator < ' self , T > > ,
594580}
595581
596582/// Lazy iterator producing elements in the set intersection (in-order)
597583pub struct Union < ' self , T > {
598- priv a: Focus < & ' self T , TreeSetIterator < ' self , T > > ,
599- priv b: Focus < & ' self T , TreeSetIterator < ' self , T > > ,
584+ priv a: Peekable < & ' self T , TreeSetIterator < ' self , T > > ,
585+ priv b: Peekable < & ' self T , TreeSetIterator < ' self , T > > ,
586+ }
587+
588+ /// Compare `x` and `y`, but return `short` if x is None and `long` if y is None
589+ fn cmp_opt < T : TotalOrd > ( x : Option < & T > , y : Option < & T > ,
590+ short : Ordering , long : Ordering ) -> Ordering {
591+ match ( x, y) {
592+ ( None , _ ) => short,
593+ ( _ , None ) => long,
594+ ( Some ( x1) , Some ( y1) ) => x1. cmp ( y1) ,
595+ }
600596}
601597
602598impl < ' self , T : TotalOrd > Iterator < & ' self T > for Difference < ' self , T > {
603599 fn next ( & mut self ) -> Option < & ' self T > {
604600 loop {
605- match ( self . a . focus , self . b . focus ) {
606- ( None , _ ) => return None ,
607- ( ret , None ) => { self . a . step ( ) ; return ret } ,
608- ( Some ( a1) , Some ( b1) ) => {
609- let cmp = a1. cmp ( b1) ;
610- if cmp == Less {
611- self . a . step ( ) ;
612- return Some ( a1) ;
613- } else {
614- if cmp == Equal { self . a . step ( ) }
615- self . b . step ( ) ;
616- }
617- }
601+ match cmp_opt ( self . a . peek ( ) , self . b . peek ( ) , Less , Less ) {
602+ Less => return self . a . next ( ) ,
603+ Equal => { self . a . next ( ) ; self . b . next ( ) ; }
604+ Greater => { self . b . next ( ) ; }
618605 }
619606 }
620607 }
@@ -623,23 +610,10 @@ impl<'self, T: TotalOrd> Iterator<&'self T> for Difference<'self, T> {
623610impl < ' self , T : TotalOrd > Iterator < & ' self T > for SymDifference < ' self , T > {
624611 fn next ( & mut self ) -> Option < & ' self T > {
625612 loop {
626- match ( self . a . focus , self . b . focus ) {
627- ( ret , None ) => { self . a . step ( ) ; return ret } ,
628- ( None , ret ) => { self . b . step ( ) ; return ret } ,
629- ( Some ( a1) , Some ( b1) ) => {
630- let cmp = a1. cmp ( b1) ;
631- if cmp == Less {
632- self . a . step ( ) ;
633- return Some ( a1) ;
634- } else {
635- self . b . step ( ) ;
636- if cmp == Greater {
637- return Some ( b1) ;
638- } else {
639- self . a . step ( ) ;
640- }
641- }
642- }
613+ match cmp_opt ( self . a . peek ( ) , self . b . peek ( ) , Greater , Less ) {
614+ Less => return self . a . next ( ) ,
615+ Equal => { self . a . next ( ) ; self . b . next ( ) ; }
616+ Greater => return self . b . next ( ) ,
643617 }
644618 }
645619 }
@@ -648,20 +622,16 @@ impl<'self, T: TotalOrd> Iterator<&'self T> for SymDifference<'self, T> {
648622impl < ' self , T : TotalOrd > Iterator < & ' self T > for Intersection < ' self , T > {
649623 fn next ( & mut self ) -> Option < & ' self T > {
650624 loop {
651- match ( self . a . focus , self . b . focus ) {
652- ( None , _ ) => return None ,
653- ( _ , None ) => return None ,
654- ( Some ( a1) , Some ( b1) ) => {
655- let cmp = a1. cmp ( b1) ;
656- if cmp == Less {
657- self . a . step ( ) ;
658- } else {
659- self . b . step ( ) ;
660- if cmp == Equal {
661- return Some ( a1) ;
662- }
663- }
664- } ,
625+ let o_cmp = match ( self . a . peek ( ) , self . b . peek ( ) ) {
626+ ( None , _ ) => None ,
627+ ( _ , None ) => None ,
628+ ( Some ( a1) , Some ( b1) ) => Some ( a1. cmp ( b1) ) ,
629+ } ;
630+ match o_cmp {
631+ None => return None ,
632+ Some ( Less ) => { self . a . next ( ) ; }
633+ Some ( Equal ) => { self . b . next ( ) ; return self . a . next ( ) }
634+ Some ( Greater ) => { self . b . next ( ) ; }
665635 }
666636 }
667637 }
@@ -670,22 +640,10 @@ impl<'self, T: TotalOrd> Iterator<&'self T> for Intersection<'self, T> {
670640impl < ' self , T : TotalOrd > Iterator < & ' self T > for Union < ' self , T > {
671641 fn next ( & mut self ) -> Option < & ' self T > {
672642 loop {
673- match ( self . a . focus , self . b . focus ) {
674- ( ret , None ) => { self . a . step ( ) ; return ret } ,
675- ( None , ret ) => { self . b . step ( ) ; return ret } ,
676- ( Some ( a1) , Some ( b1) ) => {
677- let cmp = a1. cmp ( b1) ;
678- if cmp == Greater {
679- self . b . step ( ) ;
680- return Some ( b1) ;
681- } else {
682- self . a . step ( ) ;
683- if cmp == Equal {
684- self . b . step ( ) ;
685- }
686- return Some ( a1) ;
687- }
688- }
643+ match cmp_opt ( self . a . peek ( ) , self . b . peek ( ) , Greater , Less ) {
644+ Less => return self . a . next ( ) ,
645+ Equal => { self . b . next ( ) ; return self . a . next ( ) }
646+ Greater => return self . b . next ( ) ,
689647 }
690648 }
691649 }
0 commit comments