#!/usr/bin/env python3

import sys
import re

header = \
"""/*
 * WARNING: This file is autogenerated from gen_enum_defs.py [1].
 *
 * [1] https://github.com/sched-ext/scx/blob/main/scripts/gen_enum_defs.py
 */

#ifndef __ENUM_DEFS_AUTOGEN_H__
#define __ENUM_DEFS_AUTOGEN_H__

"""

footer = \
"""
#endif /* __ENUM_DEFS_AUTOGEN_H__ */
"""

abi_header = \
"""/*
 * WARNING: This file is autogenerated from gen_enum_defs.py [1].
 *
 * scx enumerator values from the vmlinux.h this tree is built against.
 * Used as the substitution source when the running kernel's BTF lacks
 * BTF_KIND_ENUM64 encoding and 64-bit enum values are truncated.
 *
 * [1] https://github.com/sched-ext/scx/blob/main/scripts/gen_enum_defs.py
 */

#ifndef __ENUMS_ABI_AUTOGEN_H__
#define __ENUMS_ABI_AUTOGEN_H__

struct __scx_enum_abi_val {
	const char *type;
	const char *name;
	u64 val;
};

static const struct __scx_enum_abi_val __scx_enum_abi_vals[]
	__attribute__((unused)) = {
"""

abi_footer = \
"""};

#endif /* __ENUMS_ABI_AUTOGEN_H__ */
"""

abi_rs_header = \
"""// WARNING: This file is autogenerated from gen_enum_defs.py [1].
//
// scx enumerator values from the vmlinux.h this tree is built against.
// Used as the substitution source when the running kernel's BTF lacks
// BTF_KIND_ENUM64 encoding and 64-bit enum values are truncated.
//
// [1] https://github.com/sched-ext/scx/blob/main/scripts/gen_enum_defs.py

/// (enum type name, enumerator name, value) triples from vmlinux.h.
pub const ENUM_ABI_VALUES: &[(&str, &str, u64)] = &[
"""

abi_rs_footer = \
"""];
"""

def is_enum_begin(toks):
    if len(toks) == 2:
        # enum {
        if toks[0] == "enum" and toks[-1] == "{":
            return True
    elif len(toks) == 3:
        # typedef enum {
        if toks[0] == "typedef" and toks[1] == "enum" and toks[2] == "{":
            return True
        # enum XXX {
        if toks[0] == "enum" and toks[-1] == "{":
            return True
    return False

def get_enum_type_name(toks):
    # enum XXX {
    if len(toks) == 3 and toks[0] == "enum":
        return toks[1]
    return None

def get_enum_val(line):
    toks = re.split(" |\t|\n|=|,", line)
    toks = list(filter(lambda tok: tok != "", toks))
    if len(toks) == 2:
        # x = 3,
        return (toks[0], toks[1])
    elif len(toks) == 1:
        # x,
        return (toks[0], None)
    return (None, None)

def parse_enum_num(val):
    # 9223372036854775808ULL, 65536, 0x10, possibly with U/L suffixes.
    # Negative values are reduced to their 64-bit two's complement
    # representation, matching how BTF readers sign-extend them to u64.
    return int(re.sub("[uUlL]+$", "", val), 0) & 0xffffffffffffffff

def is_enum_end(toks):
    if len(toks) == 1:
        # };
        if toks[0] == "};":
            return True;
    elif len(toks) >= 2:
        # } XXX;
        # } XXX: 8;
        if toks[0] == "}" and toks[-1][-1] == ";":
            return True
    return False

def gen_enum_header(f):
    print(header, file = f, end = "")

def gen_enum_have(f, enum):
    print("#define HAVE_%s" % (enum), file = f)

def gen_enum_footer(f):
    print(footer, file = f, end = "")

def gen_abi_vals(abi_vals_h, abi_vals):
    with open(abi_vals_h, "w") as f:
        print(abi_header, file = f, end = "")
        for (etype, enum, val) in abi_vals:
            print("\t{ \"%s\", \"%s\", 0x%xLLU }," % (etype, enum, val),
                  file = f)
        print(abi_footer, file = f, end = "")

def gen_abi_vals_rs(abi_vals_rs, abi_vals):
    with open(abi_vals_rs, "w") as f:
        print(abi_rs_header, file = f, end = "")
        for (etype, enum, val) in abi_vals:
            print("    (\"%s\", \"%s\", 0x%x)," % (etype, enum, val),
                  file = f)
        print(abi_rs_footer, file = f, end = "")

def gen_enum_defs(vmlinux_h, enum_defs_h, prefixes):
    STAT_UNKNOWN = 0
    STAT_ENUM = 1
    stat = STAT_UNKNOWN
    etype = None
    next_val = 0
    abi_vals = []
    with open(vmlinux_h, "r") as fv:
        with open(enum_defs_h, "w") as fe:
            gen_enum_header(fe)
            for line in fv:
                line = line.strip()
                toks = line.split()

                if stat == STAT_UNKNOWN:
                    if is_enum_begin(toks):
                        stat = STAT_ENUM
                        etype = get_enum_type_name(toks)
                        next_val = 0
                        continue
                elif stat == STAT_ENUM:
                    if is_enum_end(toks):
                        stat = STAT_UNKNOWN
                        continue
                    (enum, val) = get_enum_val(line)
                    if enum == None:
                        continue
                    num = parse_enum_num(val) if val != None else next_val
                    next_val = num + 1
                    for prefix in prefixes:
                        if enum.startswith(prefix):
                            gen_enum_have(fe, enum)
                            if etype != None:
                                abi_vals.append((etype, enum, num))
                            break
            gen_enum_footer(fe)
    return abi_vals

def parse_args(args):
    if len(args) != 2 and len(args) != 4:
        print("usage: gen_enum_defs.py [vmlinux.h] [enum_defs.autogen.h] "
              "([enums_abi.autogen.h] [enums_abi.autogen.rs])")
        print("")
        print("Helper script for autogenerating an enum definition header from vmlinux.h.")
        print("With 4 args, also generates C and Rust tables of enumerator ABI")
        print("values used to repair enum64 values truncated by kernel BTF.")
        exit(1)
    return args

"""
    Helper script for autogenerating an enum definition header.
"""
if __name__ == "__main__":
    args = parse_args(sys.argv[1:])
    (vmlinux_h, enum_defs_h) = (args[0], args[1])
    abi_vals = gen_enum_defs(vmlinux_h, enum_defs_h, ["SCX_", "__SCX_"])
    if len(args) == 4:
        gen_abi_vals(args[2], abi_vals)
        gen_abi_vals_rs(args[3], abi_vals)
