From 8b8e59090b47e85e8d9af6c1815a2395961bd373 Mon Sep 17 00:00:00 2001 From: Alex Trotta Date: Tue, 9 Jun 2026 11:38:33 -0400 Subject: [PATCH] Add src_deps option for mojo_library The idea here is to speed up the critical path of Mojo compilation chains by doing `-I ` instead of `-I `. This comes at the cost of individual compile actions time. Not sure if we're going to merge this, while it _does_ speed up the critical path in our monorepo some, all actions are so much slower (at least 2x) that it doesn't really feel like a win. --- mojo/mojo_import.bzl | 7 ++++- mojo/mojo_library.bzl | 39 +++++++++++++++++++++-- mojo/private/utils.bzl | 54 ++++++++++++++++++++++++++++++++ mojo/providers.bzl | 2 ++ tests/main.mojo | 5 +++ tests/src_deps_a/a/BUILD.bazel | 7 +++++ tests/src_deps_a/a/__init__.mojo | 2 ++ tests/src_deps_b/b/BUILD.bazel | 8 +++++ tests/src_deps_b/b/__init__.mojo | 4 +++ 9 files changed, 125 insertions(+), 3 deletions(-) create mode 100644 tests/main.mojo create mode 100644 tests/src_deps_a/a/BUILD.bazel create mode 100644 tests/src_deps_a/a/__init__.mojo create mode 100644 tests/src_deps_b/b/BUILD.bazel create mode 100644 tests/src_deps_b/b/__init__.mojo diff --git a/mojo/mojo_import.bzl b/mojo/mojo_import.bzl index efdfbc6..254e17d 100644 --- a/mojo/mojo_import.bzl +++ b/mojo/mojo_import.bzl @@ -1,16 +1,21 @@ """Import a precompiled mojo file for use in other Mojo targets.""" load("//mojo:providers.bzl", "MojoInfo") -load("//mojo/private:utils.bzl", "collect_mojoinfo") +load("//mojo/private:utils.bzl", "collect_mojoinfo", "collect_src_mojoinfo") def _mojo_import_impl(ctx): mojo_deps = ctx.files.mojodeps import_paths, transitive_mojodeps = collect_mojoinfo(ctx.attr.deps) + src_import_paths, src_mojodeps = collect_src_mojoinfo(ctx.attr.deps) return [ DefaultInfo(files = depset(mojo_deps, transitive = [transitive_mojodeps])), MojoInfo( import_paths = depset([pkg.dirname for pkg in mojo_deps], transitive = [import_paths]), mojodeps = depset([pkg for pkg in mojo_deps], transitive = [transitive_mojodeps]), + # A precompiled import has no sources, so source-mode consumers fall + # back to its .mojoc files and their directories. + src_import_paths = depset([pkg.dirname for pkg in mojo_deps], transitive = [src_import_paths]), + src_mojodeps = depset(mojo_deps, transitive = [src_mojodeps]), ), ] diff --git a/mojo/mojo_library.bzl b/mojo/mojo_library.bzl index 7393bc7..a075dd8 100644 --- a/mojo/mojo_library.bzl +++ b/mojo/mojo_library.bzl @@ -3,7 +3,7 @@ other Mojo targets.""" load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo") load("//mojo:providers.bzl", "MojoInfo") -load("//mojo/private:utils.bzl", "MOJO_EXTENSIONS", "collect_mojoinfo", "is_exec_config") +load("//mojo/private:utils.bzl", "MOJO_EXTENSIONS", "collect_mojoinfo", "collect_src_mojoinfo", "is_exec_config", "source_import_path") def _format_include(arg): return ["-I", arg.dirname] @@ -29,6 +29,16 @@ def _mojo_library_implementation(ctx): import_paths, transitive_mojodeps = collect_mojoinfo(ctx.attr.deps + mojo_toolchain.implicit_deps) root_directory = ctx.files.srcs[0].dirname + # The directory to put on the import path so this library is importable by + # its module name when consumed from source (see source_import_path). + own_src_import_path = source_import_path(ctx.files.srcs, root_directory) + + # Source-mode dependencies: compile against their sources directly rather + # than their precompiled .mojoc files. Both fields are depsets, so this + # already carries the full transitive closure of every src_dep. + src_import_paths, src_mojodeps = collect_src_mojoinfo(ctx.attr.src_deps) + args.add_all(src_import_paths, before_each = "-I") + file_args = ctx.actions.args() for file in ctx.files.srcs: if not file.dirname.startswith(root_directory): @@ -44,9 +54,17 @@ def _mojo_library_implementation(ctx): file_args.add_all(transitive_mojodeps, map_each = _format_include) file_args.add(root_directory) + ctx.actions.run( executable = mojo_toolchain.mojo, - inputs = depset(ctx.files.srcs + ctx.files.additional_compiler_inputs, transitive = [transitive_mojodeps]), + inputs = depset( + # Direct items must be Files: use ctx.files.* (Files), never + # ctx.attr.* (Targets) — mixing the two is what breaks the depset. + ctx.files.srcs + ctx.files.additional_compiler_inputs, + # Merge other Files via transitive depsets. src_mojodeps is the + # transitive closure of every src_dep's sources. + transitive = [transitive_mojodeps, src_mojodeps], + ), tools = mojo_toolchain.all_tools, outputs = precompile_outputs, arguments = [args, file_args], @@ -76,6 +94,11 @@ def _mojo_library_implementation(ctx): MojoInfo( import_paths = depset([mojo_precmp_file.dirname], transitive = [import_paths]), mojodeps = depset([mojo_precmp_file], transitive = [transitive_mojodeps]), + # Always populate the source-mode fields (regardless of which mode + # this target was built with) so downstream src_deps can pull this + # target's sources and import path, plus everything it depends on. + src_import_paths = depset([own_src_import_path], transitive = [src_import_paths]), + src_mojodeps = depset(ctx.files.srcs, transitive = [src_mojodeps]), ), OutputGroupInfo(**output_group_kwargs), ] @@ -111,6 +134,18 @@ precompile' since it does not accept many flags. "deps": attr.label_list( providers = [MojoInfo], ), + "src_deps": attr.label_list( + providers = [MojoInfo], + doc = """\ +Like deps, but imports these dependencies from source instead of from +precompiled files. + +Providing dependencies here may improve clean build time (as it allows parallel +builds), but at the detriment of incremental build time. Note that these do not +provide transitive dependencies. It is advised to not heavily mix and match +this with `deps`. + """, + ), "data": attr.label_list(), "_mojo_precompile_copts": attr.label( default = Label("//:mojo_precompile_copt"), diff --git a/mojo/private/utils.bzl b/mojo/private/utils.bzl index 1ef7e8a..c7c0364 100644 --- a/mojo/private/utils.bzl +++ b/mojo/private/utils.bzl @@ -24,6 +24,60 @@ def collect_mojoinfo(deps): return depset(transitive = import_paths), depset(transitive = mojodeps) +def collect_src_mojoinfo(deps): + """Collect the source-mode (src_deps) data from the passed dependencies. + + Unlike collect_mojoinfo, this gathers the raw sources and their import + paths so the consuming target can compile against them directly instead of + against precompiled .mojoc files. + + Args: + deps: A list of dependencies to collect MojoInfo from. + + Returns: + A tuple (src_import_paths, src_mojodeps) where src_import_paths is a + depset of directory strings and src_mojodeps is a depset of Files. + """ + src_import_paths = [] + src_mojodeps = [] + for dep in deps: + if MojoInfo in dep: + info = dep[MojoInfo] + + # Tolerate MojoInfo producers (e.g. third-party rules) that predate + # the source-mode fields by treating them as empty. + src_import_paths.append(getattr(info, "src_import_paths", depset())) + src_mojodeps.append(getattr(info, "src_mojodeps", depset())) + + return depset(transitive = src_import_paths), depset(transitive = src_mojodeps) + +def source_import_path(srcs, root_directory): + """Compute the -I directory that makes a library importable from source. + + Precompiled mode flattens a library into a single .mojoc and puts the + directory holding it on the import path, so the unit is always a file. Source + mode hands the compiler the real tree, where Mojo distinguishes two shapes: + + * a package -- a directory containing __init__.mojo -- is importable when + its *parent* directory is on the import path, and + * a flat module -- a single foo.mojo -- is importable when *its own* + directory is on the import path. + + In both cases the library is imported as its on-disk name, which is expected + to match ctx.label.name (the same name precompiled mode gives .mojoc). + + Args: + srcs: The source Files of the library (ctx.files.srcs). + root_directory: The directory containing the sources (srcs[0].dirname). + + Returns: + The directory string to pass with -I. + """ + is_package = any([f.basename == "__init__.mojo" for f in srcs]) + if is_package: + return root_directory.rsplit("/", 1)[0] if "/" in root_directory else "." + return root_directory + def is_exec_config(ctx): """Determines whether the current configuration is an exec configuration. diff --git a/mojo/providers.bzl b/mojo/providers.bzl index f7d8927..c027cd5 100644 --- a/mojo/providers.bzl +++ b/mojo/providers.bzl @@ -5,6 +5,8 @@ MojoInfo = provider( fields = { "import_paths": "Directories that should be passed with -I to mojo", "mojodeps": "The precompiled mojo files required by the target", + "src_import_paths": "depset of directories (strings) that should be passed with -I to mojo when building with src_deps", + "src_mojodeps": "depset of source Files of this target and all of its (transitive) dependencies", }, ) diff --git a/tests/main.mojo b/tests/main.mojo new file mode 100644 index 0000000..7fe61e6 --- /dev/null +++ b/tests/main.mojo @@ -0,0 +1,5 @@ +from a import get_num_a +from b import get_num_b + +def main() -> None: + assert get_num_a() + get_num_b() == 4 diff --git a/tests/src_deps_a/a/BUILD.bazel b/tests/src_deps_a/a/BUILD.bazel new file mode 100644 index 0000000..d5cc1f2 --- /dev/null +++ b/tests/src_deps_a/a/BUILD.bazel @@ -0,0 +1,7 @@ +load("//mojo:mojo_library.bzl", "mojo_library") + +mojo_library( + name = "a", + srcs = ["__init__.mojo"], + visibility = ["//tests:__subpackages__"], +) diff --git a/tests/src_deps_a/a/__init__.mojo b/tests/src_deps_a/a/__init__.mojo new file mode 100644 index 0000000..8dcd0ab --- /dev/null +++ b/tests/src_deps_a/a/__init__.mojo @@ -0,0 +1,2 @@ +def get_num_a() -> Int: + return 1 diff --git a/tests/src_deps_b/b/BUILD.bazel b/tests/src_deps_b/b/BUILD.bazel new file mode 100644 index 0000000..e8a9c03 --- /dev/null +++ b/tests/src_deps_b/b/BUILD.bazel @@ -0,0 +1,8 @@ +load("//mojo:mojo_library.bzl", "mojo_library") + +mojo_library( + name = "b", + srcs = ["__init__.mojo"], + src_deps = ["//tests/src_deps_a/a"], + visibility = ["//tests:__subpackages__"], +) diff --git a/tests/src_deps_b/b/__init__.mojo b/tests/src_deps_b/b/__init__.mojo new file mode 100644 index 0000000..8c3d6e9 --- /dev/null +++ b/tests/src_deps_b/b/__init__.mojo @@ -0,0 +1,4 @@ +from a import get_num_a + +def get_num_b() -> Int: + return get_num_a() + 2