split_col

#!/usr/bin/perl

use strict;
use warnings;

use Getopt::Long;
use File::Basename;

use Errno qw(EMFILE ENFILE);

# --column
# --template
### --substr
# --header

# TODO:
#
# * Add option to remove the selected column from the output.
#
# * Add option to sanitize the column content, e.g. replace / by _,
# maybe also handle other filename-unfriendly chars.

my $template = undef;
my $prefix = undef;
my $suffix = undef;
my $template_substr = '{}';
my $header = 0;
my $column = 1;
my $delim = "\t";
my $headerline = undef;
my $help = 0;
my $current_template;
my $current_infile;
my $current_fields;

sub err_exit {
    my $fmt = shift;
    my $txt = sprintf($fmt, @_);
    $txt .= "\n" unless $txt =~ m/\n$/;
    printf STDERR $txt;
    exit(1);
}

sub print_usage {
    err_exit("usage: %s [--template=a{}b] [--column=<n>] files", $0);
}

sub print_help {
    print <<HELP;
$0 - split a file using column contents

Options:
  --column n      Use column n (counting 1-based), default 1
  --header        Treat the first line of each input file as a header line.
                  That line will be copied verbatim to each output file.
  --template str  Use str as the template for output files names. str
                  should contain at least one replacement string (see
                  below) - this will be replaced by the contents of
                  the appropriate column.

If --template is not given, the output from an input file foo.bar.ext
will use the template foo.bar_{}.ext. If the input filename contains
no ., the implicitly used template is filename_{}.

The only supported column delimiter is tab (\\t).

For example, suppose foo contains these lines:

    horse     4
    spider    8
    cow       4
    octopus   8
    human     2
    millipede 1000

Running

    $0 --column 2 foo

will produce four files, foo_2, foo_4, foo_8, foo_1000, containing 1,
2, 2, and 1 lines, respectively. Had the option "--template
'animals_with_{}_legs'" been given, the output files would have had
more descriptive names.

Sanitizing the column contents or the resulting file names is not
attempted. It is obviously dangerous if the column may contain
slashes, so this should be avoided, as should any character which one
wouldn't want to end up in a filename.

Use - to read from stdin (or just omit trailing file names). When
reading from stdin, the filename implicitly used for generating output
file names is STDIN, so it is recommended to use the --template option
in that case.

Apart from {}, there are a few other replacement strings available:

{.}  Current input file without extension:
     foo.jpg becomes foo
     subdir/foo.jpg becomes subdir/foo
     sub.dir/foo.jpg becomes sub.dir/foo
     sub.dir/bar remains sub.dir/bar

{..} Extension of current input file, including the dot.  If the
     current input file does not have an extension, this is empty. In
     the above four examples, we would have ".jpg", ".jpg", ".jpg",
     and "", respectively.

     {.}{..} is always equal to the current input file name.

{n}  Column number n. One can split based on multiple columns by simply
     using multiple occurrences of these (and then the --column option
     is redundant). For example

         split_col --template 'foo{}bar' --column 4

     is equivalent to

         split_col --template 'foo{4}bar'

     {} is simply a synonym for {c}, where c is the value of the
     --column option (and which defaults to 1).

The default template described above is actually {.}_{}{..}

HELP
    exit(0);
}

GetOptions("template=s" => \$template,
	   "header!"    => \$header,
	   "column=i"   => \$column,
	   "help|h"     => \$help)
    or print_usage();

if ($help) {
    print_help();
}

err_exit("column index must be positive") if ($column <= 0);
# $column is given 1-based, but we use 0-based indexing internally
$column--;

$template = '{.}_{}{..}' unless defined $template;

err_exit("at least one replacement string ({} or {n}) must occur in the template")
    if (!($template =~ m/\{(?:[1-9][0-9]*)?\}/));


sub compute_current_template {
    my ($base, $ext);

    $current_template = $template;

    if ($current_infile =~ m#^(.*)(\.[^/]*)$#) {
	$base = $1;
	$ext = $2;
    } else {
	$base = $current_infile;
	$ext = '';
    }
    $current_template =~ s/\Q{.}\E/$base/g;
    $current_template =~ s/\Q{..}\E/$ext/g;
}


my %outfiles = ();
my $listhead = {};
$listhead->{next} = $listhead->{prev} = $listhead;
my $wopen_count = 0; # number of files open for writing (or appending)
my $wopen_max = 1000; # Never exceed this number.

sub get_replacement {
    my ($repstring) = @_;
    my $c;

    # Repstring is either undef (for {}) or a positive integer (for {n}, for some n).
    $c = (defined $repstring) ? $repstring - 1 : $column;

    exists $current_fields->[$c]
	or die "not enough columns on line $. of $current_infile";

    return $current_fields->[$c];
}

sub compute_outfn {
    my $ret = $current_template;
    $ret =~ s/\{([1-9][0-9]*)?\}/get_replacement($1)/ge;
    return $ret;
}

sub do_input_file {
    my $infile = shift;
    my $h = $infile eq '-' ? \*STDIN : open_read($infile);
    $current_infile = ($infile eq '-') ? 'STDIN' : $infile;
    my ($line, $outfn, $outhandle);

    if ($header) {
	$headerline = <$h>;
    }

    compute_current_template();

    while (<$h>) {
	# Save a copy of the text and then chomp $_. We don't want the
	# terminating newline to show up in $x if we are splitting
	# based on the last column.
	$line = $_;
	chomp;
	$current_fields = [split /\Q$delim\E/o];
	$outfn = compute_outfn();
	$outhandle = get_output_handle($outfn);
	print $outhandle $line;
    }
    close($h);
}


# We work around EMFILE/ENFILE as follows:
#
# Files open for output are tracked in %outfiles. For a given filename
# $fn, $outfiles{$fn} is a ref to a hash with three members: next, prev, handle.
#
# ->{handle} is, if defined, an file handle opened for writing to
# $fn. The two other fields are used for maintaining a LRU list of
# open file handles; the use of this is explained below.
#
# When we read a line of input and figure out what file $fn that line
# should go to, we first try to look up $fn in %outfiles. If it exists
# and ->{handle} is defined, we move $outfiles{$fn} to the front of
# the LRU list and use the handle directly.
#
# If $outfiles{$fn} exists but ->{handle} is not defined, it means
# that the file has at some point been opened, but we needed to close
# it to free up a file descriptor. In that case, we will open the file
# again, but in append mode so as not to overwrite the already written
# contents.
#
# If $outfiles{$fn} does not exist, we have never encountered $fn
# before. So open $fn for writing and insert an appropriate anon href
# in %outfiles.
#
# In either case, the just (re)opened file is inserted in the front of
# the list of open files.
#
# Now, if (re)opening a file for writing (appending) fails, and errno
# (aka $! in perl) is ENFILE or EMFILE, we may be able to solve the
# problem by closing an open file and trying again. We always close
# the least recently used file, since real input has a tendency to be
# somewhat sorted. We keep trying to close file handles as long as we
# have any file open for output (in the case of ENFILE, some other
# process may have successfully opened a file after we closed one but
# before we got around to trying again). If we run out of files to
# close, or if we encounter an error other than E[MN]FILE, we die with
# an error message.
#

sub list_insert_head {
    my $href = shift;
    die unless (!defined $href->{prev} && !defined $href->{next});
    $href->{next} = $listhead->{next};
    $href->{next}->{prev} = $href;
    $href->{prev} = $listhead;
    $listhead->{next} = $href;
}
sub list_remove {
    my $href = shift;
    die if $href == $listhead;
    $href->{prev}->{next} = $href->{next};
    $href->{next}->{prev} = $href->{prev};
    $href->{prev} = $href->{next} = undef;
}
sub list_move_to_front {
    my $href = shift;
    return if $listhead->{next} == $href;
    list_remove($href);
    list_insert_head($href);
}
sub list_remove_last {
    my $href = $listhead->{prev};
    return undef if $href == $listhead;
    list_remove($href);
    return $href;
}

sub get_output_handle {
    my $fn = shift;
    my $href = exists $outfiles{$fn} ? $outfiles{$fn} : undef;
    if (defined $href && defined $href->{handle}) {
	list_move_to_front($href);
	return $href->{handle}
    }
    if (defined $href) {
	$href->{handle} = open_append($fn);
    }
    else {
	$href = {prev => undef, next => undef, handle => open_write($fn)};
	$outfiles{$fn} = $href;
    }
    list_insert_head($href);
    $wopen_count++;
    return $href->{handle};
}

sub open_read {
    my $fn = shift;
    return open_common($fn, '<', 'reading');
}
sub open_write {
    my $fn = shift;
    my $h = open_common($fn, '>', 'writing');
    print $h $headerline if defined $headerline;
    return $h;
}
sub open_append {
    my $fn = shift;
    return open_common($fn, '>>', 'appending');
}

sub open_common {
    my ($fn, $mode, $verb) = @_;
    my ($h, $r);
    close_lru_file() while ($wopen_count >= $wopen_max);
    while (!($r = open($h, $mode, $fn)) && ($! == EMFILE || $! == ENFILE)) {
	last if $wopen_count == 0;
	close_lru_files($! == EMFILE ? $wopen_count - 1 : int($wopen_count/2));
    }
    if (!$r) {
	die "unable to open ${fn} for ${verb}: $!";
    }
    return $h;
}

sub close_lru_file {
    if ($wopen_count <= 0) {
	die "BUG! ", (caller(0))[3], "() should not be called when wopen_count == $wopen_count";
    }
    my $href = list_remove_last();
    if (!defined $href) {
	die "BUG! ", (caller(0))[3], "() called while no files appear to be open for writing";
    }
    close($href->{handle}) or die "closing file handle failed: $!";
    $href->{handle} = undef;
    $wopen_count--;
}

sub close_lru_files {
    my $target = shift;
    close_lru_file() while ($wopen_count > $target);
}


push @ARGV, '-' if (@ARGV == 0);
for my $inp (@ARGV) {
    do_input_file($inp);
}

# Close all output files. We reuse close_lru_files() so that errors are reported.
close_lru_files(0);