From 95d054b185503e649754640c5f27e3d62b0b9b5a Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Sat, 5 Nov 2022 17:37:21 -0700
Subject: [PATCH] allocate data for paste data if its too large

---
 src/Window.zig         |  8 ++---
 src/termio/Thread.zig  |  4 +++
 src/termio/message.zig | 72 ++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 78 insertions(+), 6 deletions(-)

diff --git a/src/Window.zig b/src/Window.zig
index 817ee2a10..d9c864372 100644
--- a/src/Window.zig
+++ b/src/Window.zig
@@ -816,10 +816,10 @@ fn keyCallback(
                             }, .{ .forever = {} });
                         }
 
-                        // TODO: NO! ALLOCATE THIS
-                        _ = win.io_thread.mailbox.push(.{
-                            .write_stable = data,
-                        }, .{ .forever = {} });
+                        _ = win.io_thread.mailbox.push(termio.message.IO.writeReq(
+                            win.alloc,
+                            data,
+                        ) catch unreachable, .{ .forever = {} });
 
                         if (bracketed) {
                             _ = win.io_thread.mailbox.push(.{
diff --git a/src/termio/Thread.zig b/src/termio/Thread.zig
index 7201101d7..9859138a1 100644
--- a/src/termio/Thread.zig
+++ b/src/termio/Thread.zig
@@ -164,6 +164,10 @@ fn drainMailbox(self: *Thread) !void {
                 .resize => |v| try self.impl.resize(v.grid_size, v.screen_size),
                 .write_small => |v| try self.impl.queueWrite(v.data[0..v.len]),
                 .write_stable => |v| try self.impl.queueWrite(v),
+                .write_alloc => |v| {
+                    defer v.alloc.free(v.data);
+                    try self.impl.queueWrite(v.data);
+                },
             }
         }
     }
diff --git a/src/termio/message.zig b/src/termio/message.zig
index e6ce2b287..45c4596a4 100644
--- a/src/termio/message.zig
+++ b/src/termio/message.zig
@@ -1,9 +1,14 @@
 const std = @import("std");
+const assert = std.debug.assert;
 const Allocator = std.mem.Allocator;
 const renderer = @import("../renderer.zig");
 const terminal = @import("../terminal/main.zig");
 
 /// The messages that can be sent to an IO thread.
+///
+/// This is not a tiny structure (~40 bytes at the time of writing this comment),
+/// but the messages are IO thread sends are also very few. At the current size
+/// we can queue 26,000 messages before consuming a MB of RAM.
 pub const IO = union(enum) {
     /// Resize the window.
     resize: struct {
@@ -16,12 +21,52 @@ pub const IO = union(enum) {
 
     /// Write where the data pointer is stable.
     write_stable: []const u8,
+
+    /// Write where the data is allocated and must be freed.
+    write_alloc: WriteReq.Alloc,
+
+    /// Return a write request for the given data. This will use
+    /// write_small if it fits or write_alloc otherwise. This should NOT
+    /// be used for stable pointers which can be manually set to write_stable.
+    pub fn writeReq(alloc: Allocator, data: anytype) !IO {
+        switch (@typeInfo(@TypeOf(data))) {
+            .Pointer => |info| {
+                assert(info.size == .Slice);
+                assert(info.child == u8);
+
+                // If it fits in our small request, do that.
+                if (data.len <= WriteReq.Small.Max) {
+                    var buf: WriteReq.Small.Array = undefined;
+                    std.mem.copy(u8, &buf, data);
+                    return IO{
+                        .write_small = .{
+                            .data = buf,
+                            .len = @intCast(u8, data.len),
+                        },
+                    };
+                }
+
+                // Otherwise, allocate
+                var buf = try alloc.dupe(u8, data);
+                errdefer alloc.free(buf);
+                return IO{
+                    .write_alloc = .{
+                        .alloc = alloc,
+                        .data = buf,
+                    },
+                };
+            },
+
+            else => unreachable,
+        }
+    }
 };
 
 /// Represents a write request.
 pub const WriteReq = union(enum) {
     pub const Small = struct {
-        pub const Array = [22]u8;
+        pub const Max = 38;
+        pub const Array = [Max]u8;
         data: Array,
         len: u8,
     };
@@ -43,8 +88,31 @@ pub const WriteReq = union(enum) {
     alloc: Alloc,
 };
 
+test {
+    std.testing.refAllDecls(@This());
+}
+
 test {
     // Ensure we don't grow our IO message size without explicitly wanting to.
     const testing = std.testing;
-    try testing.expectEqual(@as(usize, 24), @sizeOf(IO));
+    try testing.expectEqual(@as(usize, 40), @sizeOf(IO));
+}
+
+test "IO.writeReq small" {
+    const testing = std.testing;
+    const alloc = testing.allocator;
+
+    const input = "hello!";
+    const io = try IO.writeReq(alloc, @as([]const u8, input));
+    try testing.expect(io == .write_small);
+}
+
+test "IO.writeReq alloc" {
+    const testing = std.testing;
+    const alloc = testing.allocator;
+
+    const input = "hello! " ** 100;
+    const io = try IO.writeReq(alloc, @as([]const u8, input));
+    try testing.expect(io == .write_alloc);
+    io.write_alloc.alloc.free(io.write_alloc.data);
 }