From 2809d637dba25d4d720c389af50b28b128fc8348 Mon Sep 17 00:00:00 2001
From: ridiculousfish <rf@fishshell.com>
Date: Wed, 29 Jul 2020 19:00:04 -0700
Subject: [PATCH] Introduce wwrite_to_fd

wwrite_to_fd() is a function which writes a wide string to a file
descriptor without performing any heap allocations.
---
 src/fish_tests.cpp | 44 ++++++++++++++++++++++++++++++++++++++++
 src/wutil.cpp      | 50 ++++++++++++++++++++++++++++++++++++++++++++++
 src/wutil.h        | 11 ++++++++++
 3 files changed, 105 insertions(+)

diff --git a/src/fish_tests.cpp b/src/fish_tests.cpp
index 6f4bf065f..a3692f1b8 100644
--- a/src/fish_tests.cpp
+++ b/src/fish_tests.cpp
@@ -5025,6 +5025,49 @@ static void test_wcstring_tok() {
     }
 }
 
+static void test_wwrite_to_fd() {
+    say(L"Testing wwrite_to_fd");
+    char t[] = "/tmp/fish_test_wwrite.XXXXXX";
+    if (!mktemp(t)) {
+        err(L"Unable to create temporary file");
+        return;
+    }
+    size_t sizes[] = {0, 1, 2, 3, 5, 13, 23, 64, 128, 255, 4096, 4096 * 2};
+    for (size_t size : sizes) {
+        autoclose_fd_t fd{open(t, O_RDWR | O_TRUNC | O_CREAT, 0666)};
+        if (!fd.valid()) {
+            wperror(L"open");
+            err(L"Unable to open temporary file");
+            return;
+        }
+        wcstring input{};
+        for (size_t i = 0; i < size; i++) {
+            input.push_back(wchar_t(random()));
+        }
+
+        ssize_t amt = wwrite_to_fd(input, fd.fd());
+        if (amt < 0) {
+            wperror(L"write");
+            err(L"Unable to write to temporary file");
+            return;
+        }
+        std::string narrow = wcs2string(input);
+        size_t expected_size = narrow.size();
+        do_test(amt == expected_size);
+
+        if (lseek(fd.fd(), 0, SEEK_SET) < 0) {
+            wperror(L"seek");
+            err(L"Unable to seek temporary file");
+            return;
+        }
+
+        std::string contents(expected_size, '\0');
+        ssize_t read_amt = read(fd.fd(), &contents[0], expected_size);
+        do_test(read_amt >= 0 && static_cast<size_t>(read_amt) == expected_size);
+    }
+    (void)remove(t);
+}
+
 static void test_pcre2_escape() {
     say(L"Testing escaping strings as pcre2 literals");
     // plain text should not be needlessly escaped
@@ -5928,6 +5971,7 @@ int main(int argc, char **argv) {
 
     if (should_test_function("utility_functions")) test_utility_functions();
     if (should_test_function("wcstring_tok")) test_wcstring_tok();
+    if (should_test_function("wwrite_to_fd")) test_wwrite_to_fd();
     if (should_test_function("env_vars")) test_env_vars();
     if (should_test_function("env")) test_env_snapshot();
     if (should_test_function("str_to_num")) test_str_to_num();
diff --git a/src/wutil.cpp b/src/wutil.cpp
index f046acfac..ffb212871 100644
--- a/src/wutil.cpp
+++ b/src/wutil.cpp
@@ -536,6 +536,56 @@ int wrename(const wcstring &old, const wcstring &newv) {
     return rename(old_narrow.c_str(), new_narrow.c_str());
 }
 
+ssize_t wwrite_to_fd(const wchar_t *input, size_t input_len, int fd) {
+    // Accumulate data in a local buffer.
+    char accum[512];
+    size_t accumlen{0};
+    constexpr size_t maxaccum = sizeof accum / sizeof *accum;
+
+    // Helper to perform a write to 'fd', looping as necessary.
+    // \return true on success, false on error.
+    ssize_t total_written = 0;
+    auto do_write = [fd, &total_written](const char *cursor, size_t remaining) {
+        while (remaining > 0) {
+            ssize_t samt = write(fd, cursor, remaining);
+            if (samt < 0) return false;
+            total_written += samt;
+            size_t amt = static_cast<size_t>(samt);
+            assert(amt <= remaining && "Wrote more than requested");
+            remaining -= amt;
+            cursor += amt;
+        }
+        return true;
+    };
+
+    // Helper to flush the accumulation buffer.
+    auto flush_accum = [&] {
+        if (!do_write(accum, accumlen)) return false;
+        accumlen = 0;
+        return true;
+    };
+
+    bool success = wcs2string_callback(input, input_len, [&](const char *buff, size_t len) {
+        if (len + accumlen > maxaccum) {
+            // We have to flush.
+            // Note this modifies 'accumlen'.
+            if (!flush_accum()) return false;
+        }
+        if (len + accumlen <= maxaccum) {
+            // Accumulate more.
+            memmove(accum + accumlen, buff, len);
+            accumlen += len;
+            return true;
+        } else {
+            // Too much data to even fit, just write it immediately.
+            return do_write(buff, len);
+        }
+    });
+    // Flush any remaining.
+    if (success) success = flush_accum();
+    return success ? total_written : -1;
+}
+
 /// Return one if the code point is in a Unicode private use area.
 int fish_is_pua(wint_t wc) {
     if (PUA1_START <= wc && wc < PUA1_END) return 1;
diff --git a/src/wutil.h b/src/wutil.h
index db98b7221..75a0db690 100644
--- a/src/wutil.h
+++ b/src/wutil.h
@@ -113,6 +113,17 @@ int wmkdir(const wcstring &name, int mode);
 /// Wide character version of rename.
 int wrename(const wcstring &oldName, const wcstring &newv);
 
+/// Write a wide string to a file descriptor. This avoids doing any additional allocation.
+/// This does NOT retry on EINTR or EAGAIN, it simply returns.
+/// \return -1 on error in which case errno will have been set. In this event, the number of bytes
+/// actually written cannot be obtained.
+ssize_t wwrite_to_fd(const wchar_t *s, size_t len, int fd);
+
+/// Variant of above that accepts a wcstring.
+inline ssize_t wwrite_to_fd(const wcstring &s, int fd) {
+    return wwrite_to_fd(s.c_str(), s.size(), fd);
+}
+
 #define PUA1_START 0xE000
 #define PUA1_END 0xF900
 #define PUA2_START 0xF0000