Merge pull request #83 from alexnask/master

Support custom package imports

README.md

@@ -51,6 +51,7 @@ The following options are currently available.
 | `enable_snippets` | `bool` | `false` | Enables snippet completion, set to false for compatibility with language clients that do not support snippets (such as ale). |
 | `zig_lib_path` | `?[]const u8` | `null` | zig library path, e.g. `/path/to/zig/lib/zig`, used to analyze std library imports. |
 | `warn_style` | `bool` | `false` | Enables warnings for style *guideline* mismatches |
+| `build_runner_path` | `?[]const u8` | `null` | Path to the build_runner.zig file provided by zls. This option must be present in one of the global configuration files to have any effect. `null` is equivalent to `${executable_directory}/build_runner.zig` |
 
 ## Usage
 

build.zig

@@ -29,10 +29,14 @@ pub fn build(b: *std.build.Builder) !void {
         b.option(bool, "allocation_info", "Enable use of debugging allocator and info logging.") orelse false,
     );
 
+    exe.addPackage(.{ .name = "known-folders", .path = "src/known-folders/known-folders.zig" });
+
     exe.setTarget(target);
     exe.setBuildMode(mode);
     exe.install();
 
+    b.installFile("src/special/build_runner.zig", "bin/build_runner.zig");
+
     const run_cmd = exe.run();
     run_cmd.step.dependOn(b.getInstallStep());
 

src/config.zig

@@ -9,3 +9,7 @@ zig_lib_path: ?[]const u8 = null,
 /// Whether to pay attention to style issues. This is opt-in since the style
 /// guide explicitly states that the style info provided is a guideline only.
 warn_style: bool = false,
+
+/// Path to the build_runner.zig file. This option must be present in one of
+/// the global configuration directories to have any effect.
+build_runner_path: ?[]const u8 = null,

src/document_store.zig

@@ -5,12 +5,26 @@ const analysis = @import("analysis.zig");
 
 const DocumentStore = @This();
 
+const BuildFile = struct {
+    const Pkg = struct {
+        name: []const u8,
+        uri: []const u8,
+    };
+
+    refs: usize,
+    uri: []const u8,
+    packages: std.ArrayListUnmanaged(Pkg),
+};
+
 pub const Handle = struct {
     document: types.TextDocument,
     count: usize,
     import_uris: std.ArrayList([]const u8),
     tree: *std.zig.ast.Tree,
 
+    associated_build_file: ?*BuildFile,
+    is_build_file: ?*BuildFile,
+
     pub fn uri(handle: Handle) []const u8 {
         return handle.document.uri;
     }
@@ -18,16 +32,111 @@ pub const Handle = struct {
 
 allocator: *std.mem.Allocator,
 handles: std.StringHashMap(*Handle),
+has_zig: bool,
+build_files: std.ArrayListUnmanaged(*BuildFile),
+build_runner_path: []const u8,
 
-pub fn init(self: *DocumentStore, allocator: *std.mem.Allocator) !void {
+pub fn init(
+    self: *DocumentStore,
+    allocator: *std.mem.Allocator,
+    has_zig: bool,
+    build_runner_path: []const u8,
+) !void {
     self.allocator = allocator;
     self.handles = std.StringHashMap(*Handle).init(allocator);
-    errdefer self.handles.deinit();
+    self.has_zig = has_zig;
+    self.build_files = .{};
+    self.build_runner_path = build_runner_path;
+}
+
+const LoadPackagesContext = struct {
+    build_file: *BuildFile,
+    allocator: *std.mem.Allocator,
+    build_runner_path: []const u8,
+};
+
+fn loadPackages(context: LoadPackagesContext) !void {
+    const allocator = context.allocator;
+    const build_file = context.build_file;
+    const build_runner_path = context.build_runner_path;
+
+    const directory_path = try URI.parse(allocator, build_file.uri[0 .. build_file.uri.len - "build.zig".len]);
+    defer allocator.free(directory_path);
+
+    const target_path = try std.fs.path.resolve(allocator, &[_][]const u8{ directory_path, "build_runner.zig" });
+    defer allocator.free(target_path);
+
+    // For example, instead of testing if a file exists and then opening it, just
+    // open it and handle the error for file not found.
+    var file_exists = true;
+    check_file_exists: {
+        var fhandle = std.fs.cwd().openFile(target_path, .{.read = true, .write = false }) catch |err| switch (err) {
+            error.FileNotFound => {
+                file_exists = false;
+                break :check_file_exists;
+            },
+            else => break :check_file_exists,
+        };
+        fhandle.close();
+    }
+
+    if (file_exists) return error.BuildRunnerFileExists;
+
+    try std.fs.copyFileAbsolute(build_runner_path, target_path, .{});
+    defer std.fs.deleteFileAbsolute(target_path) catch {};
+
+    const zig_run_result = try std.ChildProcess.exec(.{
+        .allocator = allocator,
+        .argv = &[_][]const u8{ "zig", "run", "build_runner.zig" },
+        .cwd = directory_path,
+    });
+
+    defer {
+        allocator.free(zig_run_result.stdout);
+        allocator.free(zig_run_result.stderr);
+    }
+
+    switch (zig_run_result.term) {
+        .Exited => |exit_code| {
+            if (exit_code == 0) {
+                std.debug.warn("Finished zig run for build file {}\n", .{build_file.uri});
+
+                for (build_file.packages.items) |old_pkg| {
+                    allocator.free(old_pkg.name);
+                    allocator.free(old_pkg.uri);
+                }
+
+                build_file.packages.shrink(allocator, 0);
+                var line_it = std.mem.split(zig_run_result.stdout, "\n");
+                while (line_it.next()) |line| {
+                    if (std.mem.indexOfScalar(u8, line, '\x00')) |zero_byte_idx| {
+                        const name = line[0..zero_byte_idx];
+                        const rel_path = line[zero_byte_idx + 1 ..];
+
+                        const pkg_abs_path = try std.fs.path.resolve(allocator, &[_][]const u8{ directory_path, rel_path });
+                        defer allocator.free(pkg_abs_path);
+
+                        const pkg_uri = try URI.fromPath(allocator, pkg_abs_path);
+                        errdefer allocator.free(pkg_uri);
+
+                        const duped_name = try std.mem.dupe(allocator, u8, name);
+                        errdefer allocator.free(duped_name);
+
+                        (try build_file.packages.addOne(allocator)).* = .{
+                            .name = duped_name,
+                            .uri = pkg_uri,
+                        };
+                    }
+                }
+            }
+        },
+        else => {},
+    }
 }
 
 /// This function asserts the document is not open yet and takes ownership
 /// of the uri and text passed in.
-fn newDocument(self: *DocumentStore, uri: []const u8, text: []u8) !*Handle {
+fn newDocument(self: *DocumentStore, uri: []const u8, text: []u8) anyerror!*Handle {
     std.debug.warn("Opened document: {}\n", .{uri});
 
     var handle = try self.allocator.create(Handle);
@@ -42,8 +151,77 @@ fn newDocument(self: *DocumentStore, uri: []const u8, text: []u8) !*Handle {
             .mem = text,
         },
         .tree = try std.zig.parse(self.allocator, text),
+        .associated_build_file = null,
+        .is_build_file = null,
     };
 
+    // TODO: Better logic for detecting std or subdirectories?
+    const in_std = std.mem.indexOf(u8, uri, "/std/") != null;
+    if (self.has_zig and std.mem.endsWith(u8, uri, "/build.zig") and !in_std) {
+        std.debug.warn("Document is a build file, extracting packages...\n", .{});
+        // This is a build file.
+        var build_file = try self.allocator.create(BuildFile);
+        errdefer self.allocator.destroy(build_file);
+
+        build_file.* = .{
+            .refs = 1,
+            .uri = try std.mem.dupe(self.allocator, u8, uri),
+            .packages = .{},
+        };
+
+        try self.build_files.append(self.allocator, build_file);
+        handle.is_build_file = build_file;
+
+        // TODO: Do this in a separate thread?
+        // It can take quite long.
+        loadPackages(.{
+            .build_file = build_file,
+            .allocator = self.allocator,
+            .build_runner_path = self.build_runner_path,
+        }) catch {
+            std.debug.warn("Failed to load packages of build file {}\n", .{build_file.uri});
+        };
+    } else if (self.has_zig and !in_std) associate_build_file: {
+        // Look into build files to see if we already have one that fits
+        for (self.build_files.items) |build_file| {
+            const build_file_base_uri = build_file.uri[0 .. std.mem.lastIndexOfScalar(u8, build_file.uri, '/').? + 1];
+
+            if (std.mem.startsWith(u8, uri, build_file_base_uri)) {
+                std.debug.warn("Found an associated build file: {}\n", .{build_file.uri});
+                build_file.refs += 1;
+                handle.associated_build_file = build_file;
+                break :associate_build_file;
+            }
+        }
+        // Otherwise, try to find a build file.
+        var curr_path = try URI.parse(self.allocator, uri);
+        defer self.allocator.free(curr_path);
+        while (true) {
+            if (curr_path.len == 0) break :associate_build_file;
+
+            if (std.mem.lastIndexOfScalar(u8, curr_path[0 .. curr_path.len - 1], std.fs.path.sep)) |idx| {
+                // This includes the last separator
+                curr_path = curr_path[0 .. idx + 1];
+                var candidate_path = try std.mem.concat(self.allocator, u8, &[_][]const u8{ curr_path, "build.zig" });
+                defer self.allocator.free(candidate_path);
+                // Try to open the file, read it and add the new document if we find it.
+                var file = std.fs.cwd().openFile(candidate_path, .{ .read = true, .write = false }) catch continue;
+                defer file.close();
+
+                const build_file_text = try file.inStream().readAllAlloc(self.allocator, std.math.maxInt(usize));
+                errdefer self.allocator.free(build_file_text);
+
+                const build_file_uri = try URI.fromPath(self.allocator, candidate_path);
+                errdefer self.allocator.free(build_file_uri);
+
+                const build_file_handle = try self.newDocument(build_file_uri, build_file_text);
+
+                handle.associated_build_file = build_file_handle.is_build_file;
+                break;
+            } else break :associate_build_file;
+        }
+    }
+
     try self.handles.putNoClobber(uri, handle);
     return handle;
 }
@@ -52,6 +230,9 @@ pub fn openDocument(self: *DocumentStore, uri: []const u8, text: []const u8) !*H
     if (self.handles.get(uri)) |entry| {
         std.debug.warn("Document already open: {}, incrementing count\n", .{uri});
         entry.value.count += 1;
+        if (entry.value.is_build_file) |build_file| {
+            build_file.refs += 1;
+        }
         std.debug.warn("New count: {}\n", .{entry.value.count});
         return entry.value;
     }
@@ -64,9 +245,41 @@ pub fn openDocument(self: *DocumentStore, uri: []const u8, text: []const u8) !*H
     return try self.newDocument(duped_uri, duped_text);
 }
 
+fn decrementBuildFileRefs(self: *DocumentStore, build_file_ptr: *?*BuildFile) void {
+    const build_file = build_file_ptr.*.?;
+
+    build_file.refs -= 1;
+    if (build_file.refs == 0) {
+        std.debug.warn("Freeing build file {}\n", .{build_file.uri});
+        // Free the build file, set the pointer to null.
+        for (build_file.packages.items) |pkg| {
+            self.allocator.free(pkg.name);
+            self.allocator.free(pkg.uri);
+        }
+
+        build_file.packages.deinit(self.allocator);
+
+        // Decrement count of the document since one count comes
+        // from the build file existing.
+        self.decrementCount(build_file.uri);    
+        self.allocator.free(build_file.uri);
+
+        // Remove the build file from the array list
+        _ = self.build_files.swapRemove(std.mem.indexOfScalar(*BuildFile, self.build_files.items, build_file).?);
+
+        self.allocator.destroy(build_file);
+        build_file_ptr.* = null;
+    }
+}
+
 fn decrementCount(self: *DocumentStore, uri: []const u8) void {
     if (self.handles.get(uri)) |entry| {
         entry.value.count -= 1;
+
+        if (entry.value.associated_build_file != null) {
+            self.decrementBuildFileRefs(&entry.value.associated_build_file);
+        }
+
         if (entry.value.count > 0)
             return;
 
@@ -214,6 +427,15 @@ pub fn applyChanges(
     }
 
     try self.refreshDocument(handle, zig_lib_path);
+    if (handle.is_build_file) |build_file| {
+        loadPackages(.{
+            .build_file = build_file,
+            .allocator = self.allocator,
+            .build_runner_path = self.build_runner_path,
+        }) catch {
+            std.debug.warn("Failed to load packages of build file {}\n", .{build_file.uri});
+        };
+    }
 }
 
 pub fn uriFromImportStr(
@@ -231,7 +453,14 @@ pub fn uriFromImportStr(
     } else if (std.mem.eql(u8, import_str, "builtin")) {
         return null; // TODO find the correct zig-cache folder
     } else if (!std.mem.endsWith(u8, import_str, ".zig")) {
-        return null; // TODO find packages based on build.zig
+        if (handle.associated_build_file) |build_file| {
+            for (build_file.packages.items) |pkg| {
+                if (std.mem.eql(u8, import_str, pkg.name)) {
+                    return try std.mem.dupe(allocator, u8, pkg.uri);
+                }
+            }
+        }
+        return null;
     } else {
         // Find relative uri
         const path = try URI.parse(allocator, handle.uri());
@@ -414,4 +643,15 @@ pub fn deinit(self: *DocumentStore) void {
     }
 
     self.handles.deinit();
+
+    for (self.build_files.items) |build_file| {
+        for (build_file.packages.items) |pkg| {
+            self.allocator.free(pkg.name);
+            self.allocator.free(pkg.uri);
+        }
+        self.allocator.free(build_file.uri);
+        self.allocator.destroy(build_file);
+    }
+
+    self.build_files.deinit(self.allocator);
 }

src/main.zig

@@ -870,7 +870,8 @@ pub fn main() anyerror!void {
     defer std.json.parseFree(Config, config, config_parse_options);
 
     config_read: {
-        const known_folders = @import("known-folders/known-folders.zig");
+        const known_folders = @import("known-folders");
+
         const res = try known_folders.getPath(allocator, .local_configuration);
         if (res) |local_config_path| {
             defer allocator.free(local_config_path);
@@ -880,15 +881,59 @@ pub fn main() anyerror!void {
             }
         }
 
-        var exec_dir_bytes: [std.fs.MAX_PATH_BYTES]u8 = undefined;
+        var exe_dir_bytes: [std.fs.MAX_PATH_BYTES]u8 = undefined;
-        const exec_dir_path = std.fs.selfExeDirPath(&exec_dir_bytes) catch break :config_read;
+        const exe_dir_path = std.fs.selfExeDirPath(&exe_dir_bytes) catch break :config_read;
 
-        if (loadConfig(exec_dir_path)) |conf| {
+        if (loadConfig(exe_dir_path)) |conf| {
             config = conf;
         }
     }
 
-    try document_store.init(allocator);
+    // Find the zig executable in PATH
+    var has_zig = false;
+
+    // TODO: Should we just spawn a child process that calls "zig version" or something
+    // and check that way?
+    find_zig: {
+        const env_path = std.process.getEnvVarOwned(allocator, "PATH") catch |err| switch (err) {
+            error.EnvironmentVariableNotFound => {
+                std.debug.warn("Could not get PATH.\n", .{});
+                break :find_zig;
+            },
+            else => return err,
+        };
+        defer allocator.free(env_path);
+
+        const exe_extension = @as(std.zig.CrossTarget, .{}).exeFileExt();
+        const zig_exe = try std.fmt.allocPrint(allocator, "zig{}", .{exe_extension});
+        defer allocator.free(zig_exe);
+
+        var it = std.mem.tokenize(env_path, &[_]u8{std.fs.path.delimiter});
+        while (it.next()) |path| {
+            const full_path = try std.fs.path.join(allocator, &[_][]const u8{
+                path,
+                zig_exe,
+            });
+            defer allocator.free(full_path);
+
+            var buf: [std.fs.MAX_PATH_BYTES]u8 = undefined;
+            const zig_path = std.os.realpath(full_path, &buf) catch continue;
+            std.debug.warn("Found zig in PATH: {}\n", .{zig_path});
+            has_zig = true;
+            break :find_zig;
+        }
+    }
+
+    if (config.build_runner_path) |build_runner_path| {
+        try document_store.init(allocator, has_zig, try std.mem.dupe(allocator, u8, build_runner_path));
+    } else {
+        var exe_dir_bytes: [std.fs.MAX_PATH_BYTES]u8 = undefined;
+        const exe_dir_path = try std.fs.selfExeDirPath(&exe_dir_bytes);
+
+        const build_runner_path = try std.fs.path.resolve(allocator, &[_][]const u8{ exe_dir_path, "build_runner.zig" });
+        try document_store.init(allocator, has_zig, build_runner_path);
+    }
+
     defer document_store.deinit();
 
     workspace_folder_configs = std.StringHashMap(?Config).init(allocator);

src/special/build_runner.zig

@@ -0,0 +1,53 @@
+const root = @import("build.zig");
+const std = @import("std");
+const io = std.io;
+const fmt = std.fmt;
+const Builder = std.build.Builder;
+const Pkg = std.build.Pkg;
+const InstallArtifactStep = std.build.InstallArtifactStep;
+const ArrayList = std.ArrayList;
+
+///! This is a modified build runner to extract information out of build.zig
+///! Modified from the std.special.build_runner
+
+pub fn main() !void {
+    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
+    defer arena.deinit();
+
+    const allocator = &arena.allocator;
+
+    const builder = try Builder.create(allocator, "", "", "");
+    defer builder.destroy();
+
+    try runBuild(builder);
+
+    const stdout_stream = io.getStdOut().outStream();
+
+    // TODO: We currently add packages from every LibExeObj step that the install step depends on.
+    //       Should we error out or keep one step or something similar?
+    // We also flatten them, we should probably keep the nested structure.
+    for (builder.getInstallStep().dependencies.items) |step| {
+        if (step.cast(InstallArtifactStep)) |install_exe| {
+            for (install_exe.artifact.packages.items) |pkg| {
+                try processPackage(stdout_stream, pkg);
+            }
+        }
+    }
+}
+
+fn processPackage(out_stream: var, pkg: Pkg) anyerror!void {
+    try out_stream.print("{}\x00{}\n", .{ pkg.name, pkg.path });
+    if (pkg.dependencies) |dependencies| {
+        for (dependencies) |dep| {
+            try processPackage(out_stream, dep);
+        }
+    }
+}
+
+fn runBuild(builder: *Builder) anyerror!void {
+    switch (@typeInfo(@TypeOf(root.build).ReturnType)) {
+        .Void => root.build(builder),
+        .ErrorUnion => try root.build(builder),
+        else => @compileError("expected return type of build to be 'void' or '!void'"),
+    }
+}

src/uri.zig

@@ -41,6 +41,13 @@ pub fn fromPath(allocator: *std.mem.Allocator, path: []const u8) ![]const u8 {
     return buf.toOwnedSlice();
 }
 
+pub const UriParseError = error {
+    UriBadScheme,
+    UriBadHexChar,
+    UriBadEscape,
+    OutOfMemory,
+};
+
 // Original code: https://github.com/andersfr/zig-lsp/blob/master/uri.zig
 fn parseHex(c: u8) !u8 {
     return switch (c) {