Implementation of Haskell's takeWhile

accept-while.p6

use nqp;
use Test;

# this was easier to create than accept-while
sub Iterator-accept-until ( \iterator, &condition ) {
    class :: does Iterator {
        has $!source;
        has &!condition;
        has int $!running;

        method !SET-SELF($!source,&!condition) {
            nqp::stmts(
              ($!running = 1),
              self
            )
        }
        method new(\source,\condition){
            nqp::create(self)!SET-SELF(source,condition)
        }
        method pull-one() is raw {
            nqp::if(
              $!running,
              nqp::if(
                nqp::eqaddr(
                  (my Mu $pulled := $!source.pull-one),
                  IterationEnd
                ) || &!condition($pulled),
                nqp::stmts(
                  ($!running = 0), # stop the iterator
                  IterationEnd
                ),
                $pulled # return
              ),
              IterationEnd
            )
        }
        method is-lazy() { $!source.is-lazy }
    }.new(iterator,&condition)
}

sub accept-until ( Iterable \sequence, &condition ) {
  Seq.new: Iterator-accept-until sequence.iterator, &condition
}
sub accept-while ( Iterable \sequence, &condition ) {
  Seq.new: Iterator-accept-until sequence.iterator, -> |args { not condition |args }
}

is-deeply accept-until( ^20, * == 5  ).List, (^5).List;
is-deeply accept-until(  ^5, {False} ).List, (^5).List;
is-deeply accept-until( ^20, {True}  ).List, ();

{
  my $counter = -1;
  is-deeply accept-until(^20,{++$counter >= 5}), (^5).List;
  is $counter, 5;
}

is-deeply accept-while( ^20, * !== 5 ).List, (^5).List;
is-deeply accept-while( ^20, {False} ).List, ();
is-deeply accept-while(  ^5, {True}  ).List, (^5).List;

{
  my $counter = -1;
  is-deeply accept-while(^20,{++$counter < 5}), (^5).List;
  is $counter, 5;
}