#!/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.
"""Parses the output of vmmap to query IOSurface memory usage."""

import argparse
import collections
import logging
import operator
import re
import subprocess
import sys

IOSurface = collections.namedtuple(
  'IOSurface',
  [
    'start',
    'end',
    'virtual',
    'resident',
    'dirty',
    'swapped',
    'width',
    'height',
    'size',
  ],
)


def _FormatSize(n: int):
  _KIB = 1024
  _MIB = 1024 * _KIB
  result = ''
  if n > _MIB:
    result = '%.1fMiB' % (n / _MIB)
  else:
    result = '%0.1fkiB' % (n / _KIB)
  return result.rjust(8)


def IOSurfaceToString(iosurface: IOSurface):
  return (
    '%x-%x\tVirtual Size: %s\tResident: %s\tDirty: %s\tSwapped: %s'
    '\tDimensions: %dx%d'
  ) % (
    iosurface.start,
    iosurface.end,
    _FormatSize(iosurface.virtual),
    _FormatSize(iosurface.resident),
    _FormatSize(iosurface.dirty),
    _FormatSize(iosurface.swapped),
    iosurface.width,
    iosurface.height,
  )


def ExecuteVmmap(pid: int) -> str:
  """Runs vmmap PID and returns its output."""
  ret = subprocess.run(['vmmap', str(pid)], capture_output=True)
  ret.check_returncode()
  stdout = ret.stdout.decode('utf-8')
  return stdout


def _ParseSize(size):
  suffix = size[-1]
  if suffix == 'K':
    return int(1024 * float(size[:-1]))
  elif suffix == 'M':
    return int(1024 * 1024 * float(size[:-1]))
  else:
    return int(size)


def _PrettyPrint(size: int) -> str:
  _KIB = 1024
  _MIB = 1024 * _KIB
  _GIB = 1024 * _MIB
  if size > _GIB:
    return '%.1fGiB' % (size / _GIB)
  elif size > _MIB:
    return '%.1fMiB' % (size / _MIB)
  elif size > _KIB:
    return '%.1fkiB' % (size / _KIB)
  else:
    return str(size)


def ParseIOSurface(contents: str, quiet=False) -> list:
  """From the content of a vmmap file, returns a list of IOSurfaces."""
  io_surfaces = []
  for line in contents.split('\n'):
    _ADDRESS_RE = '[0-9a-f]{1,16}'
    # Has an optional suffix, in which case it's a float
    _SIZE_RE = '[0-9\.KM]*'
    _SIZE_DECIMAL = '[0-9]*'
    _RE = (
      '^IOSurface *'
      '(?P<start>%(address)s)-(?P<end>%(address)s) *'
      '\[ *(?P<virtual>%(size)s) *(?P<resident>%(size)s) '
      '*(?P<dirty>%(size)s) *(?P<swapped>%(size)s)\]'
      '.*SurfaceID: 0x[0-9a-f]* *'
      '(?P<width>%(size_decimal)s)x(?P<height>%(size_decimal)s)'
      ' *\([^\(]*\) *'
      '(?P<size>%(size)s)'
    ) % {
      'address': _ADDRESS_RE,
      'size': _SIZE_RE,
      'size_decimal': _SIZE_DECIMAL,
    }
    regexp = re.compile(_RE)
    _SAMPLE_LINE = (
      'IOSurface                   2a7d38000-2aae10000    '
      '[ 48.8M     0K     0K  48.8M] rw-/rw- SM=SHM PURGE=N'
      '  SurfaceID: 0x5f  2440x5196 (BGRA) 48.8M'
    )
    assert regexp.match(_SAMPLE_LINE)
    _SAMPLE_LINE2 = (
      'IOSurface                   10c2d0000-10c2d4000    '
      '[   16K     0K     0K     0K] rw-/rw- SM=SHM PURGE=N'
      '  SurfaceID: 0x22e  19x19 (BGRA) 2560'
    )
    assert regexp.match(_SAMPLE_LINE2)

    matches = regexp.match(line)
    if matches:
      if not quiet:
        print(line)
      start = int(matches.group('start'), 16)
      end = int(matches.group('end'), 16)
      virtual = _ParseSize(matches.group('virtual'))
      resident = _ParseSize(matches.group('resident'))
      dirty = _ParseSize(matches.group('dirty'))
      swapped = _ParseSize(matches.group('swapped'))
      width = int(matches.group('width'))
      height = int(matches.group('height'))
      size = _ParseSize(matches.group('size'))

      io_surface = IOSurface(
        start, end, virtual, resident, dirty, swapped, width, height, size
      )
      io_surfaces.append(io_surface)
  return io_surfaces


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument(
    '--pid', help='PID of the process to get the maps of', type=int
  )
  parser.add_argument(
    '--filename', help='Path to existing output of vmmap', type=str
  )
  args = parser.parse_args()

  contents = None
  if args.pid:
    contents = ExecuteVmmap(args.pid)
  else:
    assert args.filename
    with open(args.filename, 'r') as f:
      contents = f.read()

  io_surfaces = ParseIOSurface(contents)
  io_surfaces.sort(key=operator.attrgetter('virtual'))
  print('\nIOSurfaces sorted by virtual size:')
  for io_surface in io_surfaces:
    print('\t' + IOSurfaceToString(io_surface))

  io_surfaces.sort(key=operator.attrgetter('width'))
  print('\nIOSurfaces sorted by width:')
  for io_surface in io_surfaces:
    print('\t' + IOSurfaceToString(io_surface))

  io_surfaces.sort(key=lambda x: x.dirty + x.swapped)
  print('\nIOSurfaces sorted by dirty/swapped:')
  for io_surface in io_surfaces:
    print('\t' + IOSurfaceToString(io_surface))

  lost_to_paging = sum(
    io_surface.virtual - io_surface.size for io_surface in io_surfaces
  )
  dirty = sum(io_surface.dirty for io_surface in io_surfaces)
  swapped = sum(io_surface.swapped for io_surface in io_surfaces)
  print(
    '\nMemory lost due to page rounding = %s' % _PrettyPrint(lost_to_paging)
  )
  print('Dirty Memory = %s' % _PrettyPrint(dirty))
  print('Swapped Memory = %s' % _PrettyPrint(swapped))


if __name__ == '__main__':
  main()