Put some iterator goodness in Iterable.flat

This should cause some performance gains in some places, hard to tell where.
rakudo · Oct 30, 2015 · ba70274 · ba70274
1 parent 131af2c
commit ba70274
Showing 1 changed file with 43 additions and 23 deletions.
diff --git a/src/core/Iterable.pm b/src/core/Iterable.pm
@@ -17,44 +17,64 @@ my role Iterable {
 
     method flat(Iterable:D:) {
         Seq.new(class :: does Iterator {
-            has $!source;
-            has Iterator $!nested-iter;
+            has Iterator $!source;
+            has Iterator $!nested;
 
-            method new(\source-iter) {
-                my \iter = self.CREATE;
-                nqp::bindattr(iter, self, '$!source', source-iter);
+            method new(\source) {
+                my \iter = nqp::create(self);
+                nqp::bindattr(iter, self, '$!source', source);
                 iter
             }
 
             my constant NO_RESULT_YET = Mu.CREATE;
             method pull-one() is raw {
                 my $result := NO_RESULT_YET;
                 my $got;
-                repeat while nqp::eqaddr($result, NO_RESULT_YET) {
-                    if $!nested-iter {
-                        $got := $!nested-iter.pull-one();
-                        if nqp::eqaddr($got, IterationEnd) {
-                            $!nested-iter := Iterator;
-                        }
-                        else {
-                            $result := $got;
-                        }
+                while nqp::eqaddr($result, NO_RESULT_YET) {
+                    if $!nested {
+                        $got := $!nested.pull-one;
+                        nqp::eqaddr($got, IterationEnd)
+                          ?? ($!nested := Iterator)
+                          !! ($result := $got);
                     }
                     else {
                         $got := $!source.pull-one();
-                        if nqp::istype($got, Iterable) && !nqp::iscont($got) {
-                            $!nested-iter := $got.flat.iterator;
-                        }
-                        else {
-                            $result := $got;
-                        }
+                        nqp::istype($got, Iterable) && !nqp::iscont($got)
+                          ?? ($!nested := $got.flat.iterator)
+                          !! ($result := $got);
                     }
                 }
                 $result
             }
-
-            # This is a prime candidate for implementing most of the other
-            # methods, for speed reasons
+            method push-all($target) {
+                my $got;
+                until ($got := $!source.pull-one) =:= IterationEnd {
+                    if nqp::istype($got, Iterable) && !nqp::iscont($got) {
+                        my $nested := $got.flat.iterator;
+                        $target.push($got)
+                          until ($got := $nested.pull-one) =:= IterationEnd;
+                    }
+                    else {
+                        $target.push($got);
+                    }
+                }
+                IterationEnd
+            }
+            method count-only() {
+                my int $found;
+                my $got;
+                until ($got := $!source.pull-one) =:= IterationEnd {
+                    if nqp::istype($got, Iterable) && !nqp::iscont($got) {
+                        my $nested := $got.flat.iterator;
+                        $found = $found + 1
+                          until ($got := $nested.pull-one) =:= IterationEnd;
+                    }
+                    else {
+                        $found = $found + 1;
+                    }
+                }
+                nqp::p6box_i($found)
+            }
         }.new(self.iterator))
     }