Merge pull request #19 from alexnask/imports

Basic import support

README.md

@@ -28,6 +28,7 @@ The following options are currently available.
 | Option | Type | Default value | What it Does |
 | --- | --- | --- | --- |
 | `enable_snippets` | `bool` | `true` | 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, used to analyze std library imports. |
 
 ## Usage
 

src/analysis.zig

@@ -1,5 +1,5 @@
 const std = @import("std");
-
+const AnalysisContext = @import("document_store.zig").AnalysisContext;
 const ast = std.zig.ast;
 
 /// REALLY BAD CODE, PLEASE DON'T USE THIS!!!!!!! (only for testing)
@@ -110,6 +110,7 @@ pub fn getFunctionSnippet(allocator: *std.mem.Allocator, tree: *ast.Tree, func:
 
         if (param_decl.var_args_token) |_| {
             try buffer.appendSlice("...");
+            continue;
         }
 
         var curr_tok = param_decl.type_node.firstToken();
@@ -154,20 +155,20 @@ pub fn isPascalCase(name: []const u8) bool {
 // ANALYSIS ENGINE
 
 /// Gets the child of node
-pub fn getChild(tree: *std.zig.ast.Tree, node: *std.zig.ast.Node, name: []const u8) ?*std.zig.ast.Node {
+pub fn getChild(tree: *ast.Tree, node: *ast.Node, name: []const u8) ?*ast.Node {
     var index: usize = 0;
     while (node.iterate(index)) |child| {
         switch (child.id) {
             .VarDecl => {
-                const vari = child.cast(std.zig.ast.Node.VarDecl).?;
+                const vari = child.cast(ast.Node.VarDecl).?;
                 if (std.mem.eql(u8, tree.tokenSlice(vari.name_token), name)) return child;
             },
             .FnProto => {
-                const func = child.cast(std.zig.ast.Node.FnProto).?;
+                const func = child.cast(ast.Node.FnProto).?;
                 if (func.name_token != null and std.mem.eql(u8, tree.tokenSlice(func.name_token.?), name)) return child;
             },
             .ContainerField => {
-                const field = child.cast(std.zig.ast.Node.ContainerField).?;
+                const field = child.cast(ast.Node.ContainerField).?;
                 if (std.mem.eql(u8, tree.tokenSlice(field.name_token), name)) return child;
             },
             else => {}
@@ -178,51 +179,77 @@ pub fn getChild(tree: *std.zig.ast.Tree, node: *std.zig.ast.Node, name: []const
 }
 
 /// Resolves the type of a node
-pub fn resolveTypeOfNode(tree: *std.zig.ast.Tree, node: *std.zig.ast.Node) ?*std.zig.ast.Node {
+pub fn resolveTypeOfNode(analysis_ctx: *AnalysisContext, node: *ast.Node) ?*ast.Node {
     switch (node.id) {
         .VarDecl => {
-            const vari = node.cast(std.zig.ast.Node.VarDecl).?;
-            return resolveTypeOfNode(tree, vari.type_node orelse vari.init_node.?) orelse null;
+            const vari = node.cast(ast.Node.VarDecl).?;
+            return resolveTypeOfNode(analysis_ctx, vari.type_node orelse vari.init_node.?) orelse null;
         },
         .FnProto => {
-            const func = node.cast(std.zig.ast.Node.FnProto).?;
+            const func = node.cast(ast.Node.FnProto).?;
             switch (func.return_type) {
-                .Explicit, .InferErrorSet => |return_type| {return resolveTypeOfNode(tree, return_type);}
+                .Explicit, .InferErrorSet => |return_type| return resolveTypeOfNode(analysis_ctx, return_type),
             }
         },
         .Identifier => {
-            if (getChild(tree, &tree.root_node.base, tree.getNodeSource(node))) |child| {
-                return resolveTypeOfNode(tree, child);
+            if (getChild(analysis_ctx.tree, &analysis_ctx.tree.root_node.base, analysis_ctx.tree.getNodeSource(node))) |child| {
+                return resolveTypeOfNode(analysis_ctx, child);
             } else return null;
         },
         .ContainerDecl => {
             return node;
         },
         .ContainerField => {
-            const field = node.cast(std.zig.ast.Node.ContainerField).?;
-            return resolveTypeOfNode(tree, field.type_expr.?);
+            const field = node.cast(ast.Node.ContainerField).?;
+            return resolveTypeOfNode(analysis_ctx, field.type_expr orelse return null);
         },
         .SuffixOp => {
-            const suffix_op = node.cast(std.zig.ast.Node.SuffixOp).?;
+            const suffix_op = node.cast(ast.Node.SuffixOp).?;
             switch (suffix_op.op) {
                 .Call => {
-                    return resolveTypeOfNode(tree, suffix_op.lhs.node);
+                    return resolveTypeOfNode(analysis_ctx, suffix_op.lhs.node);
                 },
                 else => {}
             }
         },
         .InfixOp => {
-            const infix_op = node.cast(std.zig.ast.Node.InfixOp).?;
+            const infix_op = node.cast(ast.Node.InfixOp).?;
             switch (infix_op.op) {
                 .Period => {
-                    var left = resolveTypeOfNode(tree, infix_op.lhs).?;
-                    if (nodeToString(tree, infix_op.rhs)) |string| {
-                        return getChild(tree, left, string);
-                    } else return null;
+                    // Save the child string from this tree since the tree may switch when processing
+                    // an import lhs.
+                    var rhs_str = nodeToString(analysis_ctx.tree, infix_op.rhs) orelse return null;
+                    // Use the analysis context temporary arena to store the rhs string.
+                    rhs_str = std.mem.dupe(&analysis_ctx.arena.allocator, u8, rhs_str) catch return null;
+                    const left = resolveTypeOfNode(analysis_ctx, infix_op.lhs) orelse return null;
+                    return getChild(analysis_ctx.tree, left, rhs_str);
                 },
                 else => {}
             }
         },
+        .PrefixOp => {
+            const prefix_op = node.cast(ast.Node.PrefixOp).?;
+            switch (prefix_op.op) {
+                .PtrType => {
+                    return resolveTypeOfNode(analysis_ctx, prefix_op.rhs);
+                },
+                else => {}
+            }
+        },
+        .BuiltinCall => {
+            const builtin_call = node.cast(ast.Node.BuiltinCall).?;
+            if (!std.mem.eql(u8, analysis_ctx.tree.tokenSlice(builtin_call.builtin_token), "@import")) return null;
+            if (builtin_call.params.len > 1) return null;
+
+            const import_param = builtin_call.params.at(0).*;
+            if (import_param.id != .StringLiteral) return null;
+
+            const import_str = analysis_ctx.tree.tokenSlice(import_param.cast(ast.Node.StringLiteral).?.token);
+            return analysis_ctx.onImport(import_str[1 .. import_str.len - 1]) catch |err| block: {
+                std.debug.warn("Error {} while proessing import {}\n", .{err, import_str});
+                break :block null;
+            };
+        },
         else => {
             std.debug.warn("Type resolution case not implemented; {}\n", .{node.id});
         }
@@ -230,8 +257,53 @@ pub fn resolveTypeOfNode(tree: *std.zig.ast.Tree, node: *std.zig.ast.Node) ?*std
     return null;
 }
 
-pub fn getNodeFromTokens(tree: *std.zig.ast.Tree, node: *std.zig.ast.Node, tokenizer: *std.zig.Tokenizer) ?*std.zig.ast.Node {
-    var current_node = node;
+fn maybeCollectImport(tree: *ast.Tree, builtin_call: *ast.Node.BuiltinCall, arr: *std.ArrayList([]const u8)) !void {
+    if (!std.mem.eql(u8, tree.tokenSlice(builtin_call.builtin_token), "@import")) return;
+    if (builtin_call.params.len > 1) return;
+
+    const import_param = builtin_call.params.at(0).*;
+    if (import_param.id != .StringLiteral) return;
+
+    const import_str = tree.tokenSlice(import_param.cast(ast.Node.StringLiteral).?.token);
+    try arr.append(import_str[1 .. import_str.len - 1]);
+}
+
+/// Collects all imports we can find into a slice of import paths (without quotes).
+/// The import paths are valid as long as the tree is.
+pub fn collectImports(allocator: *std.mem.Allocator, tree: *ast.Tree) ![][]const u8 {
+    // TODO: Currently only detects `const smth = @import("string literal")<.SometThing>;`
+    var arr = std.ArrayList([]const u8).init(allocator);
+
+    var idx: usize = 0;
+    while (tree.root_node.iterate(idx)) |decl| : (idx += 1) {
+        if (decl.id != .VarDecl) continue;
+        const var_decl = decl.cast(ast.Node.VarDecl).?;
+        if (var_decl.init_node == null) continue;
+    
+        switch(var_decl.init_node.?.id) {
+            .BuiltinCall => {
+                const builtin_call = var_decl.init_node.?.cast(ast.Node.BuiltinCall).?;
+                try maybeCollectImport(tree, builtin_call, &arr);
+            },
+            .InfixOp => {
+                const infix_op = var_decl.init_node.?.cast(ast.Node.InfixOp).?;
+                
+                switch(infix_op.op) {
+                    .Period => {},
+                    else => continue,
+                }
+                if (infix_op.lhs.id != .BuiltinCall) continue;
+                try maybeCollectImport(tree, infix_op.lhs.cast(ast.Node.BuiltinCall).?, &arr);
+            },
+            else => {},
+        }
+    }
+
+    return arr.toOwnedSlice();
+}
+
+pub fn getFieldAccessTypeNode(analysis_ctx: *AnalysisContext, tokenizer: *std.zig.Tokenizer) ?*ast.Node {
+    var current_node = &analysis_ctx.tree.root_node.base;
 
     while (true) {
         var next = tokenizer.next();
@@ -240,14 +312,12 @@ pub fn getNodeFromTokens(tree: *std.zig.ast.Tree, node: *std.zig.ast.Node, token
                 return current_node;
             },
             .Identifier => {
-                // var root = current_node.cast(std.zig.ast.Node.Root).?;
+                // var root = current_node.cast(ast.Node.Root).?;
                 // current_node.
-                if (getChild(tree, current_node, tokenizer.buffer[next.start..next.end])) |child| {
-                    if (resolveTypeOfNode(tree, child)) |node_type| {
-                        if (resolveTypeOfNode(tree, child)) |child_type| {
-                            current_node = child_type;
-                        } else return null;
-                    }
+                if (getChild(analysis_ctx.tree, current_node, tokenizer.buffer[next.start..next.end])) |child| {
+                    if (resolveTypeOfNode(analysis_ctx, child)) |node_type| {
+                        current_node = node_type;
+                    } else return null;
                 } else return null;
             },
             .Period => {
@@ -255,8 +325,8 @@ pub fn getNodeFromTokens(tree: *std.zig.ast.Tree, node: *std.zig.ast.Node, token
                 if (after_period.id == .Eof) {
                     return current_node;
                 } else if (after_period.id == .Identifier) {
-                    if (getChild(tree, current_node, tokenizer.buffer[after_period.start..after_period.end])) |child| {
-                        if (resolveTypeOfNode(tree, child)) |child_type| {
+                    if (getChild(analysis_ctx.tree, current_node, tokenizer.buffer[after_period.start..after_period.end])) |child| {
+                        if (resolveTypeOfNode(analysis_ctx, child)) |child_type| {
                             current_node = child_type;
                         } else return null;
                     } else return null;
@@ -271,8 +341,8 @@ pub fn getNodeFromTokens(tree: *std.zig.ast.Tree, node: *std.zig.ast.Node, token
     return current_node;
 }
 
-pub fn getCompletionsFromNode(allocator: *std.mem.Allocator, tree: *std.zig.ast.Tree, node: *std.zig.ast.Node) ![]*std.zig.ast.Node {
-    var nodes = std.ArrayList(*std.zig.ast.Node).init(allocator);
+pub fn getCompletionsFromNode(allocator: *std.mem.Allocator, tree: *ast.Tree, node: *ast.Node) ![]*ast.Node {
+    var nodes = std.ArrayList(*ast.Node).init(allocator);
 
     var index: usize = 0;
     while (node.iterate(index)) |child_node| {
@@ -284,18 +354,18 @@ pub fn getCompletionsFromNode(allocator: *std.mem.Allocator, tree: *std.zig.ast.
     return nodes.items;
 }
 
-pub fn nodeToString(tree: *std.zig.ast.Tree, node: *std.zig.ast.Node) ?[]const u8 {
+pub fn nodeToString(tree: *ast.Tree, node: *ast.Node) ?[]const u8 {
     switch (node.id) {
         .ContainerField => {
-            const field = node.cast(std.zig.ast.Node.ContainerField).?;
+            const field = node.cast(ast.Node.ContainerField).?;
             return tree.tokenSlice(field.name_token);
         },
         .Identifier => {
-            const field = node.cast(std.zig.ast.Node.Identifier).?;
+            const field = node.cast(ast.Node.Identifier).?;
             return tree.tokenSlice(field.token);
         },
         .FnProto => {
-            const func = node.cast(std.zig.ast.Node.FnProto).?;
+            const func = node.cast(ast.Node.FnProto).?;
             if (func.name_token) |name_token| {
                 return tree.tokenSlice(name_token);
             }
@@ -308,7 +378,7 @@ pub fn nodeToString(tree: *std.zig.ast.Tree, node: *std.zig.ast.Node) ?[]const u
     return null;
 }
 
-pub fn nodesToString(tree: *std.zig.ast.Tree, maybe_nodes: ?[]*std.zig.ast.Node) void {
+pub fn nodesToString(tree: *ast.Tree, maybe_nodes: ?[]*ast.Node) void {
     if (maybe_nodes) |nodes| {
         for (nodes) |node| {
             std.debug.warn("- {}\n", .{nodeToString(tree, node)});

src/config.zig

@@ -2,3 +2,6 @@
 
 /// Whether to enable snippet completions
 enable_snippets: bool = true,
+
+/// zig library path
+zig_lib_path: ?[]const u8 = null,

src/document_store.zig

@@ -0,0 +1,404 @@
+const std = @import("std");
+const types = @import("types.zig");
+const URI = @import("uri.zig");
+const analysis = @import("analysis.zig");
+
+const DocumentStore = @This();
+
+pub const Handle = struct {
+    document: types.TextDocument,
+    count: usize,
+    import_uris: std.ArrayList([]const u8),
+
+    pub fn uri(handle: Handle) []const u8 {
+        return handle.document.uri;
+    }
+
+    /// Returns the zig AST resulting from parsing the document's text, even
+    /// if it contains errors.
+    pub fn dirtyTree(handle: Handle, allocator: *std.mem.Allocator) !*std.zig.ast.Tree {
+        return try std.zig.parse(allocator, handle.document.text);
+    }
+
+    /// Returns a zig AST with no errors, either from the current text or
+    /// the stored sane text, null if no such ast exists.
+    pub fn saneTree(handle: Handle, allocator: *std.mem.Allocator) !?*std.zig.ast.Tree {
+        var tree = try std.zig.parse(allocator, handle.document.text);
+        if (tree.errors.len == 0) return tree;
+
+        tree.deinit();
+        if (handle.document.sane_text) |sane| {
+            return try std.zig.parse(allocator, sane);
+        }
+        return null;
+    }
+};
+
+allocator: *std.mem.Allocator,
+handles: std.StringHashMap(Handle),
+std_uri: ?[]const u8,
+
+pub fn init(self: *DocumentStore, allocator: *std.mem.Allocator, zig_lib_path: ?[]const u8) !void {
+    self.allocator = allocator;
+    self.handles = std.StringHashMap(Handle).init(allocator);
+    errdefer self.handles.deinit();
+
+    if (zig_lib_path) |zpath| {
+        const std_path = std.fs.path.resolve(allocator, &[_][]const u8 {
+            zpath, "./std/std.zig"
+        }) catch |err| block: {
+            std.debug.warn("Failed to resolve zig std library path, error: {}\n", .{err});
+            self.std_uri = null;
+            return;
+        };
+
+        defer allocator.free(std_path);
+        // Get the std_path as a URI, so we can just append to it!
+        self.std_uri = try URI.fromPath(allocator, std_path);
+        std.debug.warn("Standard library base uri: {}\n", .{self.std_uri});
+    } else {
+        self.std_uri = null;
+    }
+}
+
+/// This function assersts the document is not open yet and takes owneship
+/// of the uri and text passed in.
+fn newDocument(self: *DocumentStore, uri: []const u8, text: []u8) !*Handle {
+    std.debug.warn("Opened document: {}\n", .{uri});
+
+    errdefer {
+        self.allocator.free(uri);
+        self.allocator.free(text);
+    }
+
+    var handle = Handle{
+        .count = 1,
+        .import_uris = std.ArrayList([]const u8).init(self.allocator),
+        .document = .{
+            .uri = uri,
+            .text = text,
+            .mem = text,
+            .sane_text = null,
+        },
+    };
+    try self.checkSanity(&handle);
+    try self.handles.putNoClobber(uri, handle);
+    return &(self.handles.get(uri) orelse unreachable).value;
+}
+
+pub fn openDocument(self: *DocumentStore, uri: []const u8, text: []const u8) !*Handle {
+    if (self.handles.get(uri)) |entry| {
+        std.debug.warn("Document already open: {}, incrementing count\n", .{uri});
+        entry.value.count += 1;
+        std.debug.warn("New count: {}\n", .{entry.value.count});
+        return &entry.value;
+    }
+
+    const duped_text = try std.mem.dupe(self.allocator, u8, text);
+    errdefer self.allocator.free(duped_text);
+    const duped_uri = try std.mem.dupe(self.allocator, u8, uri);
+    errdefer self.allocator.free(duped_uri);
+
+    return self.newDocument(duped_uri, duped_text);
+}
+
+fn decrementCount(self: *DocumentStore, uri: []const u8) void {
+    if (self.handles.get(uri)) |entry| {
+        entry.value.count -= 1;
+        if (entry.value.count > 0)
+            return;
+
+        std.debug.warn("Freeing document: {}\n", .{uri});
+        self.allocator.free(entry.value.document.mem);
+        if (entry.value.document.sane_text) |sane| {
+            self.allocator.free(sane);
+        }
+
+        for (entry.value.import_uris.items) |import_uri| {
+            self.decrementCount(import_uri);
+            self.allocator.free(import_uri);
+        }
+
+        entry.value.import_uris.deinit();
+
+        const uri_key = entry.key;
+        self.handles.removeAssertDiscard(uri);
+        self.allocator.free(uri_key);
+    }
+}
+
+pub fn closeDocument(self: *DocumentStore, uri: []const u8) void {
+    self.decrementCount(uri);
+}
+
+pub fn getHandle(self: *DocumentStore, uri: []const u8) ?*Handle {
+    if (self.handles.get(uri)) |entry| {
+        return &entry.value;
+    }
+
+    return null;
+}
+
+// Check if the document text is now sane, move it to sane_text if so.
+fn checkSanity(self: *DocumentStore, handle: *Handle) !void {
+    const dirty_tree = try handle.dirtyTree(self.allocator);
+    defer dirty_tree.deinit();
+
+    if (dirty_tree.errors.len > 0) return;
+
+    std.debug.warn("New sane text for document {}\n", .{handle.uri()});
+    if (handle.document.sane_text) |sane| {
+        self.allocator.free(sane);
+    }
+
+    handle.document.sane_text = try std.mem.dupe(self.allocator, u8, handle.document.text);
+
+    // TODO: Better algorithm or data structure?
+    // Removing the imports is costly since they live in an array list
+    // Perhaps we should use an AutoHashMap([]const u8, {}) ?
+
+    // Try to detect removed imports and decrement their counts.
+    if (handle.import_uris.items.len == 0) return;
+
+    const import_strs = try analysis.collectImports(self.allocator, dirty_tree);
+    defer self.allocator.free(import_strs);
+
+    const still_exist = try self.allocator.alloc(bool, handle.import_uris.items.len);
+    defer self.allocator.free(still_exist);
+
+    for (still_exist) |*ex| {
+        ex.* = false;
+    }
+
+    for (import_strs) |str| {
+        const uri = (try uriFromImportStr(self, handle, str)) orelse continue;
+        defer self.allocator.free(uri);
+
+        var idx: usize = 0;
+        exists_loop: while (idx < still_exist.len) : (idx += 1) {
+            if (still_exist[idx]) continue;
+
+            if (std.mem.eql(u8, handle.import_uris.items[idx], uri)) {
+                still_exist[idx] = true;
+                break :exists_loop;
+            }
+        }
+    }
+
+    // Go through still_exist, remove the items that are false and decrement their handle counts.
+    var offset: usize = 0;
+    var idx: usize = 0;
+    while (idx < still_exist.len) : (idx += 1) {
+        if (still_exist[idx]) continue;
+
+        std.debug.warn("Import removed: {}\n", .{handle.import_uris.items[idx - offset]});
+        const uri = handle.import_uris.orderedRemove(idx - offset);
+        offset += 1;
+
+        self.closeDocument(uri);
+        self.allocator.free(uri);
+    }
+}
+
+pub fn applyChanges(self: *DocumentStore, handle: *Handle, content_changes: std.json.Array) !void {
+    var document = &handle.document;
+
+    for (content_changes.items) |change| {
+        if (change.Object.getValue("range")) |range| {
+            const start_pos = types.Position{
+                .line = range.Object.getValue("start").?.Object.getValue("line").?.Integer,
+                .character = range.Object.getValue("start").?.Object.getValue("character").?.Integer
+            };
+            const end_pos = types.Position{
+                .line = range.Object.getValue("end").?.Object.getValue("line").?.Integer,
+                .character = range.Object.getValue("end").?.Object.getValue("character").?.Integer
+            };
+
+            const change_text = change.Object.getValue("text").?.String;
+            const start_index = try document.positionToIndex(start_pos);
+            const end_index = try document.positionToIndex(end_pos);
+
+            const old_len = document.text.len;
+            const new_len = old_len + change_text.len;
+            if (new_len > document.mem.len) {
+                // We need to reallocate memory.
+                // We reallocate twice the current filesize or the new length, if it's more than that
+                // so that we can reduce the amount of realloc calls.
+                // We can tune this to find a better size if needed.
+                const realloc_len = std.math.max(2 * old_len, new_len);
+                document.mem = try self.allocator.realloc(document.mem, realloc_len);
+            }
+
+            // The first part of the string, [0 .. start_index] need not be changed.
+            // We then copy the last part of the string, [end_index ..] to its
+            //    new position, [start_index + change_len .. ]
+            std.mem.copy(u8, document.mem[start_index + change_text.len..][0 .. old_len - end_index], document.mem[end_index .. old_len]);
+            // Finally, we copy the changes over.
+            std.mem.copy(u8, document.mem[start_index..][0 .. change_text.len], change_text);
+
+            // Reset the text substring.
+            document.text = document.mem[0 .. new_len];
+        } else {
+            const change_text = change.Object.getValue("text").?.String;
+            const old_len = document.text.len;
+
+            if (change_text.len > document.mem.len) {
+                // Like above.
+                const realloc_len = std.math.max(2 * old_len, change_text.len);
+                document.mem = try self.allocator.realloc(document.mem, realloc_len);
+            }
+
+            std.mem.copy(u8, document.mem[0 .. change_text.len], change_text);
+            document.text = document.mem[0 .. change_text.len];
+        }
+    }
+
+    try self.checkSanity(handle);
+}
+
+fn uriFromImportStr(store: *DocumentStore, handle: *Handle, import_str: []const u8) !?[]const u8 {
+    return if (std.mem.eql(u8, import_str, "std"))
+        if (store.std_uri) |std_root_uri| try std.mem.dupe(store.allocator, u8, std_root_uri)
+        else {
+            std.debug.warn("Cannot resolve std library import, path is null.\n", .{});
+            return null;
+        }
+    else b: {
+        // Find relative uri
+        const path = try URI.parse(store.allocator, handle.uri());
+        defer store.allocator.free(path);
+
+        const dir_path = std.fs.path.dirname(path) orelse "";
+        const import_path = try std.fs.path.resolve(store.allocator, &[_][]const u8 {
+            dir_path, import_str
+        });
+
+        defer store.allocator.free(import_path);
+
+        break :b (try URI.fromPath(store.allocator, import_path));
+    };
+}
+
+pub const AnalysisContext = struct {
+    store: *DocumentStore,
+    handle: *Handle,
+    // This arena is used for temporary allocations while analyzing,
+    // not for the tree allocations.
+    arena: *std.heap.ArenaAllocator,
+    tree: *std.zig.ast.Tree,
+
+    pub fn onImport(self: *AnalysisContext, import_str: []const u8) !?*std.zig.ast.Node {
+        const allocator = self.store.allocator;
+        const final_uri = (try uriFromImportStr(self.store, self.handle, import_str)) orelse return null;
+
+        std.debug.warn("Import final URI: {}\n", .{final_uri});
+        var consumed_final_uri = false;
+        defer if (!consumed_final_uri) allocator.free(final_uri);
+
+        // Check if we already imported this.
+        for (self.handle.import_uris.items) |uri| {
+            // If we did, set our new handle and return the parsed tree root node.
+            if (std.mem.eql(u8, uri, final_uri)) {
+                self.handle = self.store.getHandle(final_uri) orelse return null;
+
+                self.tree.deinit();
+                if (try self.handle.saneTree(allocator)) |tree| {
+                    self.tree = tree;
+                    return &self.tree.root_node.base;
+                }
+                return null;
+            }
+        }
+
+        // New import.
+        // Check if the import is already opened by others.
+        if (self.store.getHandle(final_uri)) |new_handle| {
+            // If it is, increment the count, set our new handle and return the parsed tree root node.
+            new_handle.count += 1;
+            self.handle = new_handle;
+
+            self.tree.deinit();
+            if (try self.handle.saneTree(allocator)) |tree| {
+                self.tree = tree;
+                return &self.tree.root_node.base;
+            }
+            return null;
+        }
+
+        // New document, read the file then call into openDocument.
+        const file_path = try URI.parse(allocator, final_uri);
+        defer allocator.free(file_path);
+
+        var file = std.fs.cwd().openFile(file_path, .{}) catch {
+            std.debug.warn("Cannot open import file {}\n", .{file_path});
+            return null;
+        };
+
+        defer file.close();
+        const size = std.math.cast(usize, try file.getEndPos()) catch std.math.maxInt(usize);
+
+        {
+            const file_contents = try allocator.alloc(u8, size);
+            errdefer allocator.free(file_contents);
+
+            file.inStream().readNoEof(file_contents) catch {
+                std.debug.warn("Could not read from file {}\n", .{file_path});
+                return null;
+            };
+
+            // Add to import table of current handle.
+            try self.handle.import_uris.append(final_uri);
+            consumed_final_uri = true;
+
+            // Swap handles and get new tree.
+            // This takes ownership of the passed uri and text.
+            self.handle = try newDocument(self.store, try std.mem.dupe(allocator, u8, final_uri), file_contents);
+        }
+
+        // Free old tree, add new one if it exists.
+        // If we return null, no one should access the tree.
+        self.tree.deinit();
+        if (try self.handle.saneTree(allocator)) |tree| {
+            self.tree = tree;
+            return &self.tree.root_node.base;
+        }
+        return null;
+    }
+
+    pub fn deinit(self: *AnalysisContext) void {
+        self.tree.deinit();
+    }
+};
+
+pub fn analysisContext(self: *DocumentStore, handle: *Handle, arena: *std.heap.ArenaAllocator) !?AnalysisContext {
+    const tree = (try handle.saneTree(self.allocator)) orelse return null;
+
+    return AnalysisContext{
+        .store = self,
+        .handle = handle,
+        .arena = arena,
+        .tree = tree,
+    };
+}
+
+pub fn deinit(self: *DocumentStore) void {
+    var entry_iterator = self.handles.iterator();
+    while (entry_iterator.next()) |entry| {
+        self.allocator.free(entry.value.document.mem);
+        if (entry.value.document.sane_text) |sane| {
+            self.allocator.free(sane);
+        }
+
+        for (entry.value.import_uris.items) |uri| {
+            self.allocator.free(uri);
+        }
+
+        entry.value.import_uris.deinit();
+        self.allocator.free(entry.key);
+    }
+
+    self.handles.deinit();
+    if (self.std_uri) |uri| {
+        self.allocator.free(uri);
+    }
+}

src/main.zig

@@ -2,7 +2,7 @@ const std = @import("std");
 const build_options = @import("build_options");
 
 const Config = @import("config.zig");
-const Uri = @import("uri.zig");
+const DocumentStore = @import("document_store.zig");
 const data = @import("data/" ++ build_options.data_version ++ ".zig");
 const types = @import("types.zig");
 const analysis = @import("analysis.zig");
@@ -12,8 +12,7 @@ const analysis = @import("analysis.zig");
 var stdout: std.fs.File.OutStream = undefined;
 var allocator: *std.mem.Allocator = undefined;
 
-/// Documents hashmap, types.DocumentUri:types.TextDocument
-var documents: std.StringHashMap(types.TextDocument) = undefined;
+var document_store: DocumentStore = undefined;
 
 const initialize_response = \\,"result":{"capabilities":{"signatureHelpProvider":{"triggerCharacters":["(",","]},"textDocumentSync":1,"completionProvider":{"resolveProvider":false,"triggerCharacters":[".",":","@"]},"documentHighlightProvider":false,"codeActionProvider":false,"workspace":{"workspaceFolders":{"supported":true}}}}}
 ;
@@ -78,48 +77,6 @@ fn respondGeneric(id: i64, response: []const u8) !void {
     try stdout.writeAll(response);
 }
 
-fn freeDocument(document: types.TextDocument) void {
-    allocator.free(document.uri);
-    allocator.free(document.mem);
-    if (document.sane_text) |str| {
-        allocator.free(str);
-    }
-}
-
-fn openDocument(uri: []const u8, text: []const u8) !void {
-    const duped_uri = try std.mem.dupe(allocator, u8, uri);
-    const duped_text = try std.mem.dupe(allocator, u8, text);
-
-    const res = try documents.put(duped_uri, .{
-        .uri = duped_uri,
-        .text = duped_text,
-        .mem = duped_text,
-    });
-
-    if (res) |entry| {
-        try log("Document already open: {}, closing old.", .{uri});
-        freeDocument(entry.value);
-    } else {
-        try log("Opened document: {}", .{uri});
-    }
-}
-
-fn closeDocument(uri: []const u8) !void {
-    if (documents.remove(uri)) |entry| {
-        try log("Closing document: {}", .{uri});
-        freeDocument(entry.value);
-    }
-}
-
-fn cacheSane(document: *types.TextDocument) !void {
-    try log("Caching sane text for document: {}", .{document.uri});
-
-    if (document.sane_text) |old_sane| {
-        allocator.free(old_sane);
-    }
-    document.sane_text = try std.mem.dupe(allocator, u8, document.text);
-}
-
 // TODO: Is this correct or can we get a better end? 
 fn astLocationToRange(loc: std.zig.ast.Tree.Location) types.Range {
     return .{
@@ -134,8 +91,8 @@ fn astLocationToRange(loc: std.zig.ast.Tree.Location) types.Range {
     };
 }
 
-fn publishDiagnostics(document: *types.TextDocument, config: Config) !void {
-    const tree = try std.zig.parse(allocator, document.text);
+fn publishDiagnostics(handle: DocumentStore.Handle, config: Config) !void {
+    const tree = try handle.dirtyTree(allocator);
     defer tree.deinit();
 
     // Use an arena for our local memory allocations.
@@ -163,7 +120,6 @@ fn publishDiagnostics(document: *types.TextDocument, config: Config) !void {
     }
 
     if (tree.errors.len == 0) {
-        try cacheSane(document);
         var decls = tree.root_node.decls.iterator(0);
         while (decls.next()) |decl_ptr| {
             var decl = decl_ptr.*;
@@ -214,7 +170,7 @@ fn publishDiagnostics(document: *types.TextDocument, config: Config) !void {
         .method = "textDocument/publishDiagnostics",
         .params = .{
             .PublishDiagnosticsParams = .{
-                .uri = document.uri,
+                .uri = handle.uri(),
                 .diagnostics = diagnostics.items,
             },
         },
@@ -267,18 +223,8 @@ fn nodeToCompletion(alloc: *std.mem.Allocator, tree: *std.zig.ast.Tree, decl: *s
     return null;
 }
 
-fn completeGlobal(id: i64, document: *types.TextDocument, config: Config) !void {
-    // The tree uses its own arena, so we just pass our main allocator.
-    var tree = try std.zig.parse(allocator, document.text);
-
-    if (tree.errors.len > 0) {
-        if (document.sane_text) |sane_text| {
-            tree.deinit();
-            tree = try std.zig.parse(allocator, sane_text);
-        } else return try respondGeneric(id, no_completions_response);
-    }
-    else try cacheSane(document);
-
+fn completeGlobal(id: i64, handle: DocumentStore.Handle, config: Config) !void {
+    var tree = (try handle.saneTree(allocator)) orelse return respondGeneric(id, no_completions_response);
     defer tree.deinit();
 
     // We use a local arena allocator to deallocate all temporary data without iterating
@@ -306,41 +252,12 @@ fn completeGlobal(id: i64, document: *types.TextDocument, config: Config) !void
     });
 }
 
-fn completeFieldAccess(id: i64, document: *types.TextDocument, position: types.Position, config: Config) !void {
-    if (document.sane_text) |sane_text| {
-        var tree = try std.zig.parse(allocator, sane_text);
-        defer tree.deinit();
+fn completeFieldAccess(id: i64, handle: *DocumentStore.Handle, position: types.Position, config: Config) !void {
+    var arena = std.heap.ArenaAllocator.init(allocator);
+    defer arena.deinit();
 
-        // We use a local arena allocator to deallocate all temporary data without iterating
-        var arena = std.heap.ArenaAllocator.init(allocator);
-        var completions = std.ArrayList(types.CompletionItem).init(&arena.allocator);
-        // Deallocate all temporary data.
-        defer arena.deinit();
-
-        var line = try document.getLine(@intCast(usize, position.line));
-        var tokenizer = std.zig.Tokenizer.init(line);
-
-        if (analysis.getNodeFromTokens(tree, &tree.root_node.base, &tokenizer)) |node| {
-            var index: usize = 0;
-            while (node.iterate(index)) |child_node| {
-                if (try nodeToCompletion(&arena.allocator, tree, child_node, config)) |completion| {
-                    try completions.append(completion);
-                }
-                index += 1;
-            }
-        }
-
-        try send(types.Response{
-            .id = .{.Integer = id},
-            .result = .{
-                .CompletionList = .{
-                    .isIncomplete = false,
-                    .items = completions.items,
-                },
-            },
-        });
-    } else {
-        return try send(types.Response{
+    var analysis_ctx = (try document_store.analysisContext(handle, &arena)) orelse {
+        return send(types.Response{
             .id = .{.Integer = id},
             .result = .{
                 .CompletionList = .{
@@ -349,7 +266,33 @@ fn completeFieldAccess(id: i64, document: *types.TextDocument, position: types.P
                 },
             },
         });
+    };
+    defer analysis_ctx.deinit();
+
+    var completions = std.ArrayList(types.CompletionItem).init(&arena.allocator);
+
+    var line = try handle.document.getLine(@intCast(usize, position.line));
+    var tokenizer = std.zig.Tokenizer.init(line);
+
+    if (analysis.getFieldAccessTypeNode(&analysis_ctx, &tokenizer)) |node| {
+        var index: usize = 0;
+        while (node.iterate(index)) |child_node| {
+            if (try nodeToCompletion(&arena.allocator, analysis_ctx.tree, child_node, config)) |completion| {
+                try completions.append(completion);
+            }
+            index += 1;
+        }
     }
+
+    try send(types.Response{
+        .id = .{.Integer = id},
+        .result = .{
+            .CompletionList = .{
+                .isIncomplete = false,
+                .items = completions.items,
+            },
+        },
+    });
 }
 
 // Compute builtin completions at comptime.
@@ -523,73 +466,27 @@ fn processJsonRpc(parser: *std.json.Parser, json: []const u8, config: Config) !v
         const uri = document.getValue("uri").?.String;
         const text = document.getValue("text").?.String;
 
-        try openDocument(uri, text);
-        try publishDiagnostics(&(documents.get(uri).?.value), config);
+        const handle = try document_store.openDocument(uri, text);
+        try publishDiagnostics(handle.*, config);
     } else if (std.mem.eql(u8, method, "textDocument/didChange")) {
         const text_document = params.getValue("textDocument").?.Object;
         const uri = text_document.getValue("uri").?.String;
-
-        var document = &(documents.get(uri).?.value);
         const content_changes = params.getValue("contentChanges").?.Array;
 
-        for (content_changes.items) |change| {
-            if (change.Object.getValue("range")) |range| {
-                const start_pos = types.Position{
-                    .line = range.Object.getValue("start").?.Object.getValue("line").?.Integer,
-                    .character = range.Object.getValue("start").?.Object.getValue("character").?.Integer
-                };
-                const end_pos = types.Position{
-                    .line = range.Object.getValue("end").?.Object.getValue("line").?.Integer,
-                    .character = range.Object.getValue("end").?.Object.getValue("character").?.Integer
-                };
+        const handle = document_store.getHandle(uri) orelse {
+            try log("Trying to change non existent document {}", .{uri});
+            return;
+        };
 
-                const change_text = change.Object.getValue("text").?.String;
-                const start_index = try document.positionToIndex(start_pos);
-                const end_index = try document.positionToIndex(end_pos);
-
-                const old_len = document.text.len;
-                const new_len = old_len + change_text.len;
-                if (new_len > document.mem.len) {
-                    // We need to reallocate memory.
-                    // We reallocate twice the current filesize or the new length, if it's more than that
-                    // so that we can reduce the amount of realloc calls.
-                    // We can tune this to find a better size if needed.
-                    const realloc_len = std.math.max(2 * old_len, new_len);
-                    document.mem = try allocator.realloc(document.mem, realloc_len);
-                }
-
-                // The first part of the string, [0 .. start_index] need not be changed.
-                // We then copy the last part of the string, [end_index ..] to its
-                //    new position, [start_index + change_len .. ]
-                std.mem.copy(u8, document.mem[start_index + change_text.len..][0 .. old_len - end_index], document.mem[end_index .. old_len]);
-                // Finally, we copy the changes over.
-                std.mem.copy(u8, document.mem[start_index..][0 .. change_text.len], change_text);
-
-                // Reset the text substring.
-                document.text = document.mem[0 .. new_len];
-            } else {
-                const change_text = change.Object.getValue("text").?.String;
-                const old_len = document.text.len;
-
-                if (change_text.len > document.mem.len) {
-                    // Like above.
-                    const realloc_len = std.math.max(2 * old_len, change_text.len);
-                    document.mem = try allocator.realloc(document.mem, realloc_len);
-                }
-
-                std.mem.copy(u8, document.mem[0 .. change_text.len], change_text);
-                document.text = document.mem[0 .. change_text.len];
-            }
-        }
-
-        try publishDiagnostics(document, config);
+        try document_store.applyChanges(handle, content_changes);
+        try publishDiagnostics(handle.*, config);
     } else if (std.mem.eql(u8, method, "textDocument/didSave")) {
         // noop
     } else if (std.mem.eql(u8, method, "textDocument/didClose")) {
         const document = params.getValue("textDocument").?.Object;
         const uri = document.getValue("uri").?.String;
 
-        try closeDocument(uri);
+        document_store.closeDocument(uri);
     }
     // Autocomplete / Signatures
     else if (std.mem.eql(u8, method, "textDocument/completion")) {
@@ -597,14 +494,18 @@ fn processJsonRpc(parser: *std.json.Parser, json: []const u8, config: Config) !v
         const uri = text_document.getValue("uri").?.String;
         const position = params.getValue("position").?.Object;
 
-        var document = &(documents.get(uri).?.value);
+        const handle = document_store.getHandle(uri) orelse {
+            try log("Trying to complete in non existent document {}", .{uri});
+            return;
+        };
+
         const pos = types.Position{
             .line = position.getValue("line").?.Integer,
             .character = position.getValue("character").?.Integer - 1,
         };
         if (pos.character >= 0) {
-            const pos_index = try document.positionToIndex(pos);
-            const pos_context = documentPositionContext(document.*, pos_index);
+            const pos_index = try handle.document.positionToIndex(pos);
+            const pos_context = documentPositionContext(handle.document, pos_index);
 
             if (pos_context == .builtin) {
                 try send(types.Response{
@@ -617,9 +518,9 @@ fn processJsonRpc(parser: *std.json.Parser, json: []const u8, config: Config) !v
                     },
                 });
             } else if (pos_context == .var_access or pos_context == .empty) {
-                try completeGlobal(id, document, config);
+                try completeGlobal(id, handle.*, config);
             } else if (pos_context == .field_access) {
-                try completeFieldAccess(id, document, pos, config);
+                try completeFieldAccess(id, handle, pos, config);
             } else {
                 try respondGeneric(id, no_completions_response);
             }
@@ -676,11 +577,9 @@ pub fn main() anyerror!void {
     const stdin = std.io.getStdIn().inStream();
     stdout = std.io.getStdOut().outStream();
 
-
-    documents = std.StringHashMap(types.TextDocument).init(allocator);
-
     // Read he configuration, if any.
     var config = Config{};
+    const config_parse_options = std.json.ParseOptions{ .allocator=allocator };
 
     // TODO: Investigate using std.fs.Watch to detect writes to the config and reload it.
     config_read: {
@@ -703,13 +602,16 @@ pub fn main() anyerror!void {
         if (bytes_read != conf_file_stat.size) break :config_read;
 
         // TODO: Better errors? Doesnt seem like std.json can provide us positions or context.
-        // Note that we don't need to pass an allocator to parse since we are not using pointer or slice fields.
-        // Thus, we don't need to even call parseFree.
-        config = std.json.parse(Config, &std.json.TokenStream.init(file_buf), std.json.ParseOptions{}) catch |err| {
+        config = std.json.parse(Config, &std.json.TokenStream.init(file_buf), config_parse_options) catch |err| {
             std.debug.warn("Error while parsing configuration file: {}\nUsing default config.\n", .{err});
             break :config_read;
         };
     }
+    defer std.json.parseFree(Config, config, config_parse_options);
+
+    // @TODO Check is_absolute
+    try document_store.init(allocator, config.zig_lib_path);
+    defer document_store.deinit();
 
     // This JSON parser is passed to processJsonRpc and reset.
     var json_parser = std.json.Parser.init(allocator, false);

src/uri.zig

@@ -1,7 +1,37 @@
 const std = @import("std");
 
-// Original code: https://github.com/andersfr/zig-lsp/blob/master/uri.zig
+const reserved_chars = &[_]u8 {
+    '!', '#', '$', '%', '&', '\'',
+    '(', ')', '*', '+', ',', ':',
+    ';', '=', '?', '@', '[', ']',
+};
 
+/// Returns a URI from a path, caller owns the memory allocated with `allocator`
+pub fn fromPath(allocator: *std.mem.Allocator, path: []const u8) ![]const u8 {
+    if (path.len == 0) return "";
+    const prefix = if (std.builtin.os.tag == .windows) "file:///" else "file://";
+
+    var buf = std.ArrayList(u8).init(allocator);
+    try buf.appendSlice(prefix);
+
+    var out_stream = buf.outStream();
+
+    for (path) |char| {
+        if (char == std.fs.path.sep) {
+            try buf.append('/');
+        } else if (std.mem.indexOfScalar(u8, reserved_chars, char) != null) {
+            // Write '%' + hex with uppercase
+            try buf.append('%');
+            try std.fmt.format(out_stream, "{X}", .{char});
+        } else {
+            try buf.append(std.ascii.toLower(char));
+        }
+    }
+
+    return buf.toOwnedSlice();
+}
+
+// Original code: https://github.com/andersfr/zig-lsp/blob/master/uri.zig
 fn parseHex(c: u8) !u8 {
     return switch(c) {
         '0'...'9' => c - '0',