File Coverage

blib/lib/Tree/RB.pm

Criterion	Covered	Total	%
statement	223	285	78.2
branch	112	170	65.8
condition	29	50	58.0
subroutine	33	37	89.1
pod	11	11	100.0
total	408	553	73.7

package Tree::RB;

use strict;
use Carp;

use Tree::RB::Node qw[set_color color_of parent_of left_of right_of];
use Tree::RB::Node::_Constants;
use vars qw( $VERSION @EXPORT_OK );
$VERSION = '0.500005';
$VERSION = eval $VERSION;

require Exporter;
*import    = \&Exporter::import;
@EXPORT_OK = qw[LUEQUAL LUGTEQ LULTEQ LUGREAT LULESS LUNEXT LUPREV];

use enum qw{
    LUEQUAL
    LUGTEQ
    LULTEQ
    LUGREAT
    LULESS
    LUNEXT
    LUPREV
};

# object slots
use enum qw{
    ROOT
    CMP
    SIZE
    HASH_ITER
    HASH_SEEK_ARG
};

# Node and hash Iteration

sub _mk_iter {
    my $start_fn = shift || 'min';
    my $next_fn  = shift || 'successor';
    return sub {
        my $self = shift;
        my $key  = shift;
        my $node;
        my $iter = sub {
            if($node) {
                $node = $node->$next_fn;
            }
            else {
                if(defined $key) {
                    # seek to $key
                    (undef, $node) = $self->lookup(
                        $key,
                        $next_fn eq 'successor' ? LUGTEQ : LULTEQ
                    );
                }
                else {
                    $node = $self->$start_fn;
                }
            }
            return $node;
        };
        return bless($iter => 'Tree::RB::Iterator');
    };
}

*Tree::RB::Iterator::next = sub { $_[0]->() };

*iter     = _mk_iter(qw/min successor/);
*rev_iter = _mk_iter(qw/max predecessor/);

sub hseek {
    my $self = shift;
    my $arg  = shift;
    defined $arg or croak("Can't seek to an undefined key");
    my %args;
    if(ref $arg eq 'HASH') {
        %args = %$arg;
    }
    else {
        $args{-key} = $arg;
    }

    if(@_ && exists $args{-key}) {
        my $arg = shift;
        if(ref $arg eq 'HASH') {
            %args = (%$arg, %args);
        }
    }
    if(! exists $args{-key}) {
        defined $args{'-reverse'} or croak("Expected option '-reverse' is undefined");
    }
    $self->[HASH_SEEK_ARG] = \%args;
    if($self->[HASH_ITER]) {
        $self->_reset_hash_iter;
    }
}

sub _reset_hash_iter {
    my $self = shift;
    if($self->[HASH_SEEK_ARG]) {
        my $iter = ($self->[HASH_SEEK_ARG]{'-reverse'} ? 'rev_iter' : 'iter');
        $self->[HASH_ITER] = $self->$iter($self->[HASH_SEEK_ARG]{'-key'});
    }
    else {
        $self->[HASH_ITER] = $self->iter;
    }
}

sub FIRSTKEY {
    my $self = shift;
    $self->_reset_hash_iter;

    my $node = $self->[HASH_ITER]->next
      or return;
    return $node->[_KEY];
}

sub NEXTKEY {
    my $self = shift;

    my $node = $self->[HASH_ITER]->next
      or return;
    return $node->[_KEY];
}

sub new {
    my ($class, $cmp) = @_;
    my $obj = [];
    $obj->[SIZE] = 0;
    if($cmp) {
        ref $cmp eq 'CODE'
          or croak('Invalid arg: codref expected');
        $obj->[CMP] = $cmp;
    }
    return bless $obj => $class;
}

*TIEHASH = \&new;

sub DESTROY { $_[0]->[ROOT]->DESTROY if $_[0]->[ROOT] }

sub CLEAR {
    my $self = shift;
    if($self->[ROOT]) {
        $self->[ROOT]->DESTROY;
        undef $self->[ROOT];
        undef $self->[HASH_ITER];
        $self->[SIZE] = 0;
    }
}

sub UNTIE {
    my $self = shift;
    $self->DESTROY;
    undef @$self;
}

sub resort {
    my $self = $_[0];
    my $cmp  = $_[1];
    ref $cmp eq 'CODE'
      or croak sprintf(q[Arg of type coderef required; got %s], ref $cmp || 'undef');

    my $new_tree = __PACKAGE__->new($cmp);
    $self->[ROOT]->strip(sub { $new_tree->put($_[0]) });
    $new_tree->put(delete $self->[ROOT]);
    $_[0] = $new_tree;
}

sub root { $_[0]->[ROOT] }
sub size { $_[0]->[SIZE] }

*SCALAR = \&size;

sub min {
    my $self = shift;
    return undef unless $self->[ROOT];
    return $self->[ROOT]->min;
}

sub max {
    my $self = shift;
    return undef unless $self->[ROOT];
    return $self->[ROOT]->max;
}

sub lookup {
    my $self = shift;
    my $key  = shift;
    defined $key
      or croak("Can't use undefined value as key");
    my $mode = shift || LUEQUAL;
    my $cmp = $self->[CMP];

    my $y;
    my $x = $self->[ROOT]
      or return;
    my $next_child;
    while($x) {
        $y = $x;
        if($cmp ? $cmp->($key, $x->[_KEY]) == 0
                : $key eq $x->[_KEY]) {
            # found it!
            if($mode == LUGREAT || $mode == LUNEXT) {
                $x = $x->successor;
            }
            elsif($mode == LULESS || $mode == LUPREV) {
                $x = $x->predecessor;
            }
            return wantarray
              ? ($x->[_VAL], $x)
              : $x->[_VAL];
        }
        if($cmp ? $cmp->($key, $x->[_KEY]) < 0
                : $key lt $x->[_KEY]) {
            $next_child = _LEFT;
        }
        else {
            $next_child = _RIGHT;
        }
        $x = $x->[$next_child];
    }
    # Didn't find it :(
    if($mode == LUGTEQ || $mode == LUGREAT) {
        if($next_child == _LEFT) {
            return wantarray ? ($y->[_VAL], $y) : $y->[_VAL];
        }
        else {
            my $next = $y->successor
              or return;
            return wantarray ? ($next->[_VAL], $next) : $next->[_VAL];
        }
    }
    elsif($mode == LULTEQ || $mode == LULESS) {
        if($next_child == _RIGHT) {
            return wantarray ? ($y->[_VAL], $y) : $y->[_VAL];
        }
        else {
            my $next = $y->predecessor
              or return;
            return wantarray ? ($next->[_VAL], $next) : $next->[_VAL];
        }
    }
    return;
}

*FETCH = \&lookup;
*get   = \&lookup;

sub nth {
    my ($self, $i) = @_;

    $i =~ /^-?\d+$/
      or croak('Integer index expected');
    if ($i < 0) {
        $i += $self->[SIZE];
    }
    if ($i < 0 || $i >= $self->[SIZE]) {
        return;
    }

    my ($node, $next, $moves);
    if ($i > $self->[SIZE] / 2) {
        $node = $self->max;
        $next = 'predecessor';
        $moves = $self->[SIZE] - $i - 1;
    }
    else {
        $node = $self->min;
        $next = 'successor';
        $moves = $i;
    }

    my $count = 0;
    while ($count != $moves) {
        $node = $node->$next;
        ++$count;
    }
    return $node;
}

sub EXISTS {
    my $self = shift;
    my $key  = shift;
    return defined $self->lookup($key);
}

sub put {
    my $self = shift;
    my $key_or_node = shift;
    defined $key_or_node
      or croak("Can't use undefined value as key or node");
    my $val = shift;

    my $cmp = $self->[CMP];
    my $z = (ref $key_or_node eq 'Tree::RB::Node')
              ? $key_or_node
              : Tree::RB::Node->new($key_or_node => $val);

    my $y;
    my $x = $self->[ROOT];
    while($x) {
        $y = $x;
        # Handle case of inserting node with duplicate key.
        if($cmp ? $cmp->($z->[_KEY], $x->[_KEY]) == 0
                : $z->[_KEY] eq $x->[_KEY])
        {
            my $old_val = $x->[_VAL];
            $x->[_VAL] = $z->[_VAL];
            return $old_val;
        }

        if($cmp ? $cmp->($z->[_KEY], $x->[_KEY]) < 0
                : $z->[_KEY] lt $x->[_KEY])
        {
            $x = $x->[_LEFT];
        }
        else {
            $x = $x->[_RIGHT];
        }
    }
    # insert new node
    $z->[_PARENT] = $y;
    if(not defined $y) {
        $self->[ROOT] = $z;
    }
    else {
        if($cmp ? $cmp->($z->[_KEY], $y->[_KEY]) < 0
                : $z->[_KEY] lt $y->[_KEY])
        {
            $y->[_LEFT] = $z;
        }
        else {
            $y->[_RIGHT] = $z;
        }
    }
    $self->_fix_after_insertion($z);
    $self->[SIZE]++;
}

*STORE = \&put;

sub _fix_after_insertion {
    my $self = shift;
    my $x = shift or croak('Missing arg: node');

    $x->[_COLOR] = RED;
    while($x != $self->[ROOT] && $x->[_PARENT][_COLOR] == RED) {
        my ($child, $rotate1, $rotate2);
        if(($x->[_PARENT] || 0) == ($x->[_PARENT][_PARENT][_LEFT] || 0)) {
            ($child, $rotate1, $rotate2) = (_RIGHT, '_left_rotate', '_right_rotate');
        }
        else {
            ($child, $rotate1, $rotate2) = (_LEFT, '_right_rotate', '_left_rotate');
        }
        my $y = $x->[_PARENT][_PARENT][$child];

        if($y && $y->[_COLOR] == RED) {
            $x->[_PARENT][_COLOR] = BLACK;
            $y->[_COLOR] = BLACK;
            $x->[_PARENT][_PARENT][_COLOR] = RED;
            $x = $x->[_PARENT][_PARENT];
        }
        else {
            if($x == ($x->[_PARENT][$child] || 0)) {
                $x = $x->[_PARENT];
                $self->$rotate1($x);
            }
            $x->[_PARENT][_COLOR] = BLACK;
            $x->[_PARENT][_PARENT][_COLOR] = RED;
            $self->$rotate2($x->[_PARENT][_PARENT]);
        }
    }
    $self->[ROOT][_COLOR] = BLACK;
}

sub delete {
    my ($self, $key_or_node) = @_;
    defined $key_or_node
      or croak("Can't use undefined value as key or node");

    my $z = (ref $key_or_node eq 'Tree::RB::Node')
              ? $key_or_node
              : ($self->lookup($key_or_node))[1];
    return unless $z;

    my $y;
    if($z->[_LEFT] && $z->[_RIGHT]) {
        # (Notes kindly provided by Christopher Gurnee)
        # When deleting a node 'z' which has two children from a binary search tree, the
        # typical algorithm is to delete the successor node 'y' instead (which is
        # guaranteed to have at most one child), and then to overwrite the key/values of
        # node 'z' (which is still in the tree) with the key/values (which we don't want
        # to lose) from the now-deleted successor node 'y'.

        # Since we need to return the deleted item, it's not good enough to overwrite the
        # key/values of node 'z' with those of node 'y'. Instead we swap them so we can
        # return the deleted values.

        $y = $z->successor;
        ($z->[_KEY], $y->[_KEY]) = ($y->[_KEY], $z->[_KEY]);
        ($z->[_VAL], $y->[_VAL]) = ($y->[_VAL], $z->[_VAL]);
    }
    else {
        $y = $z;
    }

    # splice out $y
    my $x = $y->[_LEFT] || $y->[_RIGHT];
    if(defined $x) {
        $x->[_PARENT] = $y->[_PARENT];
        if(! defined $y->[_PARENT]) {
            $self->[ROOT] = $x;
        }
        elsif($y == $y->[_PARENT][_LEFT]) {
            $y->[_PARENT][_LEFT] = $x;
        }
        else {
            $y->[_PARENT][_RIGHT] = $x;
        }
        # Null out links so they are OK to use by _fix_after_deletion
        delete @{$y}[_PARENT, _LEFT, _RIGHT];

        # Fix replacement
        if($y->[_COLOR] == BLACK) {
            $self->_fix_after_deletion($x);
        }
    }
    elsif(! defined $y->[_PARENT]) {
        # return if we are the only node
        delete $self->[ROOT];
    }
    else {
        # No children. Use self as phantom replacement and unlink
        if($y->[_COLOR] == BLACK) {
            $self->_fix_after_deletion($y);
        }
        if(defined $y->[_PARENT]) {
            no warnings 'uninitialized';
            if($y == $y->[_PARENT][_LEFT]) {
                delete $y->[_PARENT][_LEFT];
            }
            elsif($y == $y->[_PARENT][_RIGHT]) {
                delete $y->[_PARENT][_RIGHT];
            }
            delete $y->[_PARENT];
        }
    }
    $self->[SIZE]--;
    return $y;
}

*DELETE = \&delete;

sub _fix_after_deletion {
    my $self = shift;
    my $x = shift or croak('Missing arg: node');

    while($x != $self->[ROOT] && color_of($x) == BLACK) {
        my ($child1, $child2, $rotate1, $rotate2);
        no warnings 'uninitialized';
        if($x == left_of(parent_of($x))) {
            ($child1,    $child2,   $rotate1,       $rotate2) =
            (\&right_of, \&left_of, '_left_rotate', '_right_rotate');
        }
        else {
            ($child1,   $child2,    $rotate1,        $rotate2) =
            (\&left_of, \&right_of, '_right_rotate', '_left_rotate');
        }
        use warnings;

        my $w = $child1->(parent_of($x));
        if(color_of($w) == RED) {
            set_color($w, BLACK);
            set_color(parent_of($x), RED);
            $self->$rotate1(parent_of($x));
            $w = right_of(parent_of($x));
        }
        if(color_of($child2->($w)) == BLACK &&
           color_of($child1->($w)) == BLACK) {
            set_color($w, RED);
            $x = parent_of($x);
        }
        else {
            if(color_of($child1->($w)) == BLACK) {
                set_color($child2->($w), BLACK);
                set_color($w, RED);
                $self->$rotate2($w);
                $w = $child1->(parent_of($x));
            }
            set_color($w, color_of(parent_of($x)));
            set_color(parent_of($x), BLACK);
            set_color($child1->($w), BLACK);
            $self->$rotate1(parent_of($x));
            $x = $self->[ROOT];
        }
    }
    set_color($x, BLACK);
}

sub _left_rotate {
    my $self = shift;
    my $x = shift or croak('Missing arg: node');

    my $y = $x->[_RIGHT]
      or return;
    $x->[_RIGHT] = $y->[_LEFT];
    if($y->[_LEFT]) {
        $y->[_LEFT]->[_PARENT] = $x;
    }
    $y->[_PARENT] = $x->[_PARENT];
    if(not defined $x->[_PARENT]) {
        $self->[ROOT] = $y;
    }
    else {
        $x == $x->[_PARENT]->[_LEFT]
          ? $x->[_PARENT]->[_LEFT]  = $y
          : $x->[_PARENT]->[_RIGHT] = $y;
    }
    $y->[_LEFT]   = $x;
    $x->[_PARENT] = $y;
}

sub _right_rotate {
    my $self = shift;
    my $y = shift or croak('Missing arg: node');

    my $x = $y->[_LEFT]
      or return;
    $y->[_LEFT] = $x->[_RIGHT];
    if($x->[_RIGHT]) {
        $x->[_RIGHT]->[_PARENT] = $y
    }
    $x->[_PARENT] = $y->[_PARENT];
    if(not defined $y->[_PARENT]) {
        $self->[ROOT] = $x;
    }
    else {
        $y == $y->[_PARENT]->[_RIGHT]
          ? $y->[_PARENT]->[_RIGHT] = $x
          : $y->[_PARENT]->[_LEFT]  = $x;
    }
    $x->[_RIGHT] = $y;
    $y->[_PARENT] = $x;
}

1; # Magic true value required at end of module
__END__