Add src_deps option for mojo_library

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]),
         ),
     ]
 

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"),

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 <name>.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 <name>.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.
 

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",
     },
 )
 

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

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__"],
+)

tests/src_deps_a/a/__init__.mojo

@@ -0,0 +1,2 @@
+def get_num_a() -> Int:
+    return 1

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__"],
+)

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