wd_js_bundle.bzl

load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
load("@capnp-cpp//src/capnp:cc_capnp_library.bzl", "cc_capnp_library")

CAPNP_TEMPLATE = """@{schema_id};

# generated by @workerd//build/wd_js_bundle.bzl

using Modules = import "/workerd/jsg/modules.capnp";

const {const_name} :Modules.Bundle = (
  modules = [
{modules}
]);
"""

MODULE_TEMPLATE = """    (name = "{name}", {src_type} = embed "{path}", type = {type}, {ts_declaration})"""

def _to_name(file_name):
    return file_name.removesuffix(".js")

def _to_d_ts(file_name):
    return file_name.removesuffix(".js") + ".d.ts"

def _relative_path(file_path, dir_path):
    if not file_path.startswith(dir_path):
        fail("file_path need to start with dir_path: " + file_path + " vs " + dir_path)
    return file_path.removeprefix(dir_path)

def _gen_api_bundle_capnpn_impl(ctx):
    output_dir = ctx.outputs.out.dirname + "/"

    def _render_module(name, label, src_type, type):
        return MODULE_TEMPLATE.format(
            name = name,
            # capnp doesn't allow ".." dir escape, make paths relative.
            # this won't work for embedding paths outside of rule directory subtree.
            src_type = src_type,
            path = _relative_path(
                ctx.expand_location("$(location {})".format(label), ctx.attr.data),
                output_dir,
            ),
            type = type,
            ts_declaration = (
                "tsDeclaration = embed \"" + _relative_path(
                    ctx.expand_location("$(location {})".format(ctx.attr.declarations[name]), ctx.attr.data),
                    output_dir,
                ) + "\", "
            ) if name in ctx.attr.declarations else "",
        )

    modules = [
        _render_module(ctx.attr.builtin_modules[m], m.label, "src", "builtin")
        for m in ctx.attr.builtin_modules
    ]
    modules += [
        _render_module(ctx.attr.internal_modules[m], m.label, "src", "internal")
        for m in ctx.attr.internal_modules
    ]
    modules += [
        _render_module(ctx.attr.internal_wasm_modules[m], m.label, "wasm", "internal")
        for m in ctx.attr.internal_wasm_modules
    ]
    modules += [
        _render_module(ctx.attr.internal_data_modules[m], m.label, "data", "internal")
        for m in ctx.attr.internal_data_modules
    ]
    modules += [
        _render_module(ctx.attr.internal_json_modules[m], m.label, "json", "internal")
        for m in ctx.attr.internal_json_modules
    ]

    content = CAPNP_TEMPLATE.format(
        schema_id = ctx.attr.schema_id,
        modules = ",\n".join(modules),
        const_name = ctx.attr.const_name,
    )
    ctx.actions.write(ctx.outputs.out, content)

gen_api_bundle_capnpn = rule(
    implementation = _gen_api_bundle_capnpn_impl,
    attrs = {
        "schema_id": attr.string(mandatory = True),
        "out": attr.output(mandatory = True),
        "builtin_modules": attr.label_keyed_string_dict(allow_files = True),
        "internal_modules": attr.label_keyed_string_dict(allow_files = True),
        "internal_wasm_modules": attr.label_keyed_string_dict(allow_files = True),
        "internal_data_modules": attr.label_keyed_string_dict(allow_files = True),
        "internal_json_modules": attr.label_keyed_string_dict(allow_files = True),
        "declarations": attr.string_dict(),
        "data": attr.label_list(allow_files = True),
        "const_name": attr.string(mandatory = True),
        "deps": attr.label_list(),
    },
)

def _copy_modules(modules, declarations):
    """Copy files from the modules map to the current package.

    Returns new module map using file copies.
    This is necessary since capnp compiler doesn't allow embeds outside of current subidrectory.
    """
    result = dict()
    declarations_result = dict()
    for m in modules:
        new_filename = modules[m].replace(":", "_").replace("/", "_")
        copy_file(name = new_filename + "@copy", src = m, out = new_filename)

        m_d_ts = _to_d_ts(m)
        if m_d_ts in declarations:
            new_d_ts_filename = new_filename + ".d.ts"
            copy_file(name = new_d_ts_filename + "@copy", src = m_d_ts, out = new_d_ts_filename)
            declarations_result[modules[m]] = str(native.package_relative_label(new_d_ts_filename))

        result[new_filename] = modules[m]
    return result, declarations_result

def wd_js_bundle_capnp(
        name,
        import_name,
        schema_id,
        builtin_modules,
        internal_modules = [],
        internal_wasm_modules = [],
        internal_data_modules = [],
        internal_json_modules = [],
        declarations = [],
        deps = []):
    """Generate cc capnp library with js api bundle.

    NOTE: Due to capnpc embed limitation all modules must be in the same or sub directory of the
          actual rule usage.

    Args:
     name: cc_capnp_library rule name
     import_name: The js import specifier. builtin modules are accessible from
                  user code under `<import_name>:<module_name>`, internal modules
                  are accessible from builtin modules under
                  `<import_name>-internal:<module_name>`
                  The capnproto bundle object generated will be called
                  `import_name` + "Bundle"
     schema_id: capnpn schema id
     builtin_modules: list of js source files for builtin modules
     internal_modules: list of js source files for internal modules
     internal_wasm_modules: list of wasm source files
     internal_data_modules: list of data source files
     internal_json_modules: list of json source files
     declarations: d.ts label set
     deps: dependency list
    Returns: The set of data dependencies
    """
    builtin_modules_dict = {
        m: "{}:{}".format(import_name, _to_name(m))
        for m in builtin_modules
    }
    internal_modules_dict = {
        m: "{}-internal:{}".format(import_name, _to_name(m.removeprefix("internal/")))
        for m in internal_modules
    }
    internal_wasm_modules_dict = {
        m: "{}-internal:{}".format(import_name, m.removeprefix("internal/"))
        for m in internal_wasm_modules
    }
    internal_data_modules_dict = {
        m: "{}-internal:{}".format(import_name, m.removeprefix("internal/"))
        for m in internal_data_modules
    }
    internal_json_modules_dict = {
        m: "{}-internal:{}".format(import_name, m.removeprefix("internal/"))
        for m in internal_json_modules
    }

    builtin_modules_dict, builtin_declarations = _copy_modules(
        builtin_modules_dict,
        declarations,
    )
    internal_modules_dict, internal_declarations = _copy_modules(
        internal_modules_dict,
        declarations,
    )
    internal_wasm_modules_dict, _ = _copy_modules(
        internal_wasm_modules_dict,
        declarations,
    )
    internal_data_modules_dict, _ = _copy_modules(
        internal_data_modules_dict,
        declarations,
    )
    internal_json_modules_dict, _ = _copy_modules(
        internal_json_modules_dict,
        declarations,
    )

    data = (
        list(builtin_modules_dict) +
        list(internal_modules_dict) +
        list(internal_wasm_modules_dict) +
        list(internal_data_modules_dict) +
        list(internal_json_modules_dict) +
        list(builtin_declarations.values()) +
        list(internal_declarations.values())
    )

    gen_api_bundle_capnpn(
        name = name + "@gen",
        out = name,
        schema_id = schema_id,
        const_name = import_name + "Bundle",
        builtin_modules = builtin_modules_dict,
        internal_modules = internal_modules_dict,
        internal_wasm_modules = internal_wasm_modules_dict,
        internal_data_modules = internal_data_modules_dict,
        internal_json_modules = internal_json_modules_dict,
        declarations = builtin_declarations | internal_declarations,
        data = data,
        deps = deps,
    )
    return data


def wd_js_bundle(name, import_name, *args, **kwargs):
    data = wd_js_bundle_capnp(name + ".capnp", import_name, *args, **kwargs)
    cc_capnp_library(
        name = name,
        srcs = [name + ".capnp"],
        strip_include_prefix = "",
        visibility = ["//visibility:public"],
        data = data,
        deps = ["@workerd//src/workerd/jsg:modules_capnp"],
        include_prefix = import_name,
    )