# Copyright lowRISC contributors.
# Licensed under the Apache License, Version 2.0, see LICENSE for details.
# SPDX-License-Identifier: Apache-2.0

"""Rules to build Matcha for the RiscV target"""

load("@rules_pkg//:pkg.bzl", "pkg_tar")
load(
    "@lowrisc_opentitan//rules:rv.bzl",
    "rv_rule",
    _OPENTITAN_CPU = "OPENTITAN_CPU",
    _OPENTITAN_PLATFORM = "OPENTITAN_PLATFORM",
    _opentitan_transition = "opentitan_transition",
)
load(
    "@lowrisc_opentitan//rules:opentitan.bzl",
    "bin_to_vmem",
    "opentitan_binary",
    "pick_correct_archive_for_device",
    "scramble_flash_vmem",
    "sign_bin",
)
load("@rules_cc//cc:find_cc_toolchain.bzl", "find_cc_toolchain")

# Re-exports of names from transition.bzl; many files in the repo use opentitan.bzl
# to get to them.
OPENTITAN_CPU = _OPENTITAN_CPU
OPENTITAN_PLATFORM = _OPENTITAN_PLATFORM
opentitan_transition = _opentitan_transition
KELVIN_PLATFORM = "//platforms/riscv32:kelvin"

# Currently set to secure core's IMC extensions.
SMC_PLATFORM = OPENTITAN_PLATFORM

# List of supported riscv core targets.
VERILATOR_CORE_TARGETS = {
    "secure_core": "@lowrisc_opentitan//sw/device/lib/arch:sim_verilator",
    "smc": "//sw/device/lib/arch:smc_sim_verilator",
}

DV_CORE_TARGETS = {
    "secure_core": "@lowrisc_opentitan//sw/device/lib/arch:sim_dv",
    "smc": "//sw/device/lib/arch:smc_sim_dv",
}

NEXUS_CORE_TARGETS = {
    "secure_core": "//sw/device/lib/arch:sc_fpga_nexus",
    "smc": "//sw/device/lib/arch:smc_fpga_nexus",
}

# This helper function generates a dictionary of per-device dependencies which is used to
# generate slightly different binaries for each hardware target, including two
# simulation platforms (DV and Verilator), and one FPGA platform (Nexus),
# given a target core from VERILATOR_CORE_TARGETS.
def device_deps(core):
    per_device_deps = {
        "sim_verilator": [VERILATOR_CORE_TARGETS.get(core)],
        "sim_dv": [DV_CORE_TARGETS.get(core)],
        "fpga_nexus": [NEXUS_CORE_TARGETS.get(core)],
    }
    return per_device_deps

def _elf_to_scrambled_rom_impl(ctx):
    outputs = []
    for src in ctx.files.srcs:
        if src.extension != "elf":
            fail("only ROM images in the ELF format may be converted to the VMEM format and scrambled.")
        scrambled = ctx.actions.declare_file(
            "{}.39.scr.vmem".format(
                # Remove ".elf" from file basename.
                src.basename.replace("." + src.extension, ""),
            ),
        )
        outputs.append(scrambled)
        ctx.actions.run(
            outputs = [scrambled],
            inputs = [
                src,
                ctx.executable._scramble_tool,
                ctx.file._config,
            ],
            arguments = [
                ctx.file._config.path,
                src.path,
                scrambled.path,
            ],
            executable = ctx.executable._scramble_tool,
        )
    return [DefaultInfo(
        files = depset(outputs),
        data_runfiles = ctx.runfiles(files = outputs),
    )]

elf_to_scrambled_rom_vmem = rv_rule(
    implementation = _elf_to_scrambled_rom_impl,
    attrs = {
        "srcs": attr.label_list(allow_files = True),
        "_scramble_tool": attr.label(
            default = "@lowrisc_opentitan//hw/ip/rom_ctrl/util:scramble_image",
            executable = True,
            cfg = "exec",
        ),
        "_config": attr.label(
            default = "@//hw/top_matcha:top_gen_rom_ctrl_hjson",
            allow_single_file = True,
        ),
    },
)

def opentitan_rom_binary(
        name,
        per_device_deps = device_deps("secure_core"),
        platform = OPENTITAN_PLATFORM,
        testonly = False,
        **kwargs):
    """A helper macro for generating OpenTitan binary artifacts for ROM.

    This macro is mostly a wrapper around the `opentitan_binary` macro, but also
    creates artifacts for each of the keys in `PER_DEVICE_DEPS`. The actual
    artifacts created are outputs of the rules emitted by the `opentitan_binary`
    macro and those listed below.
    Args:
      name: The name of this rule.
      per_device_deps: List of devices and device deps to build the target for.
      platform: The target platform for the artifacts.
      testonly: The target is built for test only.
      **kwargs: Arguments to forward to `opentitan_binary`.
    Emits rules:
      For each device in per_device_deps entry:
        rules emitted by `opentitan_binary` named: see `opentitan_binary` macro
        bin_to_rom_vmem                     named: <name>_<device>_vmem
        elf_to_scrambled_rom_vmem           named: <name>_<device>_scr_vmem
        filegroup                           named: <name>_<device>
          Containing all targets for a single device for the above generated rules.
        filegroup                           named: <name>
          Containing all targets across all devices for the above generated rules.
    """
    deps = kwargs.pop("deps", [])
    all_targets = []
    for (device, dev_deps) in per_device_deps.items():
        if device not in device_deps("secure_core"):
            fail("invalid device; device must be in {}".format(device_deps("secure_core").keys()))
        devname = "{}_{}".format(name, device)
        dev_targets = []

        # Generate ELF, Binary, Disassembly, and (maybe) sim_dv logs database
        dev_targets.extend(opentitan_binary(
            name = devname,
            deps = deps + dev_deps,
            extract_sw_logs_db = device == "sim_dv",
            testonly = testonly,
            **kwargs
        ))

        # We need to generate VMEM files even for FPGA devices, because we use
        # them for bitstream splicing.
        elf_name = "{}.{}".format(devname, "elf")
        bin_name = "{}_{}".format(devname, "bin")

        # Generate Un-scrambled ROM VMEM
        vmem_name = "{}_vmem".format(devname)
        dev_targets.append(":" + vmem_name)
        bin_to_vmem(
            name = vmem_name,
            bin = bin_name,
            platform = platform,
            testonly = testonly,
            word_size = 32,
        )

        # Generate Scrambled ROM VMEM
        scr_vmem_name = "{}_scr_vmem".format(devname)
        dev_targets.append(":" + scr_vmem_name)
        elf_to_scrambled_rom_vmem(
            name = scr_vmem_name,
            srcs = [elf_name],
            platform = platform,
            testonly = testonly,
        )

        # Create a filegroup with just the current device's targets.
        native.filegroup(
            name = devname,
            srcs = dev_targets,
            testonly = testonly,
        )
        all_targets.extend(dev_targets)

    # Create a filegroup with just all targets from all devices.
    native.filegroup(
        name = name,
        srcs = all_targets,
        testonly = testonly,
    )

def flash_binary(
        name,
        per_device_deps = device_deps("secure_core"),
        platform = OPENTITAN_PLATFORM,
        signing_keys = {
            "fake_test_key_0": "@lowrisc_opentitan//sw/device/silicon_creator/rom/keys/fake:test_private_key_0",
        },
        signed = False,
        sim_otp = None,
        testonly = False,
        manifest = None,
        var_name = None,
        word_size = 32,
        **kwargs):
    """A helper macro for generating binary artifacts for a target core.

    This macro is mostly a wrapper around the `opentitan_binary` macro, but also
    creates artifacts for each of the keys in `PER_DEVICE_DEPS`, and if signing
    is enabled, each of the keys in `signing_keys`. The actual artifacts created
    artifacts created are outputs of the rules emitted by the `opentitan_binary`
    macro and those listed below.
    Args:
      name: The name of this rule.
      per_device_deps: List of devices and device deps to build the target for.
      platform: The target platform for the artifacts.
      signing_keys: The signing keys for to sign each BIN file with.
      signed: Whether or not to emit signed binary/VMEM files.
      sim_otp: OTP image that contains flash scrambling keys / enablement flag
                      (only relevant for VMEM files built for sim targets).
      testonly: The target is built for test only.
      manifest: Partially populated manifest to set boot stage/slot configs.
      var_name: The variable name of the array in the c header file.
      word_size: word_size for vmem. For SEC (flash) it can be 64, but
                        SMC can only take 32.
      **kwargs: Arguments to forward to `opentitan_binary`.
    Emits rules:
      For each device in per_device_deps entry:
        rules emitted by `opentitan_binary` named: see `opentitan_binary` macro
        bin_to_vmem                         named: <name>_<device>_vmem64
        scrambled_flash_vmem                named: <name>_<device>_scr_vmem64
        Optionally:
          sign_bin             named: <name>_<device>_bin_signed_<key_name>
          bin_to_vmem          named: <name>_<device>_vmem64_signed_<key_name>
          scrambled_flash_vmem named: <name>_<device>_scr_vmem64_signed_<key_name>
        filegroup              named: <name>_<device>
          Containing all targets for a single device for the above generated rules.
        filegroup              named: <name>
          Containing all targets across all devices for the above generated rules.
    """
    deps = kwargs.pop("deps", [])
    all_targets = []
    for (device, dev_deps) in per_device_deps.items():
        if device not in device_deps("secure_core"):
            fail("invalid device; device must be in {}".format(device_deps("secure_core").keys()))
        devname = "{}_{}".format(name, device)
        dev_targets = []

        depname = "{}_deps".format(devname)
        pick_correct_archive_for_device(
            name = depname,
            deps = deps + dev_deps,
            device = device,
            testonly = testonly,
        )

        # Generate ELF, Binary, Disassembly, and (maybe) sim_dv logs database
        dev_targets.extend(opentitan_binary(
            name = devname,
            deps = [depname],
            extract_sw_logs_db = device == "sim_dv",
            testonly = testonly,
            **kwargs
        ))
        bin_name = "{}_{}".format(devname, "bin")

        # Sign BIN (if required) and generate scrambled VMEM images.
        if signed:
            if manifest == None:
                fail("A 'manifest' must be provided in order to sign flash images.")
            for (key_name, key) in signing_keys.items():
                # Sign the Binary.
                signed_bin_name = "{}_bin_signed_{}".format(devname, key_name)
                dev_targets.append(":" + signed_bin_name)
                sign_bin(
                    name = signed_bin_name,
                    bin = bin_name,
                    key = key,
                    key_name = key_name,
                    manifest = manifest,
                    testonly = testonly,
                )

                # We only need to generate VMEM files for sim devices.
                if device in ["sim_dv", "sim_verilator"]:
                    # Generate a VMEM64 from the signed binary.
                    signed_vmem_name = "{}_vmem_signed_{}".format(
                        devname,
                        key_name,
                    )
                    dev_targets.append(":" + signed_vmem_name)
                    bin_to_vmem(
                        name = signed_vmem_name,
                        bin = signed_bin_name,
                        platform = platform,
                        testonly = testonly,
                        word_size = word_size,
                    )

                    # Scramble / compute ECC for signed VMEM64.
                    scr_signed_vmem_name = "{}_scr_vmem_signed_{}".format(
                        devname,
                        key_name,
                    )
                    dev_targets.append(":" + scr_signed_vmem_name)
                    scramble_flash_vmem(
                        name = scr_signed_vmem_name,
                        otp = sim_otp,
                        vmem = signed_vmem_name,
                        platform = platform,
                        testonly = testonly,
                    )

        # We only need to generate VMEM files for sim devices.
        if device in ["sim_dv", "sim_verilator"]:
            # Generate a VMEM64 from the binary.
            vmem_name = "{}_vmem".format(devname)
            dev_targets.append(":" + vmem_name)
            bin_to_vmem(
                name = vmem_name,
                bin = bin_name,
                platform = platform,
                testonly = testonly,
                word_size = word_size,
            )

            # Scramble / compute ECC for VMEM.
            scr_vmem_name = "{}_scr_vmem".format(devname)
            dev_targets.append(":" + scr_vmem_name)
            scramble_flash_vmem(
                name = scr_vmem_name,
                otp = sim_otp,
                vmem = vmem_name,
                platform = platform,
                testonly = testonly,
            )
        if device == "fpga_nexus" and var_name:
            # Generate c header based on the bin file
            bin_c_header = "{}_bin_c".format(devname)
            dev_targets.append(":" + bin_c_header)
            bin_to_c_file(
                name = bin_c_header,
                srcs = [bin_name],
                var_name = var_name,
            )

        # Create a filegroup with just the current device's targets.
        native.filegroup(
            name = devname,
            srcs = dev_targets,
            testonly = testonly,
        )
        all_targets.extend(dev_targets)

    # Create a filegroup with just all targets from all devices.
    native.filegroup(
        name = name,
        srcs = all_targets,
        testonly = testonly,
    )

def sec_flash_binary(
        name,
        per_device_deps = device_deps("secure_core"),
        signing_keys = {
            "fake_test_key_0": "@lowrisc_opentitan//sw/device/silicon_creator/rom/keys/fake:test_private_key_0",
        },
        signed = False,
        **kwargs):
    """A helper macro for generating secure core binary artifacts, using the flash_binary wrapper.

    Args:
      name: The name of this rule.
      signing_keys: The signing keys for to sign each BIN file with.
      signed: Whether or not to emit signed binary/VMEM files.
      per_device_deps: The deps for each of the hardware target.
      **kwargs: Arguments to forward to `flash_binary`.
    """

    flash_binary(
        name = name,
        platform = OPENTITAN_PLATFORM,
        per_device_deps = per_device_deps,
        signing_keys = signing_keys,
        signed = signed,
        word_size = 64,  # Backdoor-load VMEM image uses 64-bit words
        **kwargs
    )

def smc_flash_binary(
        name,
        per_device_deps = device_deps("smc"),
        signing_keys = {
            "fake_test_key_0": "@lowrisc_opentitan//sw/device/silicon_creator/rom/keys/fake:test_private_key_0",
        },
        signed = False,
        var_name = "smc_bin",
        **kwargs):
    """A helper macro for generating SMC binary artifacts, using the flash_binary_wrapper.

    Args:
      name: The name of this rule.
      signing_keys: The signing keys for to sign each BIN file with.
      signed: Whether or not to emit signed binary/VMEM files.
      per_device_deps: The deps for each of the hardware target.
      var_name: The variable name of the array in the c header file.
      **kwargs: Arguments to forward to `flash_binary`.

    Emit rules:
      {name}_bin_c_header: A header file generated by bin_to_c_file for FPGA
    """

    flash_binary(
        name = name,
        platform = SMC_PLATFORM,
        per_device_deps = per_device_deps,
        signing_keys = signing_keys,
        signed = signed,
        var_name = var_name,
        word_size = 32,
        **kwargs
    )

def matcha_extflash_tar(
        name,
        sc_binary,
        smc_binary = None,
        kelvin_binary = None,
        data = None,
        **kwargs):
    """A rule for creating a tarball containing program code and resources.

    Args:
        name: The name of the output tarball. The .tar extension is automatically added.
        sc_binary: The binary for the secure core. This will be named matcha-tock-bundle.bin.
        smc_binary: The binary for the system management core. This will be named smc.bin, and is optional.
        kelvin_binary: The binary for the Kelvin core. This will be named kelvin.bin, and is optional.
        data: A dictionary of extra files to add to the archive. The key is the input file, and the value is filename in the archive.
        **kwargs: Arguments forwared to pkg_tar
    """

    files = {
        sc_binary: "matcha-tock-bundle.bin",
    }
    if smc_binary:
        files[smc_binary] = "smc.bin"
    if kelvin_binary:
        files[kelvin_binary] = "kelvin.bin"
    if data:
        files.update(data)
    pkg_tar(
        name = name,
        files = files,
        **kwargs
    )

def bin_to_c_file(
        name,
        srcs,
        var_name = "data",
        testonly = False,
        **kwargs):
    """A rule to merge binary files into a C include file.

    Args:
        name: The name of the target.
        srcs: The input binary source files.
        var_name: The variable name of the array in output C include file.
        testonly: The target is built for test only.
        **kwargs: Extra arguments forwards to genrule.
    """

    xxd_cmd = """
        xxd -i $< > $@;
        sed -i -e 's#\\w*_len#{}_len#g' $@;
        sed -i -e "s#\\w*\\[\\]#{}\\[\\]#g" $@;
        sed -i -e 's/unsigned/const unsigned/g' $@
        """.format(var_name, var_name)

    outs = ["{}.{}".format(name, "h")]

    native.genrule(
        name = name,
        srcs = srcs,
        outs = outs,
        cmd = xxd_cmd,
        testonly = testonly,
        **kwargs
    )

def _kelvin_transition_impl(settings, attr):
    return {"//command_line_option:platforms": KELVIN_PLATFORM}

kelvin_transition = transition(
    implementation = _kelvin_transition_impl,
    inputs = [],
    outputs = ["//command_line_option:platforms"],
)

def kelvin_rule(**kwargs):
    """
    A wrapper over rule() for creating rules that trigger
    the transition to the kelvin platform config.
    """
    attrs = kwargs.pop("attrs", {})
    if "platform" not in attrs:
        attrs["platform"] = attr.string(default = KELVIN_PLATFORM)
    attrs["_allowlist_function_transition"] = attr.label(
        default = "@bazel_tools//tools/allowlists/function_transition_allowlist",
    )

    return rule(
        cfg = kelvin_transition,
        attrs = attrs,
        **kwargs
    )

def _kelvin_binary_impl(ctx):
    """Implements compilation for kelvin executables.

    This rule compiles and links provided input into an executable
    suitable for use on the Kelvin core. Generates both an ELF
    and a BIN.

    Args:
        srcs: Input source files.
        deps: Target libraries that the binary depends upon.
        hdrs: Header files that are local to the binary.
        copts: Flags to pass along to the compiler.
        defines: Preprocessor definitions.
        linkopts: Flags to pass along to the linker.

    Output:
        OutputGroupsInfo to allow definition of filegroups
        containing the output ELF and BIN.
    """
    cc_toolchain = find_cc_toolchain(ctx).cc
    feature_configuration = cc_common.configure_features(
        ctx = ctx,
        cc_toolchain = cc_toolchain,
        requested_features = ctx.features,
        unsupported_features = ctx.disabled_features,
    )
    compilation_contexts = []
    linking_contexts = []
    for dep in ctx.attr.deps:
        if CcInfo in dep:
            compilation_contexts.append(dep[CcInfo].compilation_context)
            linking_contexts.append(dep[CcInfo].linking_context)
    (_compilation_context, compilation_outputs) = cc_common.compile(
        actions = ctx.actions,
        cc_toolchain = cc_toolchain,
        feature_configuration = feature_configuration,
        name = ctx.label.name,
        srcs = ctx.files.srcs,
        compilation_contexts = compilation_contexts,
        private_hdrs = ctx.files.hdrs,
        user_compile_flags = ctx.attr.copts,
        defines = ctx.attr.defines,
    )
    linking_outputs = cc_common.link(
        name = "{}.elf".format(ctx.label.name),
        actions = ctx.actions,
        feature_configuration = feature_configuration,
        cc_toolchain = cc_toolchain,
        compilation_outputs = compilation_outputs,
        linking_contexts = linking_contexts,
        user_link_flags = ctx.attr.linkopts + [
            "-Wl,-T,{}".format(ctx.file.linker_script.path),
            # binutils >= 2.39 checks the loadable segments RWX setting by default.
            "-Wl,--no-warn-rwx-segments",
        ],
        additional_inputs = depset([ctx.file.linker_script] + ctx.files.linker_script_includes),
        output_type = "executable",
    )

    binary = ctx.actions.declare_file(
        "{}.bin".format(
            ctx.attr.name,
        ),
    )
    ctx.actions.run(
        outputs = [binary],
        inputs = [linking_outputs.executable] + cc_toolchain.all_files.to_list(),
        arguments = [
            "-g",
            "-O",
            "binary",
            linking_outputs.executable.path,
            binary.path,
        ],
        executable = cc_toolchain.objcopy_executable,
    )

    return [
        DefaultInfo(
            files = depset([linking_outputs.executable, binary]),
        ),
        OutputGroupInfo(
            all_files = depset([linking_outputs.executable, binary]),
            elf_file = depset([linking_outputs.executable]),
            bin_file = depset([binary]),
        ),
    ]

kelvin_binary_impl = kelvin_rule(
    implementation = _kelvin_binary_impl,
    attrs = {
        "srcs": attr.label_list(allow_files = True),
        "deps": attr.label_list(allow_empty = True, providers = [CcInfo]),
        "hdrs": attr.label_list(allow_files = [".h"], allow_empty = True),
        "copts": attr.string_list(),
        "defines": attr.string_list(),
        "linkopts": attr.string_list(),
        "linker_script": attr.label(allow_single_file = True),
        "linker_script_includes": attr.label_list(default = [], allow_files = True),
        "_cc_toolchain": attr.label(default = Label("@bazel_tools//tools/cpp:current_cc_toolchain")),
    },
    fragments = ["cpp"],
    toolchains = ["@rules_cc//cc:toolchain_type"],
)

def kelvin_binary(name, srcs, **kwargs):
    """A helper macro for generating binary artifacts for the kelvin core.

    This macro uses the kelvin toolchain, kelvin-specific starting asm,
    and kelvin linker script to build kelvin binaries.

    Args:
      name: The name of this rule.
      srcs: The c source files.
      **kwargs: Additional arguments forward to cc_binary.
    Emits rules:
      filegroup              named: <name>.bin
        Containing the binary output for the target.
      filegroup              named: <name>.elf
        Containing all elf output for the target.
    """
    srcs = srcs + [
        "//sw/device/lib/testing/test_framework:kelvin_gloss.c",
        "//sw/device/lib/testing/test_framework:kelvin_start.S",
        "@lowrisc_opentitan//sw/device/lib/crt:crt.S",
    ]
    kelvin_binary_impl(
        name = name,
        srcs = srcs,
        linker_script = "//sw/device/lib/testing/test_framework:kelvin.ld",
        linker_script_includes = [
            "//hw/top_matcha/sw/autogen:top_matcha_memory.ld",
        ],
        **kwargs
    )

    # Need to create the following filegroups to make the output discoverable.
    native.filegroup(
        name = "{}.bin".format(name),
        srcs = [name],
        output_group = "bin_file",
    )
    native.filegroup(
        name = "{}.elf".format(name),
        srcs = [name],
        output_group = "elf_file",
    )