Add union and enum support

examples/basic/index.d.ts

@@ -45,6 +45,23 @@ export declare class TestFactoryClass {
   format(): string
 }
 
+export interface MessagePayload {
+  title: string
+  count: number
+}
+
+export declare const enum Color {
+  Red = 1,
+  Green = 2,
+  Blue = 4,
+}
+
+export declare const enum StringColor {
+  Red = 'Red',
+  Green = 'Green',
+  Blue = 'Blue',
+}
+
 
 export declare function test_i32(left: number, right: number): number
 export declare function test_f32(left: number, right: number): number
@@ -81,3 +98,28 @@ export declare function create_dataview(): DataView
 export declare function get_dataview_length(view: DataView): number
 export declare function get_dataview_first_byte(view: DataView): number
 export declare function get_dataview_uint32_le(view: DataView): number
+export declare function union_identity(value: number | string): number | string
+export declare function make_union(is_number: boolean): number | string
+export declare function union_kind(value: number | string): string
+export declare function object_or_text_identity(value: MessagePayload | string): MessagePayload | string
+export declare function make_object_or_text(as_object: boolean): MessagePayload | string
+export declare function object_or_array_identity(value: MessagePayload | Array<number>): MessagePayload | Array<number>
+export declare function tuple_or_text_identity(value: [number, boolean, string] | string): [number, boolean, string] | string
+export declare function flip_flag_or_increment(value: boolean | number): boolean | number
+export declare function color_or_text_identity(value: Color | string): Color | string
+export declare function favorite_color_or_text(use_color: boolean): Color | string
+export declare function maybe_text_or_count_identity(value: string | undefined | null | number): string | undefined | null | number
+export declare function make_maybe_text_or_count(as_text: boolean): string | undefined | null | number
+export declare function buffer_or_text_identity(value: Buffer | string): Buffer | string
+export declare function make_buffer_or_text(use_buffer: boolean): Buffer | string
+export declare function arraybuffer_or_array_identity(value: ArrayBuffer | Array<number>): ArrayBuffer | Array<number>
+export declare function make_arraybuffer_or_array(use_arraybuffer: boolean): ArrayBuffer | Array<number>
+export declare function payload_or_color_identity(value: MessagePayload | Color): MessagePayload | Color
+export declare function make_payload_or_color(use_payload: boolean): MessagePayload | Color
+export declare function payload_or_string_color_identity(value: MessagePayload | StringColor): MessagePayload | StringColor
+export declare function make_payload_or_string_color(use_payload: boolean): MessagePayload | StringColor
+export declare function enum_identity(color: Color): Color
+export declare function favorite_color(): Color
+export declare function is_primary(color: Color): boolean
+export declare function string_enum_identity(color: StringColor): StringColor
+export declare function favorite_string_color(): StringColor

examples/basic/src/enum.zig

@@ -0,0 +1,35 @@
+pub const Color = enum(i32) {
+    Red = 1,
+    Green = 2,
+    Blue = 4,
+};
+
+pub fn enum_identity(color: Color) Color {
+    return color;
+}
+
+pub fn favorite_color() Color {
+    return .Green;
+}
+
+pub fn is_primary(color: Color) bool {
+    return switch (color) {
+        .Red, .Green, .Blue => true,
+    };
+}
+
+pub const StringColor = enum {
+    Red,
+    Green,
+    Blue,
+
+    pub const napi_string_enum = true;
+};
+
+pub fn string_enum_identity(color: StringColor) StringColor {
+    return color;
+}
+
+pub fn favorite_string_color() StringColor {
+    return .Blue;
+}

examples/basic/src/hello.zig

@@ -15,6 +15,8 @@ const arraybuffer = @import("arraybuffer.zig");
 const typedarray = @import("typedarray.zig");
 const dataview = @import("dataview.zig");
 const reference = @import("reference.zig");
+const union_value = @import("union.zig");
+const enum_value = @import("enum.zig");
 
 pub const test_i32 = number.test_i32;
 pub const test_f32 = number.test_f32;
@@ -69,6 +71,35 @@ pub const get_dataview_length = dataview.get_dataview_length;
 pub const get_dataview_first_byte = dataview.get_dataview_first_byte;
 pub const get_dataview_uint32_le = dataview.get_dataview_uint32_le;
 
+pub const union_identity = union_value.union_identity;
+pub const make_union = union_value.make_union;
+pub const union_kind = union_value.union_kind;
+pub const object_or_text_identity = union_value.object_or_text_identity;
+pub const make_object_or_text = union_value.make_object_or_text;
+pub const object_or_array_identity = union_value.object_or_array_identity;
+pub const tuple_or_text_identity = union_value.tuple_or_text_identity;
+pub const flip_flag_or_increment = union_value.flip_flag_or_increment;
+pub const color_or_text_identity = union_value.color_or_text_identity;
+pub const favorite_color_or_text = union_value.favorite_color_or_text;
+pub const maybe_text_or_count_identity = union_value.maybe_text_or_count_identity;
+pub const make_maybe_text_or_count = union_value.make_maybe_text_or_count;
+pub const buffer_or_text_identity = union_value.buffer_or_text_identity;
+pub const make_buffer_or_text = union_value.make_buffer_or_text;
+pub const arraybuffer_or_array_identity = union_value.arraybuffer_or_array_identity;
+pub const make_arraybuffer_or_array = union_value.make_arraybuffer_or_array;
+pub const payload_or_color_identity = union_value.payload_or_color_identity;
+pub const make_payload_or_color = union_value.make_payload_or_color;
+pub const payload_or_string_color_identity = union_value.payload_or_string_color_identity;
+pub const make_payload_or_string_color = union_value.make_payload_or_string_color;
+
+pub const Color = enum_value.Color;
+pub const StringColor = enum_value.StringColor;
+pub const enum_identity = enum_value.enum_identity;
+pub const favorite_color = enum_value.favorite_color;
+pub const is_primary = enum_value.is_primary;
+pub const string_enum_identity = enum_value.string_enum_identity;
+pub const favorite_string_color = enum_value.favorite_string_color;
+
 comptime {
     napi.NODE_API_MODULE("hello", @This());
 }

examples/basic/src/union.zig

@@ -0,0 +1,174 @@
+const napi = @import("napi");
+const enum_value = @import("enum.zig");
+
+pub const NumberOrText = union(enum) {
+    number: f64,
+    text: []const u8,
+};
+
+const MessagePayload = struct {
+    title: []const u8,
+    count: f64,
+};
+
+pub const ObjectOrText = union(enum) {
+    payload: MessagePayload,
+    text: []const u8,
+};
+
+pub const ObjectOrArray = union(enum) {
+    payload: MessagePayload,
+    list: []f32,
+};
+
+const NamedTuple = struct { f32, bool, []const u8 };
+
+pub const TupleOrText = union(enum) {
+    tuple: NamedTuple,
+    text: []const u8,
+};
+
+pub const FlagOrCount = union(enum) {
+    flag: bool,
+    count: f64,
+};
+
+pub const ColorOrText = union(enum) {
+    color: enum_value.Color,
+    text: []const u8,
+};
+
+pub const MaybeTextOrCount = union(enum) {
+    maybe_text: ?[]const u8,
+    count: f64,
+};
+
+pub const BufferOrText = union(enum) {
+    buffer: napi.Buffer,
+    text: []const u8,
+};
+
+pub const ArrayBufferOrArray = union(enum) {
+    arraybuffer: napi.ArrayBuffer,
+    list: []const f32,
+};
+
+pub const PayloadOrColor = union(enum) {
+    payload: MessagePayload,
+    color: enum_value.Color,
+};
+
+pub const PayloadOrStringColor = union(enum) {
+    payload: MessagePayload,
+    color: enum_value.StringColor,
+};
+
+pub fn union_identity(value: NumberOrText) NumberOrText {
+    return value;
+}
+
+pub fn make_union(is_number: bool) NumberOrText {
+    if (is_number) {
+        return .{ .number = 42 };
+    }
+    return .{ .text = "hello" };
+}
+
+pub fn union_kind(value: NumberOrText) []const u8 {
+    return switch (value) {
+        .number => "number",
+        .text => "text",
+    };
+}
+
+pub fn object_or_text_identity(value: ObjectOrText) ObjectOrText {
+    return value;
+}
+
+pub fn make_object_or_text(as_object: bool) ObjectOrText {
+    if (as_object) {
+        return .{ .payload = .{ .title = "hello", .count = 2 } };
+    }
+    return .{ .text = "plain" };
+}
+
+pub fn object_or_array_identity(value: ObjectOrArray) ObjectOrArray {
+    return value;
+}
+
+pub fn tuple_or_text_identity(value: TupleOrText) TupleOrText {
+    return value;
+}
+
+pub fn flip_flag_or_increment(value: FlagOrCount) FlagOrCount {
+    return switch (value) {
+        .flag => |flag| .{ .flag = !flag },
+        .count => |count| .{ .count = count + 1 },
+    };
+}
+
+pub fn color_or_text_identity(value: ColorOrText) ColorOrText {
+    return value;
+}
+
+pub fn favorite_color_or_text(use_color: bool) ColorOrText {
+    if (use_color) {
+        return .{ .color = .Blue };
+    }
+    return .{ .text = "fallback" };
+}
+
+pub fn maybe_text_or_count_identity(value: MaybeTextOrCount) MaybeTextOrCount {
+    return value;
+}
+
+pub fn make_maybe_text_or_count(as_text: bool) MaybeTextOrCount {
+    if (as_text) {
+        return .{ .maybe_text = null };
+    }
+    return .{ .count = 7 };
+}
+
+pub fn buffer_or_text_identity(value: BufferOrText) BufferOrText {
+    return value;
+}
+
+pub fn make_buffer_or_text(env: napi.Env, use_buffer: bool) !BufferOrText {
+    if (use_buffer) {
+        return .{ .buffer = try napi.Buffer.New(env, 16) };
+    }
+    return .{ .text = "buffer-fallback" };
+}
+
+pub fn arraybuffer_or_array_identity(value: ArrayBufferOrArray) ArrayBufferOrArray {
+    return value;
+}
+
+pub fn make_arraybuffer_or_array(env: napi.Env, use_arraybuffer: bool) !ArrayBufferOrArray {
+    if (use_arraybuffer) {
+        return .{ .arraybuffer = try napi.ArrayBuffer.New(env, 16) };
+    }
+    return .{ .list = &[_]f32{ 1, 2, 3 } };
+}
+
+pub fn payload_or_color_identity(value: PayloadOrColor) PayloadOrColor {
+    return value;
+}
+
+pub fn make_payload_or_color(use_payload: bool) PayloadOrColor {
+    if (use_payload) {
+        return .{ .payload = .{ .title = "mixed", .count = 9 } };
+    }
+    return .{ .color = .Red };
+}
+
+pub fn payload_or_string_color_identity(value: PayloadOrStringColor) PayloadOrStringColor {
+    return value;
+}
+
+pub fn make_payload_or_string_color(use_payload: bool) PayloadOrStringColor {
+    if (use_payload) {
+        return .{ .payload = .{ .title = "string-enum", .count = 3 } };
+    }
+    return .{ .color = .Green };
+}

src/build/napi-tsgen.zig

@@ -140,6 +140,11 @@ fn isDataViewType(comptime T: type) bool {
     return @hasDecl(T, "is_napi_dataview");
 }
 
+fn isStringEnumType(comptime T: type) bool {
+    if (@typeInfo(T) != .@"enum") return false;
+    return @hasDecl(T, "napi_string_enum") and @TypeOf(@field(T, "napi_string_enum")) == bool and @field(T, "napi_string_enum");
+}
+
 fn isArrayList(comptime T: type) bool {
     const info = @typeInfo(T);
     if (info != .@"struct") return false;
@@ -1003,6 +1008,9 @@ fn emitType(state: *State, comptime T: type) ![]const u8 {
             try emitEnumDecl(state, T);
             return shortTypeName(T);
         },
+        .@"union" => {
+            return try emitUnionType(state, T);
+        },
         else => {},
     }
 
@@ -1016,7 +1024,11 @@ fn emitEnumDecl(state: *State, comptime T: type) !void {
 
     try appendFmt(&state.declarations, "export declare const enum {s} {{\n", .{name});
     inline for (@typeInfo(T).@"enum".fields) |field| {
-        try appendFmt(&state.declarations, "  {s} = '{s}',\n", .{ field.name, field.name });
+        if (comptime isStringEnumType(T)) {
+            try appendFmt(&state.declarations, "  {s} = '{s}',\n", .{ field.name, field.name });
+        } else {
+            try appendFmt(&state.declarations, "  {s} = {d},\n", .{ field.name, field.value });
+        }
     }
     try append(&state.declarations, "}\n\n");
 }
@@ -1043,6 +1055,21 @@ fn emitInterfaceDecl(state: *State, comptime T: type) !void {
     try append(&state.declarations, "}\n\n");
 }
 
+fn emitUnionType(state: *State, comptime T: type) ![]const u8 {
+    const info = @typeInfo(T).@"union";
+    if (info.fields.len == 0) return "never";
+
+    var buf = StringBuilder.init(state.allocator);
+    defer buf.deinit();
+
+    inline for (info.fields, 0..) |field, idx| {
+        if (idx > 0) try append(&buf, " | ");
+        try append(&buf, try emitType(state, field.type));
+    }
+
+    return try buf.toOwnedSlice();
+}
+
 fn collectFunctionInfo(comptime T: type) struct {
     args_type: type,
     return_type: type,