#!/usr/bin/python
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

# This program splits up a jprof profile into multiple files based on a
# list of functions in a text file.  First, a complete profile is
# generated.  Then, for each line in the text file, a profile is
# generated containing only stacks that go through that line, and also
# excluding all stacks in earlier lines in the text file.  This means
# that the text file, from start to end, is splitting out pieces of the
# profile in their own file.  Finally, a final profile containing the
# remainder is produced.

# The program takes four arguments:
#   (1) The path to jprof.
#   (2) The path to the text file describing the splits.  The output
#       will be placed in the same directory as this file.
#   (3) The program that was profiled.
#   (4) The jprof-log file generated by the profile, to be split up.
# (Really, all arguments from (3) and later are passed through to
# jprof, so additional arguments could be provided if you want to pass
# additional arguments to jprof.)

# In slightly more detail:
#
# This script uses jprof's includes (-i) and excludes (-e) options to
# split profiles into segments.  It takes as input a single text file,
# and from that text file creates a series of jprof profiles in the
# directory the text file is in.
#
# The input file format looks like the following:
#
#   poll g_main_poll
#   GetRuleCascade CSSRuleProcessor::GetRuleCascade(nsPresContext *, nsAtom *)
#   RuleProcessorData RuleProcessorData::RuleProcessorData
#        (nsPresContext *, nsIContent *, nsRuleWalker *, nsCompatibility *)
#
#
# From this input file, the script will construct a profile called
# jprof-0.html that contains the whole profile, a profile called
# jprof-1-poll.html that includes only stacks with g_main_poll, a
# profile called jprof-2-GetRuleCascade.html that includes only stacks
# that have GetRuleCascade and do not have g_main_poll, a profile called
# jprof-3-RuleProcessorData.html that includes only stacks that have the
# RuleProcessorData constructor and do not have GetRuleCascade or
# g_main_poll, and a profile called jprof-4.html that includes only
# stacks that do not have any of the three functions in them.
#
# This means that all of the segments of the profile, except
# jprof-0.html, are mutually exclusive.  Thus clever ordering of the
# functions in the input file can lead to a logical splitting of the
# profile into segments.

from __future__ import absolute_import, print_function

import sys
import subprocess
import os.path

if len(sys.argv) < 5:
    sys.stderr.write("Expected arguments: <jprof> <split-file> <program> <jprof-log>\n")
    sys.exit(1)

jprof = sys.argv[1]
splitfile = sys.argv[2]
passthrough = sys.argv[3:]

for f in [jprof, splitfile]:
    if not os.path.isfile(f):
        sys.stderr.write("could not find file: {0}\n".format(f))
        sys.exit(1)


def read_splits(splitfile):
    """
    Read splitfile (each line of which contains a name, a space, and
    then a function name to split on), and return a list of pairs
    representing exactly that.  (Note that the name cannot contain
    spaces, but the function name can, and often does.)
    """
    def line_to_split(line):
        line = line.strip("\r\n")
        idx = line.index(" ")
        return (line[0:idx], line[idx + 1:])

    io = open(splitfile, "r")
    result = [line_to_split(line) for line in io]
    io.close()
    return result


splits = read_splits(splitfile)


def generate_profile(options, destfile):
    """
    Run jprof to generate one split of the profile.
    """
    args = [jprof] + options + passthrough
    print("Generating {}".format(destfile))
    destio = open(destfile, "w")
    # jprof expects the "jprof-map" file to be in its current working directory
    cwd = None
    for option in passthrough:
        if option.find("jprof-log"):
            cwd = os.path.dirname(option)
    if cwd is None:
        raise Exception("no jprof-log option given")
    process = subprocess.Popen(args, stdout=destio, cwd=cwd)
    process.wait()
    destio.close()
    if process.returncode != 0:
        os.remove(destfile)
        sys.stderr.write("Error {0} from command:\n  {1}\n".format(
            process.returncode, " ".join(args)))
        sys.exit(process.returncode)


def output_filename(number, splitname):
    """
    Return the filename (absolute path) we should use to output the
    profile segment with the given number and splitname.  Splitname
    should be None for the complete profile and the remainder.
    """
    def pad_count(i):
        result = str(i)
        # 0-pad to the same length
        result = "0" * (len(str(len(splits) + 1)) - len(result)) + result
        return result

    name = pad_count(number)
    if splitname is not None:
        name += "-" + splitname

    return os.path.join(os.path.dirname(splitfile),
                        "jprof-{0}.html".format(name))


# generate the complete profile
generate_profile([], output_filename(0, None))

# generate the listed splits
count = 1
excludes = []
for (splitname, splitfunction) in splits:
    generate_profile(excludes + ["-i" + splitfunction],
                     output_filename(count, splitname))
    excludes += ["-e" + splitfunction]
    count = count + 1

# generate the remainder after the splits
generate_profile(excludes, output_filename(count, None))