#!/usr/bin/perl -w
# Copyright (C) 2018 The Android Open Source Project
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#  * Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in
#    the documentation and/or other materials provided with the
#    distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
# OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
# AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.

use strict;

sub PrintHeader() {
  print <<EOT;
/*
 * Copyright (C) 2018 The Android Open Source Project
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *  * Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 *  * Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in
 *    the documentation and/or other materials provided with the
 *    distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
 * OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
 * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

// Generated by gen_malloc.pl, do not modify.

EOT
}

sub PrintMainloop() {
  print <<EOT;
void BenchmarkMalloc(MallocEntry entries[], size_t total_entries, size_t max_allocs) {
  void* ptrs[max_allocs];

  for (size_t i = 0; i < total_entries; i++) {
    switch (entries[i].type) {
    case MALLOC:
      ptrs[entries[i].idx] = malloc(entries[i].size);
      // Touch at least one byte of the allocation to make sure that
      // PSS for this allocation is counted.
      reinterpret_cast<uint8_t*>(ptrs[entries[i].idx])[0] = 10;
      break;
    case CALLOC:
      ptrs[entries[i].idx] = calloc(entries[i].arg2, entries[i].size);
      // Touch at least one byte of the allocation to make sure that
      // PSS for this allocation is counted.
      reinterpret_cast<uint8_t*>(ptrs[entries[i].idx])[0] = 20;
      break;
    case MEMALIGN:
      ptrs[entries[i].idx] = memalign(entries[i].arg2, entries[i].size);
      // Touch at least one byte of the allocation to make sure that
      // PSS for this allocation is counted.
      reinterpret_cast<uint8_t*>(ptrs[entries[i].idx])[0] = 30;
      break;
    case REALLOC:
      if (entries[i].arg2 == 0) {
        ptrs[entries[i].idx] = realloc(nullptr, entries[i].size);
      } else {
        ptrs[entries[i].idx] = realloc(ptrs[entries[i].arg2 - 1], entries[i].size);
      }
      // Touch at least one byte of the allocation to make sure that
      // PSS for this allocation is counted.
      reinterpret_cast<uint8_t*>(ptrs[entries[i].idx])[0] = 40;
      break;
    case FREE:
      free(ptrs[entries[i].idx]);
      break;
    }
  }
}

EOT
}

sub PrintDefinitions() {
  print <<EOT;
enum AllocEnum : uint8_t {
  MALLOC = 0,
  CALLOC,
  MEMALIGN,
  REALLOC,
  FREE,
};

struct MallocEntry {
  AllocEnum type;
  size_t idx;
  size_t size;
  size_t arg2;
};

EOT
}

sub PrintUsageAndExit() {
  print "USAGE: gen_malloc.pl [-d][-i][-m] THREAD_ID STRUCT_NAME MAX_SLOT_NAME < ALLOCS.txt\n";
  print "  -d\n";
  print "    Print the structure definitions.\n";
  print "  -i\n";
  print "    Ignore missing allocations.\n";
  print "  -m\n";
  print "    Print the main loop code that can reproduce the trace.\n";
  print "  THREAD_ID\n";
  print "    The thread for which entries will be printed.\n";
  print "  STRUCT_NAME\n";
  print "    The name of the structure containing all of the entries.\n";
  print "  MAX_SLOT_NAME\n";
  print "    The name of the name of the maximum slots variable.\n";
  print "  ALLOCS.txt\n";
  print "    A file generated by the malloc debug option record_allocs\n";
  exit(1);
}

sub GetSlot($) {
  my ($opts) = @_;

  if (scalar(@{$opts->{empty_slots}}) == 0) {
    return $opts->{last_slot}++;
  } else {
    return pop(@{$opts->{empty_slots}});
  }
}

sub PrintFreeSlots($) {
  my ($opts) = @_;

  if (scalar(@{$opts->{empty_slots}}) == $opts->{last_slot}) {
    return;
  }

  print "\n    // Free rest of the allocs.\n";
  my @sorted_empty_slots = sort({$a <=> $b} @{$opts->{empty_slots}});
  my $slot = 0;
  my $last_slot = $opts->{last_slot};
  while ($slot < $last_slot) {
    my $empty_slot = $last_slot;
    if (scalar(@sorted_empty_slots) != 0) {
      $empty_slot = shift(@sorted_empty_slots);
    }
    for (; $slot < $empty_slot; $slot++) {
      print "  {FREE, $slot, 0, 0},\n";
    }
    $slot++;
  }
}

sub PrintAlloc($$$$$$) {
  my ($opts, $cur_thread, $pointer, $name, $size, $arg2) = @_;

  if ($opts->{thread} eq $cur_thread) {
    my $slot = GetSlot($opts);
    $opts->{pointers}->{$pointer} = $slot;
    print "  {$name, $slot, $size, $arg2},\n";
  } else {
    $opts->{pointers}->{$pointer} = -1;
  }
}

sub PrintEntries($$) {
  my ($thread, $ignore_missing_allocations) = @_;

  my $opts = {};
  $opts->{thread} = $thread;
  $opts->{empty_slots} = [];
  $opts->{last_slot} = 0;
  $opts->{pointers} = {};

  while (<>) {
    if (!/^(\d+):\s*/) {
      continue
    }
    my $cur_thread = $1;

    $_ = $';
    if (/^malloc\s+(\S+)\s+(\d+)/) {
      my $pointer = $1;
      my $size = $2;
      PrintAlloc($opts, $cur_thread, $pointer, "MALLOC", $size, 0);
    } elsif (/^calloc\s+(\S+)\s+(\d+)\s+(\d+)/) {
      my $pointer = $1;
      my $nmemb = $2;
      my $size = $3;
      PrintAlloc($opts, $cur_thread, $pointer, "CALLOC", $size, $nmemb);
    } elsif (/^memalign\s+(\S+)\s+(\d+)\s+(\d+)/) {
      my $pointer = $1;
      my $align = $2;
      my $size = $3;
      PrintAlloc($opts, $cur_thread, $pointer, "MEMALIGN", $size, $align);
    } elsif (/^free\s+(\S+)/) {
      my $pointer = $1;
      if (!exists $opts->{pointers}->{$pointer}) {
        if ($ignore_missing_allocations) {
          warn "WARNING: $.: Unknown allocation $pointer ignored on $cur_thread\n";
          next;
        } else {
          die "$.: Unknown allocation $pointer on $cur_thread\n";
        }
      } elsif ($opts->{pointers}->{$pointer} != -1) {
        print "  {FREE, $opts->{pointers}->{$pointer}, 0, 0},\n";
        push @{$opts->{empty_slots}}, $opts->{pointers}->{$pointer};
      }
    } elsif (/^realloc\s+(\S+)\s+(\S+)\s+(\d+)/) {
      my $new_pointer = $1;
      my $old_pointer = $2;
      my $size = $3;

      if ($thread ne $cur_thread) {
        if ($new_pointer ne $old_pointer) {
          $opts->{pointers}->{$new_pointer} = -1;
          delete $opts->{pointers}->{$old_pointer};
        }
      } elsif ($old_pointer eq "0x0") {
        my $slot = GetSlot($opts);
        # This was a realloc(nullptr, size) call.
        print "  {REALLOC, $slot, $size, 0},\n";
        $opts->{pointers}->{$new_pointer} = $slot;
      } else {
        if (!exists $opts->{pointers}->{$old_pointer}) {
          if ($ignore_missing_allocations) {
            warn "WARNING: $.: Unknown realloc allocation $old_pointer ignored on $cur_thread\n";
            next;
          } else {
            die "Unknown realloc allocation $old_pointer on $cur_thread\n";
          }
        }

        if ($opts->{pointers}->{$old_pointer} != -1) {
          # Reuse the same slot, no need to get a new one.
          my $slot = $opts->{pointers}->{$old_pointer};
          printf("    {REALLOC, $slot, $size, %d},\n", $slot + 1);

          # NOTE: It is possible that old pointer and new pointer are the
          # same (a realloc returns the same pointer).
          if ($new_pointer ne $old_pointer) {
            $opts->{pointers}->{$new_pointer} = $slot;
            delete $opts->{pointers}->{$old_pointer};
          }
        }
      }
    } elsif (!/^thread_done/) {
      die "$.: Unknown line $_\n";
    }
  }

  PrintFreeSlots($opts);

  return $opts->{last_slot};
}

sub ProcessArgs($) {
  my ($opts) = @_;

  $opts->{print_definitions} = 0;
  $opts->{ignore_missing_allocations} = 0;
  $opts->{print_mainloop} = 0;
  my @args = ();
  while (scalar(@ARGV)) {
    my $arg = pop(@ARGV);
    if ($arg =~ /^-/) {
      if ($arg eq "-d") {
        $opts->{print_definitions} = 1;
      } elsif ($arg eq "-i") {
        $opts->{ignore_missing_allocations} = 1;
      } elsif ($arg eq "-m") {
        $opts->{print_mainloop} = 1;
      } else {
        print "Unknown option $arg\n";
        PrintUsageAndExit();
      }
    } else {
      unshift @args, $arg;
    }
  }

  return @args;
}

my $opts = {};
my @args = ProcessArgs($opts);
if (scalar(@args) != 3) {
  PrintUsageAndExit();
}

my $thread = $args[0];
my $struct_name = $args[1];
my $max_slot_name = $args[2];

PrintHeader();
if ($opts->{print_definitions}) {
  PrintDefinitions();
}
if ($opts->{print_mainloop}) {
  PrintMainloop();
}

print "static MallocEntry ${struct_name}[] = {\n";
my $total_slots = PrintEntries($thread, $opts->{ignore_missing_allocations});
print "};\n";
print "static constexpr size_t ${max_slot_name} = $total_slots;\n";