#! /usr/bin/env perl
use strict;
use File::Basename;

my $memlog_fn = $ARGV[0];
my $out_dir = $ARGV[1] || '.';

my $print_raw_proc_name = 0;

if (! -f $memlog_fn) {
  print "Usage: $0 <memlog file> [<output directory>]\n";
  exit 1;
}

open(MEMLOG, $memlog_fn) || die "Can't open $memlog_fn: $!";

# The first step is to determine the high-water mark.
my $max_rss = 0;
foreach my $line (<MEMLOG>) {
  chomp($line);
  my @parts = split(/\t/, $line);

  my $op = shift(@parts);
  my $state = shift(@parts);

  my ($time, $then_max_rss, $tid) = split(/\s+/, $state);
  if ($max_rss < $then_max_rss) {
    $max_rss = $then_max_rss;
  }
}

seek(MEMLOG, 0, 0);

# Scan the log for malloc/free pairings. We're interested only in active
# allocations at the time when the rss reaches the final maxrss.
my $max_rss_time = 0;
my %malloc_lines;
foreach my $line (<MEMLOG>) {
  chomp($line);
  my @parts = split(/\t/, $line);

  my $op = shift(@parts);
  my $state = shift(@parts);

  if ($op =~ /^M:/) {
    my ($size, $ptr) = ($op =~ /^M: (\d+) 0x(\w+)/);
    $malloc_lines{$ptr} = $line;
  } elsif ($op =~ /^F:/) {
    my ($ptr) = ($op =~ /^F: 0x(\w+)/);
    delete $malloc_lines{$ptr};
  } else {
    next;
  }

  # If we've reached the max rss, we've seen all we need to see.
  my ($time, $then_max_rss, $tid) = split(/\s+/, $state);
  $max_rss_time = $time;
  if ($then_max_rss == $max_rss) {
    last;
  }
}

close(MEMLOG);

# Convert maxrss, currently in KB, to bytes.
$max_rss *= 1024;

my $total_size = 0;
my %roots;
my %all_nodes;
foreach my $line (values %malloc_lines) {
  my @parts = split(/\t/, $line);

  my $op = shift(@parts);
  my $state = shift(@parts);

  # Only dealing with allocations here...
  if ($op !~ /^M:/) {
    next;
  }

  my ($size, $ptr) = ($op =~ /^M: (\d+) 0x(\w+)/);
  my ($time, $then_max_rss, $tid) = split(/\s+/, $state);

  $total_size += $size;

  sub level_parts($) {
    my $level = @_[0];
    my ($file_name, $proc_name, $off, $pc, $relpc) =
      ($level =~ /^(.*) \((.*)\+0x(\w+)\) \[0x(\w+) \(0x(\w+)\)\]/);

    return ($file_name, $proc_name, $off, $pc, $relpc);
  }

  # Put the top of the stack first.
  @parts = reverse(@parts);

  my $parent = \%roots;
  for (my $i = 0; $i < scalar(@parts); ++$i) {
    my $level = $parts[$i];
    my ($file_name, $proc_name, $off, $pc, $relpc) = level_parts($level);

    # Skip this level if we don't even know from what file it came.
    if ($file_name eq '?') {
      next;
    }

    # print STDERR "parsed: $file_name, $proc_name, $off, $pc, $relpc\n";

    if (!exists $all_nodes{$pc}) {
      $all_nodes{$pc}->{'file_name'} = $file_name;
      $all_nodes{$pc}->{'proc_name'} = $proc_name;
      $all_nodes{$pc}->{'off'} = $off;
      $all_nodes{$pc}->{'pc'} = $pc;
      $all_nodes{$pc}->{'relpc'} = $relpc;
    }

    if (!exists $parent->{$pc}) {
      $parent->{$pc} = $all_nodes{$pc};
    }

    $parent->{$pc}->{'size'} += $size;

    my ($next_file_name, $next_proc_name, $next_off, $next_pc, $next_relpc);
    if ($i < scalar(@parts)-1) {
     my $next_level = $parts[$i+1];
     ($next_file_name, $next_proc_name, $next_off, $next_pc, $next_relpc) =
       level_parts($next_level);
      $parent->{$pc}->{'child_sizes'}->{$next_pc} += $size;
    }

    if (!exists $parent->{'children'}) {
      $parent->{'children'} = {};
    }

    $parent = $parent->{'children'};
  }
}

my $dot_fn = "$out_dir/" . basename($memlog_fn) . ".dot";
my $ps_fn = "$out_dir/" . basename($memlog_fn) . ".ps";
my $pdf_fn = "$out_dir/" . basename($memlog_fn) . ".pdf";

open(DOT, ">$dot_fn") || die "Can't open $dot_fn: $!";

sub format_bytes($) {
  my @sizes = qw( B KB MB GB TB PB );
  my $size = $_[0];

  my $i = 0;
  while ($size > 1024) {
    $size /= 1024;
    ++$i;
  }

  return sprintf("%.3f $sizes[$i]", $size);
}

printf DOT ("digraph \"memlog %s (maxrss = %s after %s s)\" {\n",
            format_bytes($total_size), format_bytes($max_rss), $max_rss_time);
print DOT ("size=\"8,11\";\n");
print DOT ("node [width=0.375,height=0.25];\n");

my %cached_names;
sub get_name($) {
  my $node = $_[0];
  my $pc = $node->{'pc'};

  if (exists $cached_names{$pc}) {
    return $cached_names{$pc};
  }

  my $ret = '';

  # Prefer the relative offset (that is what we want for shared libraries), but
  # if is not available, use the full offset (which is what we want for the
  # base executable).
  my $exe_off = $node->{'relpc'};
  if (!$exe_off) {
    $exe_off = $pc;
  }

  my $file_name = $node->{'file_name'};
  my ($func, $loc) = `addr2line -e $file_name -f 0x$exe_off`;
  chomp($func);
  chomp($loc);

  if ($func !~ /^\?/) {
    # In general, this function name might look something like:
    #   00000329.plt_call.wcsnrtombs@@GLIBC_2.3+0
    $func =~ s/@.*//;
    $func =~ s/.*\.//;
    $func = `c++filt '$func'`;
    chomp($func);

    $ret .= $func . '\n';

    if ($loc !~ /^\?/) {
      $ret .= $loc . '\n';
    }
  } elsif ($node->{'proc_name'} ne '?') {
    my $proc_name = $node->{'proc_name'};
    $proc_name = `c++filt '$proc_name'`;
    chomp($proc_name);

    $ret .= $proc_name . '\n';
  }

  $ret .= $node->{'file_name'};
  if ($print_raw_proc_name and $node->{'proc_name'} ne '?') {
    $ret .= '\n' . $node->{'proc_name'} . '+0x' . $node->{'off'};
  }

  $cached_names{$pc} = $ret;
  return $ret;
}

my $skip_frac = 0.01;
my %skipped;

foreach my $pc (keys %all_nodes) {
  my $node = $all_nodes{$pc};
  my $name = get_name($node);

  my $local_size = $node->{'size'};
  if ($local_size * 1.0 / $total_size < $skip_frac) {
    $skipped{$pc} = 1;
    next;
  }

  my $fs = 8.0;
  if ($local_size > 0) {
    $fs = 50.0 * (abs($local_size * 1.0 / $total_size))**0.125;
  }

  printf DOT ("N%s [label=\"%s\\n%s\", shape=box, fontsize=%.1f%s];\n",
    $pc, $name, format_bytes($local_size), $fs);
}

foreach my $pc (keys %all_nodes) {
  my $node = $all_nodes{$pc};

  my $local_size = $node->{'size'};
  if ($skipped{$pc}) {
    next;
  }

  foreach my $cpc (keys %{$node->{'child_sizes'}}) {
    if ($skipped{$cpc}) {
      next;
    }

    my $child_size = $node->{'child_sizes'}->{$cpc};
    my $frac = $child_size * 1.0 / $local_size;

    my $weight = 100.0 * sqrt($frac);
    my $style = sprintf("setlinewidth(%f)", 8.0 * sqrt($frac));

    my $fs = 40.0 * $frac**0.125;

    printf DOT ("N%s -> N%s [label=\"%s\", weight=%d, style=\"%s\", fontsize=%.1f];\n",
      $pc, $cpc, format_bytes($child_size), $weight, $style, $fs);
  }
}

print DOT ("}\n");

close(DOT);

system("dot -Tps2 < '$dot_fn' > '$ps_fn'");
system("ps2pdf '$ps_fn' '$pdf_fn'");

exit 0;