fune/mozglue/baseprofiler/build/generate_profiling_categories.py

# 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 script generates ProfilingCategoryList.h and profiling_categories.rs
# files from profiling_categories.yaml.

import yaml

CPP_HEADER_TEMPLATE = """\
/* 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/. */

#ifndef {includeguard}
#define {includeguard}

/* This file is generated by generate_profiling_categories.py from
   profiling_categories.yaml. DO NOT EDIT! */

// Profiler sub-categories are applied to each sampled stack to describe the
// type of workload that the CPU is busy with. Only one sub-category can be
// assigned so be mindful that these are non-overlapping. The active category is
// set by pushing a label to the profiling stack, or by the unwinder in cases
// such as JITs. A profile sample in arbitrary C++/Rust will typically be
// categorized based on the top of the label stack.
//
// The list of available color names for categories is:
//    transparent
//    blue
//    green
//    grey
//    lightblue
//    magenta
//    orange
//    purple
//    yellow

// clang-format off

{contents}

// clang-format on

#endif // {includeguard}
"""

CPP_MACRO_DEFINITION = """\
#define MOZ_PROFILING_CATEGORY_LIST(BEGIN_CATEGORY, SUBCATEGORY, END_CATEGORY) \\
"""

RUST_TEMPLATE = """\
/* 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 file is generated by generate_profiling_categories.py from
   profiling_categories.yaml. DO NOT EDIT! */

{contents}\
"""

RUST_ENUM_TEMPLATE = """\
#[repr(u32)]
#[derive(Debug, Copy, Clone)]
pub enum {name} {{
{fields}
}}
"""

RUST_CONVERSION_IMPL_TEMPLATE = """\
impl {name} {{
  pub fn to_cpp_enum_value(&self) -> u32 {{
{content}
  }}
}}
"""

RUST_DEFAULT_IMPL_TEMPLATE = """\
impl Default for {name} {{
  fn default() -> Self {{
{content}
  }}
}}
"""

RUST_MATCH_SELF = """\
    match *self {{
{fields}
    }}
"""


def generate_header(c_out, includeguard, contents):
    c_out.write(
        CPP_HEADER_TEMPLATE.format(includeguard=includeguard, contents=contents)
    )


def generate_rust_file(c_out, contents):
    c_out.write(RUST_TEMPLATE.format(contents=contents))


def load_yaml(yaml_path):
    file_handler = open(yaml_path)
    return yaml.safe_load(file_handler)


def generate_category_macro(name, label, color, subcategories):
    contents = '  BEGIN_CATEGORY({name}, "{label}", "{color}") \\\n'.format(
        name=name, label=label, color=color
    )

    subcategory_items = []

    for subcategory in subcategories:
        subcat_name = subcategory["name"]
        assert isinstance(subcat_name, str)
        subcat_label = subcategory["label"]
        assert isinstance(subcat_label, str)

        subcategory_items.append(
            '    SUBCATEGORY({parent_cat}, {name}, "{label}") \\\n'.format(
                parent_cat=name, name=subcat_name, label=subcat_label
            )
        )

    contents += "".join(subcategory_items)
    contents += "  END_CATEGORY"

    return contents


def generate_macro_header(c_out, yaml_path):
    """Generate ProfilingCategoryList.h from profiling_categories.yaml.
    The generated file has a macro to generate the profiling category enums.
    """

    data = load_yaml(yaml_path)

    # Stores the macro definition of each categories.
    category_items = []

    for category in data:
        name = category["name"]
        assert isinstance(name, str)
        label = category["label"]
        assert isinstance(label, str)
        color = category["color"]
        assert isinstance(color, str)
        subcategories = category.get("subcategories", None)
        assert (
            isinstance(subcategories, list) and len(subcategories) > 0
        ), "At least one subcategory expected as default in {}.".format(name)

        category_items.append(
            generate_category_macro(name, label, color, subcategories)
        )

    contents = CPP_MACRO_DEFINITION
    contents += " \\\n".join(category_items)

    generate_header(c_out, "baseprofiler_ProfilingCategoryList_h", contents)


class RustEnum:
    """Class that keeps the rust enum fields and impls.
    This is used for generating the Rust ProfilingCategoryPair and ProfilingCategory
    enums as well as ProfilingCategoryPair's sub category enums.
    For example, this can either generate an enum with discrimant fields for sub
    category enums and ProfilingCategory:
    ```
    #[repr(u32)]
    #[derive(Debug, Copy, Clone)]
    pub enum Graphics {
        LayerBuilding = 0,
        ...
    }
    ```
    or can generate an enum with optional tuple values for ProfilingCategoryPair
     to explicitly mention their sub categories:
    ```
    #[repr(u32)]
    #[derive(Debug, Copy, Clone)]
    pub enum ProfilingCategoryPair {
        Network(Option<Network>),
        ...
    }
    ```

    And in addition to enums, it will generate impls for each enum. See one
    example below:
    ```
    impl Default for Network {
        fn default() -> Self {
            Network::Other
        }
    }
    ```
    """

    def __init__(self, name):
        # Name of the Rust enum.
        self.name = name
        # Fields of the Rust enum. This list contains elements of
        # (field_name, field_string) tuple for convenience.
        self.fields = []
        # Impls of the Rust enum. Each element is a string.
        self.impls = []
        # Default category of the Rust enum. Main enums won't have it, but all
        # sub category enums must have one. This is being checked later.
        self.default_category = None

    def append_optional_tuple_field(self, field_name):
        """Append the enum fields list with an optional tuple field."""
        field = (field_name, "  {name}(Option<{name}>),".format(name=field_name))
        self.fields.append(field)

    def append_discriminant_field(self, field_name, field_value):
        """Append the enum fields list with a discriminant field."""
        field = (
            field_name,
            "  {name} = {value},".format(name=field_name, value=field_value),
        )
        self.fields.append(field)

    def append_default_impl(self, default_category):
        """Append the enum impls list with a default implementation."""
        self.default_category = default_category

        self.impls.append(
            RUST_DEFAULT_IMPL_TEMPLATE.format(
                name=self.name,
                content="      {category}::{subcategory}".format(
                    category=self.name, subcategory=self.default_category
                ),
            )
        )

    def append_conversion_impl(self, content):
        """Append the enum impls list with a conversion implementation for cpp values."""
        self.impls.append(
            RUST_CONVERSION_IMPL_TEMPLATE.format(name=self.name, content=content)
        )

    def to_rust_string(self):
        """Serialize the enum with its impls as a string"""
        joined_fields = "\n".join(map(lambda field: field[1], self.fields))
        result = RUST_ENUM_TEMPLATE.format(name=self.name, fields=joined_fields)
        result += "\n"
        result += "\n".join(self.impls)
        return result


def generate_rust_enums(c_out, yaml_path):
    """Generate profiling_categories.rs from profiling_categories.yaml.
    The generated file has a profiling category enums and their impls.
    """

    data = load_yaml(yaml_path)

    # Each category has its own enum for keeping its subcategories. We are
    # keeping all of them here.
    enums = []
    # Parent enums for prifiling category and profiling category pair. They will
    # be appended to the end of the `enums`.
    profiling_category_pair_enum = RustEnum("ProfilingCategoryPair")
    profiling_category_enum = RustEnum("ProfilingCategory")
    profiling_category_pair_value = 0

    for cat_index, category in enumerate(data):
        cat_name = category["name"]
        assert isinstance(cat_name, str)
        cat_label = category["label"]
        assert isinstance(cat_label, str)
        # This will be used as our main enum field and sub category enum.
        cat_label = "".join(filter(str.isalnum, cat_label))
        cat_subcategories = category.get("subcategories", None)
        assert (
            isinstance(cat_subcategories, list) and len(cat_subcategories) > 0
        ), "At least one subcategory expected as default in {}.".format(cat_name)

        # Create a new enum for this sub category and append it to the enums list.
        category_enum = RustEnum(cat_label)
        enums.append(category_enum)

        for subcategory in cat_subcategories:
            subcat_name = subcategory["name"]
            assert isinstance(subcat_name, str)
            subcat_label = subcategory["label"]
            assert isinstance(subcat_label, str)
            friendly_subcat_name = None

            if cat_name == subcat_name:
                # This is the default sub-category. It should use the label as name.
                friendly_subcat_name = subcat_label
                category_enum.append_default_impl(subcat_label)
            else:
                # This is a non-default sub-category.
                underscore_pos = subcat_name.find("_")
                friendly_subcat_name = subcat_name[underscore_pos + 1 :]

            friendly_subcat_name = "".join(filter(str.isalnum, friendly_subcat_name))
            category_enum.append_discriminant_field(
                friendly_subcat_name, profiling_category_pair_value
            )
            profiling_category_pair_value += 1

        assert (
            category_enum.default_category is not None
        ), "There must be a default subcategory with the same name."

        # Append the main enums.
        profiling_category_pair_enum.append_optional_tuple_field(cat_label)
        profiling_category_enum.append_discriminant_field(cat_label, cat_index)

    # Add the main enums impls for conversion into cpp values.
    profiling_category_pair_impl_fields = "\n".join(
        "      {enum_name}::{field_name}(val) => val.unwrap_or_default() as u32,".format(
            enum_name="ProfilingCategoryPair", field_name=field
        )
        for field, _ in profiling_category_pair_enum.fields
    )
    profiling_category_pair_enum.append_conversion_impl(
        RUST_MATCH_SELF.format(fields=profiling_category_pair_impl_fields)
    )
    profiling_category_enum.append_conversion_impl("      *self as u32")

    # After adding all the sub category enums, we can add the main enums to the list.
    enums.append(profiling_category_pair_enum)
    enums.append(profiling_category_enum)

    # Print all the enums and their impls.
    contents = "\n".join(map(lambda enum: enum.to_rust_string(), enums))
    generate_rust_file(c_out, contents)