From e95b1707c1acba18e43012ff47e97381b779b8b0 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Mon, 12 Feb 2024 21:21:35 -0800
Subject: [PATCH] core: fallback to heap allocation for long preedit inputs

Fixes #1514

We previously required all preedit inputs to fit into the small copied
message size. That's true for 99% of all inputs, but if a long pre-edit
input comes in, this may not be true. We should try the small array
fast-path but fall back to heap allocation if we must.
---
 src/Surface.zig        | 50 ++++++++++++++++++++++++++++++++++--------
 src/termio/message.zig | 16 ++++++++++++++
 2 files changed, 57 insertions(+), 9 deletions(-)

diff --git a/src/Surface.zig b/src/Surface.zig
index 0422428fa..61e881cf4 100644
--- a/src/Surface.zig
+++ b/src/Surface.zig
@@ -1486,18 +1486,50 @@ pub fn keyCallback(
         break :event copy;
     };
 
-    var data: termio.Message.WriteReq.Small.Array = undefined;
-    const seq = try enc.encode(&data);
-    if (seq.len == 0) return .ignored;
+    const write_req: termio.Message.WriteReq = req: {
+        // Try to write the input into a small array. This fits almost
+        // every scenario. Larger situations can happen due to long
+        // pre-edits.
+        var data: termio.Message.WriteReq.Small.Array = undefined;
+        if (enc.encode(&data)) |seq| {
+            // Special-case: we did nothing.
+            if (seq.len == 0) return .ignored;
 
-    _ = self.io_thread.mailbox.push(.{
-        .write_small = .{
-            .data = data,
-            .len = @intCast(seq.len),
-        },
+            break :req .{ .small = .{
+                .data = data,
+                .len = @intCast(seq.len),
+            } };
+        } else |err| switch (err) {
+            // Means we need to allocate
+            error.OutOfMemory => {},
+            else => return err,
+        }
+
+        // We need to allocate. We allocate double the UTF-8 length
+        // or double the small array size, whichever is larger. That's
+        // a heuristic that should work. The only scenario I know while
+        // typing this where we don't have enough space is a long preedit,
+        // and in that case the size we need is exactly the UTF-8 length,
+        // so the double is being safe.
+        const buf = try self.alloc.alloc(u8, @max(
+            event.utf8.len * 2,
+            data.len * 2,
+        ));
+        defer self.alloc.free(buf);
+
+        // This results in a double allocation but this is such an unlikely
+        // path the performance impact is unimportant.
+        const seq = try enc.encode(buf);
+        break :req try termio.Message.WriteReq.init(self.alloc, seq);
+    };
+
+    _ = self.io_thread.mailbox.push(switch (write_req) {
+        .small => |v| .{ .write_small = v },
+        .stable => |v| .{ .write_stable = v },
+        .alloc => |v| .{ .write_alloc = v },
     }, .{ .forever = {} });
     if (insp_ev) |*ev| {
-        ev.pty = self.alloc.dupe(u8, seq) catch |err| err: {
+        ev.pty = self.alloc.dupe(u8, write_req.slice()) catch |err| err: {
             log.warn("error copying pty data for inspector err={}", .{err});
             break :err "";
         };
diff --git a/src/termio/message.zig b/src/termio/message.zig
index ce6db8238..c91f3b5d6 100644
--- a/src/termio/message.zig
+++ b/src/termio/message.zig
@@ -162,6 +162,22 @@ pub fn MessageData(comptime Elem: type, comptime small_size: comptime_int) type
                 else => unreachable,
             }
         }
+
+        pub fn deinit(self: Self) void {
+            switch (self) {
+                .small, .stable => {},
+                .alloc => |v| v.alloc.free(v.alloc.data),
+            }
+        }
+
+        /// Returns a const slice of the data pointed to by this request.
+        pub fn slice(self: Self) []const Elem {
+            return switch (self) {
+                .small => |v| v.data[0..v.len],
+                .stable => |v| v,
+                .alloc => |v| v.data,
+            };
+        }
     };
 }