Map files to their corresponding compilation steps

src/DocumentStore.zig

@@ -113,10 +113,12 @@ pub const BuildFile = struct {
         const build_config = self.tryLockConfig() orelse return false;
         defer self.unlockConfig();
 
-        try package_uris.ensureUnusedCapacity(allocator, build_config.packages.len);
-        for (build_config.packages) |package| {
-            package_uris.appendAssumeCapacity(try URI.fromPath(allocator, package.path));
+        for (build_config.compilation_unit_info) |compilation_unit| {
+            for (compilation_unit.packages) |package| {
+                try package_uris.append(allocator, try URI.fromPath(allocator, package.path));
+            }
         }
+
         return true;
     }
 
@@ -140,18 +142,19 @@ pub const BuildFile = struct {
         const build_config = self.tryLockConfig() orelse return false;
         defer self.unlockConfig();
 
-        try include_paths.ensureUnusedCapacity(allocator, build_config.include_dirs.len);
-        for (build_config.include_dirs) |include_path| {
-            const absolute_path = if (std.fs.path.isAbsolute(include_path))
-                try allocator.dupe(u8, include_path)
-            else blk: {
-                const build_file_dir = std.fs.path.dirname(self.uri).?;
-                const build_file_path = try URI.parse(allocator, build_file_dir);
-                defer allocator.free(build_file_path);
-                break :blk try std.fs.path.join(allocator, &.{ build_file_path, include_path });
-            };
-
-            include_paths.appendAssumeCapacity(absolute_path);
+        for (build_config.compilation_unit_info) |compilation_unit| {
+            for (compilation_unit.include_dirs) |include_path| {
+                const absolute_path = if (std.fs.path.isAbsolute(include_path))
+                    try allocator.dupe(u8, include_path)
+                else blk: {
+                    const build_file_dir = std.fs.path.dirname(self.uri).?;
+                    const build_file_path = try URI.parse(allocator, build_file_dir);
+                    defer allocator.free(build_file_path);
+                    break :blk try std.fs.path.join(allocator, &.{ build_file_path, include_path });
+                };
+
+                try include_paths.append(allocator, absolute_path);
+            }
         }
         return true;
     }
@@ -1303,8 +1306,10 @@ fn loadBuildConfiguration(self: *DocumentStore, build_file_uri: Uri) !std.json.P
     ) catch return error.InvalidBuildConfig;
     errdefer build_config.deinit();
 
-    for (build_config.value.packages) |*pkg| {
-        pkg.path = try std.fs.path.resolve(build_config.arena.allocator(), &.{ build_file_path, "..", pkg.path });
+    for (build_config.value.compilation_unit_info) |compilation_unit_info| {
+        for (compilation_unit_info.packages) |*pkg| {
+            pkg.path = try std.fs.path.resolve(build_config.arena.allocator(), &.{ build_file_path, "..", pkg.path });
+        }
     }
 
     return build_config;
@@ -1702,9 +1707,10 @@ pub fn collectCMacros(
             const build_config = build_file.tryLockConfig() orelse break :blk false;
             defer build_file.unlockConfig();
 
-            try c_macros.ensureUnusedCapacity(allocator, build_config.c_macros.len);
-            for (build_config.c_macros) |c_macro| {
-                c_macros.appendAssumeCapacity(try allocator.dupe(u8, c_macro));
+            for (build_config.compilation_unit_info) |compilation_unit| {
+                for (compilation_unit.c_macros) |c_macro| {
+                    try c_macros.append(allocator, try allocator.dupe(u8, c_macro));
+                }
             }
             break :blk true;
         },
@@ -1907,9 +1913,30 @@ pub fn uriFromImportStr(self: *DocumentStore, allocator: std.mem.Allocator, hand
             const build_config = build_file.tryLockConfig() orelse break :blk;
             defer build_file.unlockConfig();
 
-            for (build_config.packages) |pkg| {
-                if (std.mem.eql(u8, import_str, pkg.name)) {
-                    return try URI.fromPath(allocator, pkg.path);
+            const current_file_uri = handle.uri;
+
+            var current_cu: ?BuildConfig.CompilationUnit = null;
+
+            find_cu: for (build_config.compilation_unit_info) |compilation_unit| {
+                for (compilation_unit.files) |file_path| {
+                    const file_uri = try URI.fromPath(allocator, file_path);
+                    defer allocator.free(file_uri);
+
+                    const relative = std.fs.path.relative(allocator, file_uri, current_file_uri) catch continue;
+                    defer allocator.free(relative);
+
+                    if (relative.len == 0) {
+                        current_cu = compilation_unit;
+                        break :find_cu;
+                    }
+                }
+            }
+
+            if (current_cu) |compilation_unit| {
+                for (compilation_unit.packages) |pkg| {
+                    if (std.mem.eql(u8, import_str, pkg.name)) {
+                        return try URI.fromPath(allocator, pkg.path);
+                    }
                 }
             }
         } else if (isBuildFile(handle.uri)) blk: {

src/build_runner/0.14.0.zig

@@ -1097,6 +1097,147 @@ fn extractBuildInformation(
             }
         }
 
+        /// Ported from src/offsets.zig
+        /// Support formats:
+        /// - `foo`
+        /// - `@"foo"`
+        /// - `@foo`
+        fn identifierIndexToLoc(text: [:0]const u8, source_index: usize) struct {
+            start: usize,
+            end: usize,
+        } {
+            if (text[source_index] == '@' and text[source_index + 1] == '"') {
+                const start_index = source_index + 2;
+                var index: usize = start_index;
+                while (true) : (index += 1) {
+                    switch (text[index]) {
+                        '\n' => break,
+                        '"' => {
+                            // continue on e.g. `@"\""` but not `@"\\"`
+                            if (text[index - 1] == '\\' and text[index - 2] != '\\') continue;
+                            // include the closing quote
+                            break;
+                        },
+                        else => {},
+                    }
+                }
+                return .{ .start = start_index, .end = index };
+            } else {
+                const start: usize = source_index + @intFromBool(text[source_index] == '@');
+                var index = start;
+                while (true) : (index += 1) {
+                    switch (text[index]) {
+                        'a'...'z', 'A'...'Z', '_', '0'...'9' => {},
+                        else => break,
+                    }
+                }
+                return .{ .start = start, .end = index };
+            }
+        }
+
+        /// Ported from src/offsets.zig
+        /// Collects all `@import`'s we can find into a slice of import paths (without quotes).
+        fn collectImports(allocator: std.mem.Allocator, tree: std.zig.Ast) error{OutOfMemory}!std.ArrayListUnmanaged([]const u8) {
+            var imports: std.ArrayListUnmanaged([]const u8) = .empty;
+            errdefer imports.deinit(allocator);
+
+            for (0..tree.tokens.len) |i| {
+                if (tree.tokenTag(@intCast(i)) != .builtin)
+                    continue;
+
+                const loc = identifierIndexToLoc(tree.source, tree.tokenStart(@intCast(i)));
+                const name = tree.source[loc.start..loc.end];
+                if (!std.mem.eql(u8, name, "import")) continue;
+                if (!std.mem.startsWith(std.zig.Token.Tag, tree.tokens.items(.tag)[i + 1 ..], &.{ .l_paren, .string_literal, .r_paren })) continue;
+
+                const str = tree.tokenSlice(@intCast(i + 2));
+                try imports.append(allocator, str[1 .. str.len - 1]);
+            }
+
+            return imports;
+        }
+
+        fn getRelatedFiles(allocator: Allocator, root_source_file: ?std.Build.LazyPath) error{OutOfMemory}![]const []const u8 {
+            var files = std.StringArrayHashMap(void).init(allocator);
+            defer files.deinit();
+            errdefer for (files.keys()) |key| {
+                allocator.free(key);
+            };
+
+            var stack = std.ArrayList([]const u8).init(allocator);
+            defer stack.deinit();
+            errdefer for (stack.items) |file_path| {
+                allocator.free(file_path);
+            };
+
+            if (root_source_file == null or switch (root_source_file.?) {
+                .src_path => false,
+                else => true,
+            }) {
+                return try allocator.dupe([]const u8, files.keys());
+            }
+
+            const build_root = root_source_file.?.src_path.owner.build_root.path.?;
+
+            {
+                const resolved_root_source_file = try std.fs.path.resolve(allocator, &.{
+                    build_root,
+                    root_source_file.?.src_path.sub_path,
+                });
+                errdefer allocator.free(resolved_root_source_file);
+
+                try stack.append(resolved_root_source_file);
+            }
+
+            while (stack.pop()) |file_path| {
+                defer allocator.free(file_path);
+
+                const gop = blk: {
+                    const resolved_file_path = try std.fs.path.resolve(allocator, &.{
+                        build_root,
+                        file_path,
+                    });
+                    errdefer allocator.free(resolved_file_path);
+
+                    break :blk try files.getOrPut(resolved_file_path);
+                };
+
+                if (!gop.found_existing) {
+                    const file = std.fs.cwd().openFile(file_path, .{}) catch continue;
+                    defer file.close();
+
+                    const content = file.readToEndAllocOptions(allocator, std.math.maxInt(usize), null, .@"1", 0) catch continue;
+                    defer allocator.free(content);
+
+                    var tree = try std.zig.Ast.parse(allocator, content, .zig);
+                    defer tree.deinit(allocator);
+
+                    var imports = try collectImports(allocator, tree);
+                    defer {
+                        for (imports.items) |item| {
+                            allocator.free(item);
+                        }
+                        imports.deinit(allocator);
+                    }
+
+                    for (imports.items) |item| {
+                        if (std.mem.endsWith(u8, item, ".zig")) {
+                            const resolved_path = try std.fs.path.resolve(allocator, &.{
+                                file_path,
+                                "..",
+                                item,
+                            });
+                            errdefer allocator.free(resolved_path);
+
+                            try stack.append(resolved_path);
+                        }
+                    }
+                }
+            }
+
+            return try allocator.dupe([]const u8, files.keys());
+        }
+
         fn processItem(
             allocator: Allocator,
             module: *std.Build.Module,
@@ -1206,18 +1347,32 @@ fn extractBuildInformation(
         run,
     );
 
-    var include_dirs: std.StringArrayHashMapUnmanaged(void) = .{};
-    defer include_dirs.deinit(gpa);
-
-    var c_macros: std.StringArrayHashMapUnmanaged(void) = .{};
-    defer c_macros.deinit(gpa);
-
-    var packages: Packages = .{ .allocator = gpa };
-    defer packages.deinit();
+    var compilation_unit_info = std.AutoArrayHashMap(*std.Build.Module, BuildConfig.CompilationUnit).init(gpa);
+    defer {
+        for (compilation_unit_info.values()) |value| {
+            gpa.free(value.include_dirs);
+            gpa.free(value.c_macros);
+            gpa.free(value.packages);
+            for (value.files) |file| {
+                gpa.free(file);
+            }
+            gpa.free(value.files);
+        }
+        compilation_unit_info.deinit();
+    }
 
     // extract packages and include paths
     {
         for (steps.keys()) |step| {
+            var include_dirs: std.StringArrayHashMapUnmanaged(void) = .{};
+            defer include_dirs.deinit(gpa);
+
+            var c_macros: std.StringArrayHashMapUnmanaged(void) = .{};
+            defer c_macros.deinit(gpa);
+
+            var packages: Packages = .{ .allocator = gpa };
+            defer packages.deinit();
+
             const compile = step.cast(Step.Compile) orelse continue;
             const graph = compile.root_module.getGraph();
             try helper.processItem(gpa, compile.root_module, compile, "root", &packages, &include_dirs, &c_macros);
@@ -1226,16 +1381,61 @@ fn extractBuildInformation(
                     try helper.processItem(gpa, import, null, name, &packages, &include_dirs, &c_macros);
                 }
             }
+
+            const files = try helper.getRelatedFiles(gpa, compile.root_module.root_source_file);
+            errdefer {
+                for (files) |file_path| {
+                    gpa.free(file_path);
+                }
+                gpa.free(files);
+            }
+
+            const gop = try compilation_unit_info.getOrPut(compile.root_module);
+            if (!gop.found_existing) {
+                gop.value_ptr.* = .{
+                    .include_dirs = try gpa.dupe([]const u8, include_dirs.keys()),
+                    .c_macros = try gpa.dupe([]const u8, c_macros.keys()),
+                    .packages = try packages.toPackageList(),
+                    .files = files,
+                };
+            }
         }
 
         for (b.modules.values()) |root_module| {
+            var include_dirs: std.StringArrayHashMapUnmanaged(void) = .{};
+            defer include_dirs.deinit(gpa);
+
+            var c_macros: std.StringArrayHashMapUnmanaged(void) = .{};
+            defer c_macros.deinit(gpa);
+
+            var packages: Packages = .{ .allocator = gpa };
+            defer packages.deinit();
+
             const graph = root_module.getGraph();
             try helper.processItem(gpa, root_module, null, "root", &packages, &include_dirs, &c_macros);
             for (graph.modules) |module| {
                 for (module.import_table.keys(), module.import_table.values()) |name, import| {
                     try helper.processItem(gpa, import, null, name, &packages, &include_dirs, &c_macros);
                 }
             }
+
+            const files = try helper.getRelatedFiles(gpa, root_module.root_source_file);
+            errdefer {
+                for (files) |file_path| {
+                    gpa.free(file_path);
+                }
+                gpa.free(files);
+            }
+
+            const gop = try compilation_unit_info.getOrPut(root_module);
+            if (!gop.found_existing) {
+                gop.value_ptr.* = .{
+                    .include_dirs = try gpa.dupe([]const u8, include_dirs.keys()),
+                    .c_macros = try gpa.dupe([]const u8, c_macros.keys()),
+                    .packages = try packages.toPackageList(),
+                    .files = files,
+                };
+            }
         }
     }
 
@@ -1279,11 +1479,9 @@ fn extractBuildInformation(
     try std.json.stringify(
         BuildConfig{
             .deps_build_roots = deps_build_roots.items,
-            .packages = try packages.toPackageList(),
-            .include_dirs = include_dirs.keys(),
+            .compilation_unit_info = compilation_unit_info.values(),
             .top_level_steps = b.top_level_steps.keys(),
             .available_options = available_options,
-            .c_macros = c_macros.keys(),
         },
         .{
             .whitespace = .indent_2,

src/build_runner/shared.zig

@@ -7,17 +7,21 @@ const need_bswap = native_endian != .little;
 
 pub const BuildConfig = struct {
     deps_build_roots: []DepsBuildRoots,
-    packages: []Package,
-    include_dirs: []const []const u8,
+    compilation_unit_info: []CompilationUnit,
     top_level_steps: []const []const u8,
     available_options: std.json.ArrayHashMap(AvailableOption),
-    c_macros: []const []const u8 = &.{},
 
     pub const DepsBuildRoots = Package;
     pub const Package = struct {
         name: []const u8,
         path: []const u8,
     };
+    pub const CompilationUnit = struct {
+        include_dirs: []const []const u8,
+        c_macros: []const []const u8,
+        packages: []Package,
+        files: []const []const u8,
+    };
     pub const AvailableOption = std.meta.FieldType(std.meta.FieldType(std.Build, .available_options_map).KV, .value);
 };