std.os.windows: UTF-8 to UTF-16 conversion for Windows console output

Author: Prcuvu

lib/std/os.zig

@@ -1062,6 +1062,10 @@ pub const WriteError = error{
 
     /// Connection reset by peer.
     ConnectionResetByPeer,
+
+    /// This error occurs when trying to write UTF-8 text to a Windows console,
+    /// and the UTF-8 to UTF-16 conversion fails. Windows-only.
+    InvalidUtf8,
 } || UnexpectedError;
 
 /// Write to a file descriptor.

lib/std/os/windows.zig

@@ -445,6 +445,9 @@ pub fn GetQueuedCompletionStatusEx(
 
 pub fn CloseHandle(hObject: HANDLE) void {
     assert(ntdll.NtClose(hObject) == .SUCCESS);
+    if (IsConsoleHandle(hObject)) {
+        _ = RemoveConsoleHandleData(hObject) catch {};
+    }
 }
 
 pub fn FindClose(hFindFile: HANDLE) void {
@@ -547,6 +550,9 @@ pub const WriteFileError = error{
     /// a portion of the file.
     LockViolation,
     Unexpected,
+    /// This error occurs when trying to write UTF-8 text to a Windows console,
+    /// and the UTF-8 to UTF-16 conversion fails.
+    InvalidUtf8,
 };
 
 pub fn WriteFile(
@@ -617,21 +623,108 @@ pub fn WriteFile(
             break :blk &overlapped_data;
         } else null;
         const adjusted_len = math.cast(u32, bytes.len) orelse maxInt(u32);
-        if (kernel32.WriteFile(handle, bytes.ptr, adjusted_len, &bytes_written, overlapped) == 0) {
+        if (IsConsoleHandle(handle)) {
+            assert(offset == null);
+            bytes_written = WriteConsoleWithUtf8ToUtf16Conversion(handle, bytes) catch |err| return err;
+        } else {
+            if (kernel32.WriteFile(handle, bytes.ptr, adjusted_len, &bytes_written, overlapped) == 0) {
+                switch (kernel32.GetLastError()) {
+                    .INVALID_USER_BUFFER => return error.SystemResources,
+                    .NOT_ENOUGH_MEMORY => return error.SystemResources,
+                    .OPERATION_ABORTED => return error.OperationAborted,
+                    .NOT_ENOUGH_QUOTA => return error.SystemResources,
+                    .IO_PENDING => unreachable,
+                    .BROKEN_PIPE => return error.BrokenPipe,
+                    .INVALID_HANDLE => return error.NotOpenForWriting,
+                    .LOCK_VIOLATION => return error.LockViolation,
+                    else => |err| return unexpectedError(err),
+                }
+            }
+        }
+        return bytes_written;
+    }
+}
+
+fn WriteConsoleWithUtf8ToUtf16Conversion(handle: HANDLE, bytes: []const u8) WriteFileError!DWORD {
+    const handle_data: *ConsoleHandleData = GetConsoleHandleData(handle) catch |err| switch (err) {
+        error.ConsoleHandleLimitReached => @panic("Reached maximum number of 64 console handles."),
+        else => return error.Unexpected,
+    };
+    var bytes_written: DWORD = 0;
+    var byte_index: DWORD = 0;
+    while (byte_index < bytes.len) {
+        var utf16_buffer: [2]u16 = undefined;
+        var utf16_code_units: usize = undefined;
+        if (handle_data.utf8_buffer.bytes_used == 0) {
+            const utf8_byte_sequence_length: u3 = std.unicode.utf8ByteSequenceLength(bytes[byte_index]) catch return error.InvalidUtf8;
+            const bytes_available: usize = bytes.len - byte_index;
+            if (bytes_available < utf8_byte_sequence_length) {
+                for (0..bytes_available) |index| {
+                    handle_data.utf8_buffer.data[index] = bytes[index];
+                }
+                bytes_written += @truncate(bytes_available);
+                return bytes_written;
+            } else {
+                utf16_code_units = std.unicode.utf8ToUtf16Le(&utf16_buffer, bytes[byte_index..byte_index + utf8_byte_sequence_length]) catch return error.InvalidUtf8;
+                byte_index += utf8_byte_sequence_length;
+            }
+        } else {
+            const utf8_byte_sequence_length: u3 = std.unicode.utf8ByteSequenceLength(handle_data.utf8_buffer.data[0]) catch return error.InvalidUtf8;
+            assert(utf8_byte_sequence_length > 1 and utf8_byte_sequence_length > handle_data.utf8_buffer.bytes_used);
+            const bytes_available: usize = bytes.len - byte_index;
+            const bytes_needed: u3 = utf8_byte_sequence_length - handle_data.utf8_buffer.bytes_used;
+            if (bytes_available < bytes_needed) {
+                assert(handle_data.utf8_buffer.bytes_used + bytes_available < utf8_byte_sequence_length);
+                for (0..bytes_available) |index| {
+                    handle_data.utf8_buffer.data[handle_data.utf8_buffer.bytes_used + index] = bytes[index];
+                }
+                bytes_written += @truncate(bytes_available);
+                return bytes_written;
+            } else {
+                for (0..bytes_needed) |index| {
+                    handle_data.utf8_buffer.data[handle_data.utf8_buffer.bytes_used + index] = bytes[index];
+                }
+                utf16_code_units = std.unicode.utf8ToUtf16Le(&utf16_buffer, handle_data.utf8_buffer.data[0..utf8_byte_sequence_length]) catch return error.InvalidUtf8;
+                byte_index += bytes_needed;
+            }
+        }
+        switch (utf16_buffer[0]) {
+            0x000D => {
+                handle_data.last_character_written_is_CR = true;
+            },
+            0x000A => {
+                if (handle_data.last_character_written_is_CR) {
+                    handle_data.last_character_written_is_CR = false;
+                } else {
+                    utf16_buffer = .{ 0x000D, 0x000A };
+                    utf16_code_units = 2;
+                }
+            },
+            else => {
+                handle_data.last_character_written_is_CR = false;
+            },
+        }
+        var utf16_code_units_written: DWORD = undefined;
+        if (kernel32.WriteConsoleW(handle, &utf16_buffer, @truncate(utf16_code_units), &utf16_code_units_written, null) == FALSE) {
             switch (kernel32.GetLastError()) {
                 .INVALID_USER_BUFFER => return error.SystemResources,
                 .NOT_ENOUGH_MEMORY => return error.SystemResources,
                 .OPERATION_ABORTED => return error.OperationAborted,
                 .NOT_ENOUGH_QUOTA => return error.SystemResources,
                 .IO_PENDING => unreachable,
-                .BROKEN_PIPE => return error.BrokenPipe,
+                .BROKEN_PIPE => unreachable,
                 .INVALID_HANDLE => return error.NotOpenForWriting,
                 .LOCK_VIOLATION => return error.LockViolation,
                 else => |err| return unexpectedError(err),
             }
         }
-        return bytes_written;
+        if (utf16_code_units_written < utf16_code_units) {
+            return bytes_written;
+        } else {
+            bytes_written = byte_index;
+        }
     }
+    return bytes_written;
 }
 
 pub const SetCurrentDirectoryError = error{
@@ -5073,3 +5166,86 @@ pub fn ProcessBaseAddress(handle: HANDLE) ProcessBaseAddressError!HMODULE {
     const ppeb: *const PEB = @ptrCast(@alignCast(peb_out.ptr));
     return ppeb.ImageBaseAddress;
 }
+
+pub fn IsConsoleHandle(handle: HANDLE) bool {
+    var out: DWORD = undefined;
+    return kernel32.GetConsoleMode(handle, &out) != FALSE;
+}
+
+const ConsoleHandleData = struct {
+    is_assigned: bool = false,
+
+    handle: ?HANDLE = null,
+
+    /// On Windows NT, UTF-8 encoded strings should be converted to UTF-16 before writing to the
+    /// native console. Since write() might be called with a string fragment or even a single byte,
+    /// we have to store residual UTF-8 byte(s) without returning error. UTF-16 code unit(s) will be
+    /// generated when we have enough bytes to complete a code point.
+    utf8_buffer: Utf8Buffer = .{},
+
+    /// Used for LF to CRLF conversion before writing to the native console.
+    last_character_written_is_CR: bool = false,
+
+    pub const Utf8Buffer = struct {
+        data: [4]u8 = .{ 0x00, 0x00, 0x00, 0x00 },
+        bytes_used: u3 = 0,
+    };
+};
+
+const console_handle_data_limit = 64;
+
+var console_handle_data_array: switch (builtin.os.tag) {
+    .windows => [console_handle_data_limit]ConsoleHandleData,
+    else => void,
+} = switch (builtin.os.tag) {
+    .windows => [_]ConsoleHandleData{.{}} ** console_handle_data_limit,
+    else => void{},
+};
+
+const ConsoleHandleDataError = error{
+    OsUnsupported,
+    DataNotFound,
+    ConsoleHandleLimitReached,
+};
+
+fn GetConsoleHandleData(handle: HANDLE) ConsoleHandleDataError!*ConsoleHandleData {
+    if (builtin.os.tag == .windows) {
+        var found_unassigned: bool = false;
+        var first_unassigned_index: usize = undefined;
+        for (0..console_handle_data_limit) |index| {
+            if (console_handle_data_array[index].is_assigned) {
+                if (console_handle_data_array[index].handle == handle) {
+                    return &console_handle_data_array[index];
+                }
+            } else if (!found_unassigned) {
+                found_unassigned = true;
+                first_unassigned_index = index;
+            }
+        }
+        if (found_unassigned) {
+            console_handle_data_array[first_unassigned_index].is_assigned = true;
+            console_handle_data_array[first_unassigned_index].handle = handle;
+            console_handle_data_array[first_unassigned_index].utf8_buffer.bytes_used = 0;
+            console_handle_data_array[first_unassigned_index].last_character_written_is_CR = false;
+            return &console_handle_data_array[first_unassigned_index];
+        } else {
+            return error.ConsoleHandleLimitReached;
+        }
+    } else {
+        return error.OsUnsupported;
+    }
+}
+
+fn RemoveConsoleHandleData(handle: HANDLE) ConsoleHandleDataError!usize {
+    if (builtin.os.tag == .windows) {
+        for (0..console_handle_data_limit) |index| {
+            if (console_handle_data_array[index].is_assigned and console_handle_data_array[index].handle == handle) {
+                console_handle_data_array[index].is_assigned = false;
+                return index;
+            }
+        }
+        return error.DataNotFound;
+    } else {
+        return error.OsUnsupported;
+    }
+}

lib/std/os/windows/kernel32.zig

@@ -174,6 +174,8 @@ pub extern "kernel32" fn FillConsoleOutputCharacterW(hConsoleOutput: HANDLE, cCh
 pub extern "kernel32" fn FillConsoleOutputAttribute(hConsoleOutput: HANDLE, wAttribute: WORD, nLength: DWORD, dwWriteCoord: COORD, lpNumberOfAttrsWritten: *DWORD) callconv(WINAPI) BOOL;
 pub extern "kernel32" fn SetConsoleCursorPosition(hConsoleOutput: HANDLE, dwCursorPosition: COORD) callconv(WINAPI) BOOL;
 
+pub extern "kernel32" fn WriteConsoleW(hConsoleOutput: HANDLE, lpBuffer: *const anyopaque, nNumberOfCharsToWrite: DWORD, lpNumberOfCharsWritten: ?*DWORD, lpReserved: ?LPVOID) callconv(WINAPI) BOOL;
+
 pub extern "kernel32" fn GetCurrentDirectoryW(nBufferLength: DWORD, lpBuffer: ?[*]WCHAR) callconv(WINAPI) DWORD;
 
 pub extern "kernel32" fn GetCurrentThread() callconv(WINAPI) HANDLE;

src/link.zig

@@ -499,6 +499,7 @@ pub const File = struct {
         DeviceBusy,
         InvalidArgument,
         HotSwapUnavailableOnHostOperatingSystem,
+        InvalidUtf8,
     };
 
     /// Called from within the CodeGen to lower a local variable instantion as an unnamed