#!/usr/bin/env python3
# Copyright 2022 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Maps the superpage occupancy, and the allocated object sizes on the heap.

To run this:

1. Build pa_dump_heap at a close-enough revision to the Chrome instance you're
   interested in.
2. Run pa_dump_heap --pid=PID_OF_RUNNING_CHROME_PROCESS --json=output.json
3. Run plot_superpages.py --json=output.json --output=output.png
"""

import argparse
import bisect
import math
import json

import matplotlib
from matplotlib import pylab as plt
import numpy as np

# Valid for the "before slot" variant.
_BRP_OVERHEAD = 4


def ParseJson(filename: str) -> dict:
  with open(filename, 'r') as f:
    data = json.load(f)
    return data


def PlotSuperPageData(data: dict, output_filename: str):
  num_superpages = len(data)
  num_partition_pages = len(data[0]['partition_pages'])
  data_np = np.zeros((num_superpages, num_partition_pages, 3))

  address_to_superpage = {superpage['address']: superpage for superpage in data}
  addresses = sorted(address_to_superpage.keys())

  BLACK = (0.0, 0.0, 0.0)
  GRAY = (0.5, 0.5, 0.5)
  RED = (1.0, 0.0, 0.0)
  GREEN = (0.0, 0.8, 0.0)
  WHITE = (1.0, 1.0, 1.0)
  cmap = matplotlib.cm.get_cmap('coolwarm')

  for superpage_index in range(len(addresses)):
    superpage = address_to_superpage[addresses[superpage_index]]
    for index, partition_page in enumerate(superpage['partition_pages']):
      value = None
      if partition_page['type'] == 'metadata':
        value = BLACK
      elif partition_page['type'] == 'guard':
        value = GRAY
      elif partition_page['all_zeros'] and not 'is_active' in partition_page:
        # Otherwise it may be the subsequent partition page of a decommitted
        # slot span.
        if partition_page['page_index_in_span'] == 0:
          value = WHITE

      if value is not None:
        data_np[superpage_index, index] = value
        continue

      assert partition_page['type'] == 'payload'
      if partition_page['page_index_in_span'] == 0:
        num_partition_pages_in_slot_span = math.ceil(
          partition_page['num_system_pages_per_slot_span'] / 4
        )
        if partition_page['is_decommitted'] or partition_page['is_empty']:
          value = GREEN
        else:
          fullness = (
            partition_page['num_allocated_slots']
            / partition_page['slots_per_span']
          )
          value = cmap(fullness)[:3]
        data_np[
          superpage_index, index : index + num_partition_pages_in_slot_span, :
        ] = value

  plt.figure(figsize=(20, len(address_to_superpage) / 2))
  plt.imshow(data_np)
  plt.title('Super page map')
  plt.yticks(ticks=range(len(addresses)), labels=addresses)
  plt.xlabel('PartitionPage index')

  handles = [
    matplotlib.patches.Patch(facecolor=BLACK, edgecolor='k', label='Metadata'),
    matplotlib.patches.Patch(facecolor=GRAY, edgecolor='k', label='Guard'),
    matplotlib.patches.Patch(facecolor=WHITE, edgecolor='k', label='Empty'),
    matplotlib.patches.Patch(
      facecolor=GREEN, edgecolor='k', label='Decommitted'
    ),
    matplotlib.patches.Patch(
      facecolor=cmap(0.0), edgecolor='k', label='Committed Empty'
    ),
    matplotlib.patches.Patch(
      facecolor=cmap(1.0), edgecolor='k', label='Committed Full'
    ),
  ]
  plt.legend(handles=handles, loc='lower left', fontsize=12, framealpha=0.7)

  plt.savefig(output_filename, bbox_inches='tight')


def PlotSuperPageCompressionData(data: dict, output_filename: str):
  num_superpages = len(data)
  num_pages = len(data[0]['page_sizes'])
  data_np = np.zeros((num_superpages, num_pages, 3))

  address_to_superpage = {superpage['address']: superpage for superpage in data}
  addresses = sorted(address_to_superpage.keys())

  WHITE = (1.0, 1.0, 1.0)
  cmap = matplotlib.cm.get_cmap('coolwarm')

  total_compressed_size = 0
  total_uncompressed_size = 0
  for superpage_index in range(len(addresses)):
    superpage = address_to_superpage[addresses[superpage_index]]
    for index, page in enumerate(superpage['page_sizes']):
      total_uncompressed_size += page['uncompressed']
      total_compressed_size += page['compressed']
      if page['uncompressed'] == 0:
        value = WHITE
      else:
        value = cmap(page['compressed'] / page['uncompressed'])[:3]
      data_np[superpage_index, index] = value

  overall_compression_ratio = (
    (100.0 * total_compressed_size / (total_uncompressed_size))
    if total_uncompressed_size
    else 0
  )
  plt.figure(figsize=(20, len(address_to_superpage)))
  plt.imshow(data_np, aspect=8)
  plt.title(
    'Super page compression ratio map - Ratio = %.1f%% '
    '- Compressed Size = %.1fMiB'
    % (overall_compression_ratio, total_compressed_size / (1 << 20))
  )
  plt.yticks(ticks=range(len(addresses)), labels=addresses)
  plt.xlabel('Page index in SuperPage')

  handles = [
    matplotlib.patches.Patch(
      facecolor=WHITE, edgecolor='k', label='Empty / Decommitted'
    ),
    matplotlib.patches.Patch(
      facecolor=cmap(0.0), edgecolor='k', label='Compressible'
    ),
    matplotlib.patches.Patch(
      facecolor=cmap(1.0), edgecolor='k', label='Incompressible'
    ),
  ]
  plt.legend(handles=handles, loc='lower left', fontsize=12, framealpha=0.7)
  plt.savefig(output_filename, bbox_inches='tight')


def _PlotFragmentationCommon(bucket_to_allocated: dict, output_filename: str):
  slot_sizes = sorted(bucket_to_allocated.keys())
  allocated = []
  waste = []
  for slot_size in slot_sizes:
    slots = np.array(bucket_to_allocated[slot_size])
    allocated.append(np.sum(slots))
    waste.append(np.sum(slot_size - slots))

  plt.figure(figsize=(18, 8))
  indices = range(len(slot_sizes))
  plt.bar(indices, allocated, label='Requested Memory')
  b = plt.bar(
    indices, waste, bottom=allocated, label='Waste due to padding + bucketing'
  )

  waste_percentage = [
    ('%d%%' % (int(100.0 * w / (w + a)))) if a else ''
    for (w, a) in zip(waste, allocated)
  ]

  plt.bar_label(b, labels=waste_percentage)
  plt.xticks(indices, slot_sizes, rotation='vertical')
  plt.xlim(left=-0.5, right=len(indices))
  plt.xlabel('Bucket size')
  plt.ylabel('Memory (bytes)')
  plt.legend()

  total_allocated = np.sum(allocated)
  total_waste = np.sum(waste)
  plt.title(
    'Absolute and relative waste due to bucketing and padding'
    ' per bucket - Total Allocated = %dMiB, Waste = %d%%'
    % (
      total_allocated / (1 << 20),
      int(100 * total_waste / (total_allocated + total_waste)),
    )
  )
  plt.savefig(output_filename, bbox_inches='tight')


def _AdjustSizes(data: dict, bucket_sizes: list, adjustment_size: int) -> dict:
  requested_sizes = []
  for slot_span in data:
    requested_sizes += slot_span['allocated_sizes']

  # Assumes that all buckets have *some* live allocations.
  bucket_sizes = sorted(bucket_sizes)
  bucket_to_allocated = {slot_size: list() for slot_size in bucket_sizes}

  # Map the requested sizes without paddingb to buckets.
  for requested_size in requested_sizes:
    adjusted_size = requested_size + adjustment_size
    bucket_index = bisect.bisect_left(bucket_sizes, adjusted_size)
    slot_size = bucket_sizes[bucket_index]
    assert bucket_sizes[bucket_index] >= adjusted_size
    bucket_to_allocated[slot_size].append(adjusted_size)

  return bucket_to_allocated


def PlotSimulatedFragmentationData(
  data: dict, bucket_sizes: list, output_filename: str
):
  # No adjustment, want to check waste without any padding.
  bucket_to_allocated = _AdjustSizes(data, bucket_sizes, 0)
  _PlotFragmentationCommon(bucket_to_allocated, output_filename)


def PlotFragmentationData(data: dict, bucket_sizes: list, output_filename: str):
  # "Before allocation" takes only 4 bytes, but the instrumentation added
  # expands this to 8 bytes. To reconstruct what it would have been without it,
  # take the requested size and add 4 bytes to it.
  bucket_to_allocated = _AdjustSizes(data, bucket_sizes, _BRP_OVERHEAD)
  _PlotFragmentationCommon(bucket_to_allocated, output_filename)


def PlotDelta(data: dict, bucket_sizes: list, output_filename: str):
  bucket_to_allocated_brp = _AdjustSizes(data, bucket_sizes, _BRP_OVERHEAD)
  bucket_to_allocated_no_brp = _AdjustSizes(data, bucket_sizes, 0)
  slot_sizes = sorted(bucket_to_allocated_no_brp)

  allocated_per_bucket_brp = {
    slot_size: len(bucket_to_allocated_brp[slot_size]) * slot_size
    for slot_size in bucket_to_allocated_brp
  }
  allocated_per_bucket_no_brp = {
    slot_size: len(bucket_to_allocated_no_brp[slot_size]) * slot_size
    for slot_size in bucket_to_allocated_no_brp
  }
  delta = [
    allocated_per_bucket_brp[slot_size] - allocated_per_bucket_no_brp[slot_size]
    for slot_size in slot_sizes
  ]

  plt.figure(figsize=(18, 8))
  indices = range(len(slot_sizes))
  plt.bar(indices, delta, label='Delta')
  plt.xticks(indices, slot_sizes, rotation='vertical')
  plt.xlim(left=-0.5, right=len(indices))
  plt.xlabel('Bucket size')
  plt.ylabel('Memory delta (bytes)')
  plt.title(
    'Per-bucket size difference BRP/No-BRP - Total = %.1fMiB'
    % (sum(delta) / (1 << 20))
  )
  plt.savefig(output_filename, bbox_inches='tight')


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument(
    '--json', help='JSON dump from pa_tcache_inspect', required=True
  )
  parser.add_argument('--output', help='Output filename prefix', required=True)

  args = parser.parse_args()
  data = ParseJson(args.json)

  PlotSuperPageData(data['superpages'], args.output + '_superpage.png')
  PlotSuperPageCompressionData(
    data['superpages'], args.output + '_compression.png'
  )

  # Allocated object data is not always available.
  if 'allocated_sizes' in data:
    bucket_sizes = [x['slot_size'] for x in data['buckets']]
    PlotFragmentationData(
      data['allocated_sizes'], bucket_sizes, args.output + '_waste.png'
    )
    PlotSimulatedFragmentationData(
      data['allocated_sizes'],
      bucket_sizes,
      args.output + '_waste_simulated.png',
    )
    PlotDelta(data['allocated_sizes'], bucket_sizes, args.output + '_delta.png')


if __name__ == '__main__':
  main()