File Coverage

blib/lib/Compress/BGZF/Writer.pm

Criterion	Covered	Total	%
statement	119	119	100.0
branch	27	40	67.5
condition	5	5	100.0
subroutine	23	23	100.0
pod	7	7	100.0
total	181	194	93.3

package Compress::BGZF::Writer;

use strict;
use warnings;

use Carp;
use Compress::Zlib;
use IO::Compress::RawDeflate qw/rawdeflate $RawDeflateError/;

use constant HEAD_BYTES  => 18;
use constant FOOT_BYTES  => 8;
use constant FLUSH_SIZE  => 2**16 - HEAD_BYTES - FOOT_BYTES - 1;
use constant BGZF_HEADER => pack "H*", '1f8b08040000000000ff060042430200';
use constant BGZF_EOF    => pack "H*",
    '1f8b08040000000000ff0600424302001b0003000000000000000000';

## no critic
# allow for filehandle tie'ing
sub TIEHANDLE { Compress::BGZF::Writer::new(@_) }
sub PRINT     { Compress::BGZF::Writer::_queue(@_) }
sub CLOSE     { Compress::BGZF::Writer::finalize(@_) }
## use critic

sub new_filehandle {

    #-------------------------------------------------------------------------
    # ARG 0 : (optional) output filename
    #-------------------------------------------------------------------------
    # RET 0 : filehandle GLOB
    #-------------------------------------------------------------------------

    my ($class, $fn_out) = @_;

    open my $fh, '<', undef;
    tie *$fh, $class, $fn_out or croak "failed to tie filehandle";
    return $fh;

}

sub new {

    #-------------------------------------------------------------------------
    # ARG 0 : (optional) output filename
    #-------------------------------------------------------------------------
    # RET 0 : Compress::BGZF::Writer object
    #-------------------------------------------------------------------------

    my ($class, $fn_out) = @_;
    my $self = bless {}, $class;

    # initialize
    if (defined $fn_out) {
        open $self->{fh}, '>', $fn_out
            or croak "Error opening file for writing";
    }
    else {
        $self->{fh} = \*STDOUT;
    }
    binmode $self->{fh};

    $self->{c_level} = Z_DEFAULT_COMPRESSION;
    $self->{buffer}  = ''; # contents waiting to be compressed/written

    # these variables are tracked to allow for virtual offset calculation
    $self->{block_offset}  = 0;  # offset of current block in compressed data
    $self->{buffer_offset} = 0;  # offset of current pos in uncompressed block

    # these variables are tracked to allow for index creation
    $self->{u_offset} = 0; #uncompressed file offset
    $self->{idx} = [];
    $self->{write_eof} = 0;

    return $self;

}

sub set_level {

    #-------------------------------------------------------------------------
    # ARG 0 : compression level desired
    #-------------------------------------------------------------------------
    # no returns
    #-------------------------------------------------------------------------

    my ($self, $level) = @_;

    croak "Invalid compression level (allowed 0-9)"
        if ($level !~ /^\d$/);
    $self->{c_level} = $level;

    return;

}

sub set_write_eof {

    # Sets whether to include htslib-style EOF empty block at end of file

    #-------------------------------------------------------------------------
    # ARG 0 : (optional) boolean
    #-------------------------------------------------------------------------
    # no returns
    #-------------------------------------------------------------------------

    my ($self, $bool) = @_;

    $bool //= 1;
    $self->{write_eof} = $bool ? 1 : 0;

    return;

}

sub add_data {

    # a wrapper around the queue() function that returns the virtual offset
    # of the chunk added

    #-------------------------------------------------------------------------
    # ARG 0 : data chunk to queue for compression
    #-------------------------------------------------------------------------
    # RET 1 : virtual offset of data written
    #-------------------------------------------------------------------------

    my ($self, $content) = @_;

    my $vo =  ($self->{block_offset} << 16) | $self->{buffer_offset};
    $self->_queue( $content );

    return $vo;

}

sub _queue {

    #-------------------------------------------------------------------------
    # ARG 0 : data chunk to queue for compression
    #-------------------------------------------------------------------------
    # no returns
    #-------------------------------------------------------------------------

    my ($self, $content) = @_;

    $self->{buffer} .= $content;

    # compress/write in 64k chunks
    while (length($self->{buffer}) >= FLUSH_SIZE) {

        my $chunk = substr $self->{buffer}, 0, FLUSH_SIZE, '';
        my $unwritten = $self->_write_block($chunk);
        $self->{buffer} = $unwritten . $self->{buffer}
            if ( length($unwritten) );

    }
    $self->{buffer_offset} = length $self->{buffer};

    return;

}

sub _write_block {

    #-------------------------------------------------------------------------
    # ARG 0 : independent data block to compress
    #-------------------------------------------------------------------------
    # RET 0 : remaining data that wasn't written
    #-------------------------------------------------------------------------

    my ($self, $chunk) = @_;

    my $chunk_len = length($chunk);

    # payload is compressed with DEFLATE
    rawdeflate(\$chunk, \my $payload, -Level => $self->{c_level})
        or croak "deflate failed: $RawDeflateError\n";

    # very rarely, a DEFLATEd string may be larger than input. This can result
    # in a block size > 2**16, which violates the BGZF specification and
    # causes corruption of the BC field. Fix those edge cases here (somewhat
    # slow but shouldn't happen often) and send the rest back to the buffer
    my $trimmed = '';
    while (length($payload) > FLUSH_SIZE) {
        my $trim_len = int( $chunk_len * 0.05 );
        $trimmed = substr($chunk, -$trim_len, $trim_len, '') . $trimmed;
        rawdeflate(\$chunk, \$payload, -Level => $self->{c_level})
            or croak "deflate failed: $RawDeflateError\n";
        $chunk_len = length($chunk);
    }

    my $block_size = length($payload) + HEAD_BYTES + FOOT_BYTES;

    croak "Internal error: block size > 65536" if ($block_size > 2**16);

    # payload is wrapped with appropriate headers and footers
    print { $self->{fh} } pack(
        "a*va*VV",
        BGZF_HEADER,
        $block_size - 1,
        $payload,
        crc32($chunk),
        $chunk_len,
    ) or croak "Error writing compressed block";

    # increment the current offsets
    $self->{block_offset} += $block_size;
    $self->{u_offset}     += $chunk_len;
    push @{ $self->{idx} }, [$self->{block_offset}, $self->{u_offset}];

    return $trimmed;

}

sub finalize {

    #-------------------------------------------------------------------------
    # no arguments
    #-------------------------------------------------------------------------
    # no returns
    #-------------------------------------------------------------------------

    my ($self) = @_;

    while (length($self->{buffer}) > 0) {

        croak "file closed but buffer not empty"
            if ( ! defined fileno($self->{fh}) );

        my $chunk = substr $self->{buffer}, 0, FLUSH_SIZE, '';
        my $unwritten = $self->_write_block($chunk);
        $self->{buffer} = $unwritten . $self->{buffer}
            if ( length($unwritten) );

    }
    # write EOF block if requested (only first time finalize() is run)
    if ($self->{write_eof} && defined fileno($self->{fh})) {
        print { $self->{fh} } BGZF_EOF;
    }
    if (defined fileno($self->{fh}) ) {
        close $self->{fh}
            or croak "Error closing compressed file";
    }

    return;

}

sub write_index {

    #-------------------------------------------------------------------------
    # ARG 0 : index output filename
    #-------------------------------------------------------------------------
    # No returns
    #-------------------------------------------------------------------------

    my ($self, $fn_out) = @_;

    $self->finalize(); # always clear remaining buffer to fully populate index
    croak "missing index output filename" if (! defined $fn_out);
    open my $fh_out, '>:raw', $fn_out
        or croak "Error opening index file for writing";

    my @offsets = @{ $self->{idx} };
    pop @offsets; # last offset is EOF
    print {$fh_out} pack('Q<', scalar(@offsets))
        or croak "Error printing to index file";
    for (@offsets) {
        print {$fh_out} pack('Q
            or croak "Error printing offset to index file";
    }

    close $fh_out
        or croak "Error closing index file after writing";
    return;

}

sub DESTROY {

    my ($self) = @_;

    # make sure we call finalize in case the caller forgot
    $self->finalize();

    return;

}

1;


__END__