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

my $total_size = 0;
my %roots;
my @all_nodes;
foreach my $line ( <STDIN> ) {
  chomp($line);
  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, $maxrss, $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 $parent->{$pc}) {
      $parent->{$pc}->{'file_name'} = $file_name;
      $parent->{$pc}->{'proc_name'} = $proc_name;
      $parent->{$pc}->{'off'} = $off;
      $parent->{$pc}->{'pc'} = $pc;
      $parent->{$pc}->{'relpc'} = $relpc;

      push(@all_nodes, $parent->{$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'};
  }
}

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("digraph \"memlog %s\" {\n", format_bytes($total_size));
print("size=\"8,11\";\n");
print("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 !~ /^\?/) {
    $func =~ s/.*\.//;
    $func = `c++filt '$func'`;
    chomp($func);

    $ret .= $func . '\n' . $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 .= basename($node->{'file_name'});
  if ($node->{'proc_name'} ne '?') {
    $ret .= ' (' . $node->{'proc_name'} . '+0x' . $node->{'off'} . ')';
  }

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

my $skip_frac = 0.01;
my %skipped;

foreach my $node (@all_nodes) {
  my $name = get_name($node);

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

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

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

foreach my $node (@all_nodes) {
  my $local_size = $node->{'size'};
  if ($skipped{$node->{'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)", 3.0 * sqrt($frac));

    printf("N%s -> N%s [label=\"%s\", weight=%d, style=\"%s\"];\n",
      $node->{'pc'}, $cpc, format_bytes($child_size), $weight, $style);
  }
}

print("}\n");