Commit 5a313192e6

@@ -130,6 +130,77 @@ pub fn get_child(decl: *const Decl, name: []const u8) ?Decl.Index {
             const child_node = scope.get_child(name) orelse return null;
             return file.node_decls.get(child_node);
         },
+        .type_function => {
+            // Find a decl with this function as the parent, with a name matching `name`
+            for (Walk.decls.items, 0..) |*candidate, i| {
+                if (candidate.parent != .none and candidate.parent.get() == decl and std.mem.eql(u8, candidate.extra_info().name, name)) {
+                    return @enumFromInt(i);
+                }
+            }
+
+            return null;
+        },
+        else => return null,
+    }
+}
+
+/// If the type function returns another type function, return the index of that type function.
+pub fn get_type_fn_return_type_fn(decl: *const Decl) ?Decl.Index {
+    if (decl.get_type_fn_return_expr()) |return_expr| {
+        const ast = decl.file.get_ast();
+        const node_tags = ast.nodes.items(.tag);
+
+        switch (node_tags[return_expr]) {
+            .call, .call_comma, .call_one, .call_one_comma => {
+                const node_data = ast.nodes.items(.data);
+                const function = node_data[return_expr].lhs;
+                const token = ast.nodes.items(.main_token)[function];
+                const name = ast.tokenSlice(token);
+                if (decl.lookup(name)) |function_decl| {
+                    return function_decl;
+                }
+            },
+            else => {},
+        }
+    }
+    return null;
+}
+
+/// Gets the expression after the `return` keyword in a type function declaration.
+pub fn get_type_fn_return_expr(decl: *const Decl) ?Ast.Node.Index {
+    switch (decl.categorize()) {
+        .type_function => {
+            const ast = decl.file.get_ast();
+            const node_tags = ast.nodes.items(.tag);
+            const node_data = ast.nodes.items(.data);
+            const body_node = node_data[decl.ast_node].rhs;
+            if (body_node == 0) return null;
+
+            switch (node_tags[body_node]) {
+                .block, .block_semicolon => {
+                    const statements = ast.extra_data[node_data[body_node].lhs..node_data[body_node].rhs];
+                    // Look for the return statement
+                    for (statements) |stmt| {
+                        if (node_tags[stmt] == .@"return") {
+                            return node_data[stmt].lhs;
+                        }
+                    }
+                    return null;
+                },
+                .block_two, .block_two_semicolon => {
+                    if (node_tags[node_data[body_node].lhs] == .@"return") {
+                        return node_data[node_data[body_node].lhs].lhs;
+                    }
+                    if (node_data[body_node].rhs != 0 and
+                        node_tags[node_data[body_node].rhs] == .@"return")
+                    {
+                        return node_data[node_data[body_node].rhs].lhs;
+                    }
+                    return null;
+                },
+                else => return null,
+            }
+        },
         else => return null,
     }
 }

@@ -380,16 +380,48 @@ export fn decl_params(decl_index: Decl.Index) Slice(Ast.Node.Index) {
 }
 
 fn decl_fields_fallible(decl_index: Decl.Index) ![]Ast.Node.Index {
+    const decl = decl_index.get();
+    const ast = decl.file.get_ast();
+
+    switch (decl.categorize()) {
+        .type_function => {
+            const node_tags = ast.nodes.items(.tag);
+
+            // Find the return statement
+            if (decl.get_type_fn_return_expr()) |return_expr| {
+                switch (node_tags[return_expr]) {
+                    .call, .call_comma, .call_one, .call_one_comma => {
+                        const node_data = ast.nodes.items(.data);
+                        const function = node_data[return_expr].lhs;
+                        const token = ast.nodes.items(.main_token)[function];
+                        const name = ast.tokenSlice(token);
+                        if (decl.lookup(name)) |function_decl| {
+                            return decl_fields_fallible(function_decl);
+                        }
+                    },
+                    .container_decl, .container_decl_trailing, .container_decl_two, .container_decl_two_trailing, .container_decl_arg, .container_decl_arg_trailing => {
+                        return ast_decl_fields_fallible(ast, return_expr);
+                    },
+                    else => {},
+                }
+            }
+            return &.{};
+        },
+        else => {
+            const value_node = decl.value_node() orelse return &.{};
+            return ast_decl_fields_fallible(ast, value_node);
+        },
+    }
+}
+
+fn ast_decl_fields_fallible(ast: *Ast, ast_index: Ast.Node.Index) ![]Ast.Node.Index {
     const g = struct {
         var result: std.ArrayListUnmanaged(Ast.Node.Index) = .empty;
     };
     g.result.clearRetainingCapacity();
-    const decl = decl_index.get();
-    const ast = decl.file.get_ast();
     const node_tags = ast.nodes.items(.tag);
-    const value_node = decl.value_node() orelse return &.{};
     var buf: [2]Ast.Node.Index = undefined;
-    const container_decl = ast.fullContainerDecl(&buf, value_node) orelse return &.{};
+    const container_decl = ast.fullContainerDecl(&buf, ast_index) orelse return &.{};
     for (container_decl.ast.members) |member_node| switch (node_tags[member_node]) {
         .container_field_init,
         .container_field_align,
@@ -883,6 +915,13 @@ export fn categorize_decl(decl_index: Decl.Index, resolve_alias_count: usize) Wa
 }
 
 export fn type_fn_members(parent: Decl.Index, include_private: bool) Slice(Decl.Index) {
+    const decl = parent.get();
+
+    // If the type function returns another type function, get the members of that function
+    if (decl.get_type_fn_return_type_fn()) |function_decl| {
+        return namespace_members(function_decl, include_private);
+    }
+
     return namespace_members(parent, include_private);
 }