Resolve peer enum/union types

Author: FnControlOption

src/analysis.zig

@@ -982,6 +982,13 @@ pub fn resolveUnwrapErrorUnionType(analyser: *Analyser, ty: Type, side: ErrorUni
 }
 
 fn resolveUnionTag(analyser: *Analyser, ty: Type) error{OutOfMemory}!?Type {
+    const tag_type = try analyser.resolveUnionTagType(ty) orelse
+        return null;
+
+    return try tag_type.instanceTypeVal(analyser);
+}
+
+fn resolveUnionTagType(analyser: *Analyser, ty: Type) error{OutOfMemory}!?Type {
     if (!ty.is_type_val)
         return null;
 
@@ -1000,12 +1007,10 @@ fn resolveUnionTag(analyser: *Analyser, ty: Type) error{OutOfMemory}!?Type {
         return null;
 
     if (container_decl.ast.enum_token != null)
-        return .{ .data = .{ .union_tag = try analyser.allocType(ty) }, .is_type_val = false };
+        return .{ .data = .{ .union_tag = try analyser.allocType(ty) }, .is_type_val = true };
 
-    if (container_decl.ast.arg.unwrap()) |arg| {
-        const tag_type = (try analyser.resolveTypeOfNode(.of(arg, handle))) orelse return null;
-        return try tag_type.instanceTypeVal(analyser) orelse return null;
-    }
+    if (container_decl.ast.arg.unwrap()) |arg|
+        return try analyser.resolveTypeOfNode(.of(arg, handle));
 
     return null;
 }
@@ -6979,7 +6984,59 @@ fn resolvePeerTypesInner(analyser: *Analyser, peer_tys: []?Type) !?Type {
             return opt_cur_ty.?;
         },
 
-        .enum_or_union => return null, // TODO
+        .enum_or_union => {
+            var opt_cur_ty: ?Type = null;
+
+            for (peer_tys) |opt_ty| {
+                const ty = opt_ty orelse continue;
+                switch (ty.zigTypeTag(analyser).?) {
+                    .enum_literal, .@"enum", .@"union" => {},
+                    else => return null,
+                }
+                const cur_ty = opt_cur_ty orelse {
+                    opt_cur_ty = ty;
+                    continue;
+                };
+
+                if ((cur_ty.isUnionType() and !cur_ty.isTaggedUnion()) or
+                    (ty.isUnionType() and !ty.isTaggedUnion()))
+                {
+                    if (cur_ty.eql(ty)) continue;
+                    return null;
+                }
+
+                switch (cur_ty.zigTypeTag(analyser).?) {
+                    .enum_literal => {
+                        opt_cur_ty = ty;
+                    },
+                    .@"enum" => switch (ty.zigTypeTag(analyser).?) {
+                        .enum_literal => {},
+                        .@"enum" => {
+                            if (!ty.eql(cur_ty)) return null;
+                        },
+                        .@"union" => {
+                            const tag_ty = try analyser.resolveUnionTagType(ty);
+                            if (!tag_ty.?.eql(cur_ty)) return null;
+                            opt_cur_ty = ty;
+                        },
+                        else => unreachable,
+                    },
+                    .@"union" => switch (ty.zigTypeTag(analyser).?) {
+                        .enum_literal => {},
+                        .@"enum" => {
+                            const cur_tag_ty = try analyser.resolveUnionTagType(cur_ty);
+                            if (!ty.eql(cur_tag_ty.?)) return null;
+                        },
+                        .@"union" => {
+                            if (!ty.eql(cur_ty)) return null;
+                        },
+                        else => unreachable,
+                    },
+                    else => unreachable,
+                }
+            }
+            return opt_cur_ty.?;
+        },
 
         .int_or_float => {
             var ip_indices: std.ArrayListUnmanaged(InternPool.Index) = try .initCapacity(analyser.gpa, peer_tys.len);

tests/analysis/peer_type_resolution.zig

@@ -5,11 +5,36 @@ const S = struct {
     int: i64,
     float: f32,
 };
+
 const s: S = .{
     .int = 0,
     .float = 1.2,
 };
 
+const E = enum {
+    foo,
+    bar,
+    baz,
+};
+
+const e: E = .bar;
+
+const T = union(E) {
+    foo: void,
+    bar: void,
+    baz: u16,
+};
+
+const t: T = .{ .baz = 3 };
+
+const U = union {
+    foo: void,
+    bar: void,
+    baz: void,
+};
+
+const u: U = .{ .baz = {} };
+
 pub fn main() !void {
     var runtime_bool: bool = true;
 
@@ -277,6 +302,57 @@ fn void_fn() void {}
 const optional_fn = if (comptime_bool) @as(?fn () void, void_fn) else void_fn;
 //    ^^^^^^^^^^^ (?fn () void)()
 
+const optional_enum_literal = if (comptime_bool) @as(?@Type(.enum_literal), .foo) else .bar;
+//    ^^^^^^^^^^^^^^^^^^^^^ (?@Type(.enum_literal))()
+
+const enum_literal_and_enum_literal = if (comptime_bool) .foo else .bar;
+//    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ (@Type(.enum_literal))()
+
+const enum_literal_and_enum = if (comptime_bool) .foo else e;
+//    ^^^^^^^^^^^^^^^^^^^^^ (E)()
+
+const enum_literal_and_union = if (comptime_bool) .foo else u;
+//    ^^^^^^^^^^^^^^^^^^^^^^ (either type)()
+
+const enum_literal_and_tagged_union = if (comptime_bool) .foo else t;
+//    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ (T)()
+
+const enum_and_enum_literal = if (comptime_bool) e else .foo;
+//    ^^^^^^^^^^^^^^^^^^^^^ (E)()
+
+const enum_and_enum = if (comptime_bool) e else e;
+//    ^^^^^^^^^^^^^ (E)()
+
+const enum_and_union = if (comptime_bool) e else u;
+//    ^^^^^^^^^^^^^^ (either type)()
+
+const enum_and_tagged_union = if (comptime_bool) e else t;
+//    ^^^^^^^^^^^^^^^^^^^^^ (T)()
+
+const union_and_enum_literal = if (comptime_bool) u else .foo;
+//    ^^^^^^^^^^^^^^^^^^^^^^ (either type)()
+
+const union_and_enum = if (comptime_bool) u else e;
+//    ^^^^^^^^^^^^^^ (either type)()
+
+const union_and_union = if (comptime_bool) u else u;
+//    ^^^^^^^^^^^^^^^ (U)()
+
+const union_and_tagged_union = if (comptime_bool) u else t;
+//    ^^^^^^^^^^^^^^^^^^^^^^ (either type)()
+
+const tagged_union_and_enum_literal = if (comptime_bool) t else .foo;
+//    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ (T)()
+
+const tagged_union_and_enum = if (comptime_bool) t else e;
+//    ^^^^^^^^^^^^^^^^^^^^^ (T)()
+
+const tagged_union_and_union = if (comptime_bool) t else u;
+//    ^^^^^^^^^^^^^^^^^^^^^^ (either type)()
+
+const tagged_union_and_tagged_union = if (comptime_bool) t else t;
+//    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ (T)()
+
 const f32_and_u32 = if (comptime_bool) @as(f32, 0) else @as(i32, 0);
 //    ^^^^^^^^^^^ (either type)()