botw/tools/generate_gdt_common_flags.py

#!/usr/bin/env python3
import argparse
from pathlib import Path
import typing as tp

import oead


def add_development_remnant_flags(flags: tp.Dict[str, str]):
    _flags = {
        "AoC_DragonFireChallengeRing_Advent": "bool_data",
        "AoC_RandomSpawnTreasure_Contents": "string64_data",
        "AoC_RandomSpawnTreasure_IsRandomized": "bool_data",
        "AoC_TestProg_Imoto_Flag_00": "bool_data",
        "AoC_TestProg_Imoto_TagCount_00": "s32_data",
        "AocTestEx_Omosako_IsPastWorld": "bool_data",
        "AocTestEx_Omosako_ReturnToMainField_Position": "vector3f_data",
        "AocTestEx_Omosako_ReturnToMainField_Rotation": "f32_data",
        "AocTestEx_Omosako_SandOfTime_Num": "s32_data",
        "AocTestEx_Omosako_SandOfTime_Rate": "f32_data",
        "Location_DarkDungeon01": "s32_data",
        "Location_DarkDungeon02": "s32_data",
        "Location_DarkDungeon03": "s32_data",
        "Location_DarkDungeon04": "s32_data",
        "SpurGear_revolve_01": "bool_data",
        "SpurGear_revolve_02": "bool_data",
    }
    flags.update(_flags)


def load_flag_types(root: Path) -> tp.Dict[str, str]:
    flag_types: tp.Dict[str, str] = dict()
    add_development_remnant_flags(flag_types)

    gdt_dir = root / "GameData"
    for path in gdt_dir.glob("Flag/*.yml"):
        flag_list = oead.byml.from_text(path.read_text(encoding="utf-8"))
        keys = list(flag_list.keys())
        assert len(keys) == 1
        flag_type = keys[0]
        for flag in flag_list[flag_type]:
            flag_types[flag["DataName"]] = flag_type

    return flag_types


def write_struct_chunk(f: tp.TextIO, flags: tp.Collection[str], i: int) -> None:
    f.write(f"""\
struct CommonFlags{i} {{
""")
    for name in flags:
        f.write(f"    FlagHandle flag_{name} = InvalidHandle;\n")
    f.write(f"""\
    u32 _pad = 0;
}};

[[gnu::visibility("hidden")]] extern CommonFlags{i} sCommonFlags{i};
""")


FLAGS_PER_CHUNK = 1023


def chunk_flag_iterator(flags: tp.Iterator[str]):
    while True:
        chunk = []
        for i in range(FLAGS_PER_CHUNK):
            try:
                chunk.append(next(flags))
            except StopIteration:
                break

        if not chunk:
            return

        yield chunk


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("root", help="Path to the root of the GameROM source")
    parser.add_argument("exe_flag_list", help="Path to a file containing a list of flags")
    args = parser.parse_args()

    src_root = Path(__file__).parent.parent
    src_gdt = src_root / "src" / "KingSystem" / "GameData"
    root = Path(args.root)
    exe_flag_list = Path(args.exe_flag_list).read_text().splitlines()
    flag_types = load_flag_types(root)

    with (src_gdt / "gdtCommonFlags.h").open("w") as f:
        f.write("""\
#pragma once

// DO NOT EDIT. This file is automatically generated.

#include "KingSystem/GameData/gdtManager.h"

namespace ksys::gdt {

// clang-format off

namespace detail {

""")
        for i, chunk in enumerate(chunk_flag_iterator(iter(exe_flag_list))):
            write_struct_chunk(f, chunk, i)
            f.write("\n")
        f.write("""\
}  // namespace detail

void initCommonFlags();

""")

        for i, name in enumerate(exe_flag_list):
            chunk_idx: int = i // FLAGS_PER_CHUNK
            f.write(f"inline FlagHandle& flag_{name}() {{ return detail::sCommonFlags{chunk_idx}.flag_{name}; }}\n")
            pass

        f.write("""\

// clang-format on

}  // namespace ksys::gdt
""")

    # Generate the implementation.

    fn_names = {
        "bool_data": "getBoolHandle",
        "s32_data": "getS32Handle",
        "f32_data": "getF32Handle",
        "string_data": "getStrHandle",
        "string64_data": "getStr64Handle",
        "string256_data": "getStr256Handle",
        "vector2f_data": "getVec2fHandle",
        "vector3f_data": "getVec3fHandle",
        "vector4f_data": "getVec4fHandle",

        "bool_array_data": "getBoolArrayHandle",
        "s32_array_data": "getS32ArrayHandle",
        "f32_array_data": "getF32ArrayHandle",
        "string_array_data": "getStrArrayHandle",
        "string64_array_data": "getStr64ArrayHandle",
        "string256_array_data": "getStr256ArrayHandle",
        "vector2f_array_data": "getVec2fArrayHandle",
        "vector3f_array_data": "getVec3fArrayHandle",
        "vector4f_array_data": "getVec4fArrayHandle",
    }

    with (src_gdt / "gdtCommonFlags.cpp").open("w") as f:
        f.write("""\
// DO NOT EDIT. This file is automatically generated.

#include "KingSystem/GameData/gdtCommonFlags.h"

namespace ksys::gdt {

namespace detail {

""")
        for i in range(len(exe_flag_list) // FLAGS_PER_CHUNK + 1):
            f.write(f"CommonFlags{i} sCommonFlags{i};\n")
        f.write("""
}  // namespace detail

void initCommonFlags_();

void initCommonFlags() {
    initCommonFlags_();
}

void initCommonFlags_() {
    auto* mgr = Manager::instance();
    if (!mgr)
        return;

    // clang-format off

""")
        for flag_name in exe_flag_list:
            f.write(f"    flag_{flag_name}() = mgr->{fn_names[flag_types[flag_name]]}(\"{flag_name}\");\n")

        f.write("""\

    // clang-format on
}

}  // namespace ksys::gdt
""")


if __name__ == "__main__":
    main()