Bzlmod support for rules_ispc

Question

Currently, rules_ispc only supports the traditional WORKSPACE approach. I want to add Bzlmod support to this ruleset. When looking at other rulesets, such as rules_ruby the final usage of those rules could look like this in MODULE.bazel file:

// add this to your MODULE.bazel file
bazel_dep(name = "rules_ispc", version = "0.0.5")

ispc = use_extension("@rules_ispc//ispc:extension.bzl", "ispc")
ispc.toolchain(
  name = "ispc",
  version = "1.23.0",
)
use_repo(ispc, "ispc")

register_toolchains("@ispc_toolchains//:all")

In a BUILD file, I want then to use something like this:

load("@rules_ispc//:ispc.bzl", "ispc_cc_library")

ispc_cc_library(
    name = "square",
    srcs = ["square.ispc"],
    out = "square.h",
    ispc_main_source_file = "square.ispc",
)

cc_binary(
    name = "main",
    srcs = ["main.cpp"],
    deps = [":square"],
)

I wonder if this is the proper way to go in bzlmod?

I have created an extension.bzl file:

"""Module extensions used by bzlmod"""

ispc_toolchain = tag_class(attrs = {
    "name": attr.string(doc = "Base name for generated repositories, allowing multiple to be registered."),
    "version": attr.string(doc = "Explicit version of ISPC."),
})

def _ispc_module_extension(module_ctx):
    registrations = {}
    
    # Does it make sense to call here register_toolchains?
    # native.register_toolchains(...)
    # How can I do this for different pre-compiled binaries of ISPC?


ispc = module_extension(
    implementation = _ispc_module_extension,
    tag_classes = {
        "toolchain": ispc_toolchain,
    },
)

I have also created a toolchain.bzl file that looks like this:

"""Toolchain for ISPC compiler."""

IspcToolchainInfo = provider(
    doc = "Information about how to invoke ISPC.",
    fields = [
        "ispc_path",
        "default_target",
        "default_target_os",
        "default_arch",
    ],
)

def _ispc_toolchain_impl(ctx):
    expand_ispc_path = ctx.expand_location(ctx.attr.ispc_cmd, ctx.attr.data)
    toolchain_info = platform_common.ToolchainInfo(
        ispc_info = IspcToolchainInfo(
            ispc_path = expand_ispc_path,
            default_target = ctx.attr.default_target,
            default_target_os = ctx.attr.default_target_os,
            default_arch = ctx.attr.default_arch,
        ),
    )
    return [toolchain_info]

ispc_toolchain = rule(
    implementation = _ispc_toolchain_impl,
    attrs = {
        "ispc_cmd": attr.string(),
        "default_target": attr.string(),
        "default_target_os": attr.string(),
        "default_arch": attr.string(),
        "data": attr.label_list(allow_files = True),
    },
)

I am unsure how to connect the dots to get this working. Also, I am not sure if this is the right approach. Any hints are welcome.