rework nodesAtLoc

src/ast.zig

@@ -1584,41 +1584,92 @@ pub fn nodeChildrenRecursiveAlloc(allocator: std.mem.Allocator, tree: Ast, node:
 pub fn nodesAtLoc(allocator: std.mem.Allocator, tree: Ast, loc: offsets.Loc) error{OutOfMemory}![]Ast.Node.Index {
     std.debug.assert(loc.start <= loc.end and loc.end <= tree.source.len);
 
-    var nodes = std.ArrayListUnmanaged(Ast.Node.Index){};
+    const Context = struct {
-    errdefer nodes.deinit(allocator);
+        allocator: std.mem.Allocator,
+        nodes: std.ArrayListUnmanaged(Ast.Node.Index) = .{},
+        locs: std.ArrayListUnmanaged(offsets.Loc) = .{},
+
+        pub fn append(self: *@This(), ast: Ast, node: Ast.Node.Index) !void {
+            if (node == 0) return;
+            try self.nodes.append(self.allocator, node);
+            try self.locs.append(self.allocator, offsets.nodeToLoc(ast, node));
+        }
+    };
+    var context: Context = .{ .allocator = allocator };
+    defer context.nodes.deinit(allocator);
+    defer context.locs.deinit(allocator);
+
+    try context.nodes.ensureTotalCapacity(allocator, 32);
+
     var parent: Ast.Node.Index = 0; // root node
-
-    try nodes.ensureTotalCapacity(allocator, 32);
-
     while (true) {
-        const children = try nodeChildrenAlloc(allocator, tree, parent);
+        try iterateChildren(tree, parent, &context, error{OutOfMemory}, Context.append);
-        defer allocator.free(children);
 
-        var children_loc: ?offsets.Loc = null;
+        if (smallestEnclosingSubrange(context.locs.items, loc)) |subslice| {
-        for (children) |child_node| {
+            std.debug.assert(subslice.len != 0);
-            const child_loc = offsets.nodeToLoc(tree, child_node);
+            const nodes = context.nodes.items[subslice.start .. subslice.start + subslice.len];
+            if (nodes.len == 1) { // recurse over single child node
+                parent = nodes[0];
+                context.nodes.clearRetainingCapacity();
+                context.locs.clearRetainingCapacity();
+                continue;
+            } else { // end-condition: found enclosing children
+                return try allocator.dupe(Ast.Node.Index, nodes);
+            }
+        } else { // the children cannot enclose the given source location
+            context.nodes.clearRetainingCapacity();
+            context.nodes.appendAssumeCapacity(parent); // capacity is never 0
+            return try context.nodes.toOwnedSlice(allocator);
+        }
+    }
+}
 
-            const merge_child = offsets.locIntersect(loc, child_loc) or offsets.locInside(child_loc, loc);
+/// the following code can be described as the the following problem:
+/// @param children a non-intersecting list of source ranges
+/// @param loc be a source range
+///
+/// @return a slice of #children
+///
+/// Return the smallest possible subrange of #children whose
+/// combined source range is inside #loc.
+/// If #children cannot contain #loc i.e #loc is too large, return null.
+/// @see tests/utility.ast.zig for usage examples
+pub fn smallestEnclosingSubrange(children: []const offsets.Loc, loc: offsets.Loc) ?struct {
+    start: usize,
+    len: usize,
+} {
+    switch (children.len) {
+        0 => return null,
+        1 => return if (offsets.locInside(loc, children[0])) .{ .start = 0, .len = 1 } else null,
+        else => {
+            for (children[0 .. children.len - 1], children[1..]) |previous_loc, current_loc| {
+                std.debug.assert(previous_loc.end <= current_loc.start); // must by sorted
+                std.debug.assert(!offsets.locIntersect(previous_loc, current_loc)); // must be non-intersecting
+            }
+        },
+    }
 
-            if (merge_child) {
+    var i: usize = 0;
-                children_loc = if (children_loc) |l| offsets.locMerge(l, child_loc) else child_loc;
+    const start: usize = while (i < children.len) : (i += 1) {
-                try nodes.append(allocator, child_node);
+        const child_loc = children[i];
+        if (child_loc.end < loc.start) continue;
+
+        if (child_loc.start <= loc.start) {
+            break i;
+        } else if (i != 0) {
+            break i - 1;
         } else {
-                if (nodes.items.len != 0) break;
+            return null;
-            }
         }
+    } else return null;
 
-        if (children_loc == null or !offsets.locInside(loc, children_loc.?)) {
+    const end = while (i < children.len) : (i += 1) {
-            nodes.clearRetainingCapacity();
+        const child_loc = children[i];
-            nodes.appendAssumeCapacity(parent); // capacity is never 0
+        if (loc.end <= child_loc.end) break i + 1;
-            return try nodes.toOwnedSlice(allocator);
+    } else return null;
-        }
 
-        if (nodes.items.len == 1) {
+    return .{
-            parent = nodes.items[0];
+        .start = start,
-            nodes.clearRetainingCapacity();
+        .len = end - start,
-        } else {
+    };
-            return try nodes.toOwnedSlice(allocator);
-        }
-    }
 }