TextInput: handle multiple text inputs correctly

The protocol states that we must send enter and leave to all text input
objects if the client has created multiple.

Only one text input is allowed to be activated by the client per seat
however.

 river/InputManager.zig |  2 +-
 river/InputRelay.zig   | 81 +++++++++++++++++++++++++-------------------------
 river/Seat.zig         |  2 +-
 river/TextInput.zig    | 39 ++++++++++++------------
 4 files changed, 62 insertions(+), 62 deletions(-)

diff --git a/river/InputManager.zig b/river/InputManager.zig
index c202c28..1e4df90 100644
--- a/river/InputManager.zig
+++ b/river/InputManager.zig
@@ -197,7 +197,7 @@ fn handleNewInputMethod(
     log.debug("new input method on seat {s}", .{relay.seat.wlr_seat.name});
 
     if (seat.focused.surface()) |surface| {
-        relay.setSurfaceFocus(surface);
+        relay.focus(surface);
     }
 }
 
diff --git a/river/InputRelay.zig b/river/InputRelay.zig
index f5fad43..99d4f34 100644
--- a/river/InputRelay.zig
+++ b/river/InputRelay.zig
@@ -40,6 +40,8 @@ seat: *Seat,
 text_inputs: std.TailQueue(TextInput) = .{},
 
 input_method: ?*wlr.InputMethodV2 = null,
+/// The currently enabled text input for the currently focused surface.
+text_input: ?*TextInput = null,
 
 input_method_commit: wl.Listener(*wlr.InputMethodV2) =
     wl.Listener(*wlr.InputMethodV2).init(handleInputMethodCommit),
@@ -63,7 +65,7 @@ fn handleInputMethodCommit(
     assert(input_method == self.input_method);
 
     if (!input_method.client_active) return;
-    const text_input = self.getFocusedTextInput() orelse return;
+    const text_input = self.text_input orelse return;
 
     if (input_method.current.preedit.text) |preedit_text| {
         text_input.wlr_text_input.sendPreeditString(
@@ -102,9 +104,7 @@ fn handleInputMethodDestroy(
 
     self.input_method = null;
 
-    if (self.getFocusedTextInput()) |text_input| {
-        text_input.wlr_text_input.sendLeave();
-    }
+    self.focus(null);
 }
 
 fn handleInputMethodGrabKeyboard(
@@ -131,30 +131,23 @@ fn handleInputMethodGrabKeyboardDestroy(
     }
 }
 
-pub fn getFocusedTextInput(self: *Self) ?*TextInput {
-    var it = self.text_inputs.first;
-    return while (it) |node| : (it = node.next) {
-        const text_input = &node.data;
-        if (text_input.wlr_text_input.focused_surface != null) break text_input;
-    } else null;
-}
+pub fn disableTextInput(self: *Self) void {
+    assert(self.text_input != null);
 
-pub fn disableTextInput(self: *Self, text_input: *TextInput) void {
-    if (self.input_method) |im| {
-        im.sendDeactivate();
-    } else {
-        log.debug("disabling text input but input method is gone", .{});
-        return;
+    if (self.input_method) |input_method| {
+        input_method.sendDeactivate();
+        input_method.sendDone();
     }
 
-    self.sendInputMethodState(text_input.wlr_text_input);
+    self.text_input = null;
 }
 
-pub fn sendInputMethodState(self: *Self, wlr_text_input: *wlr.TextInputV3) void {
-    const input_method = self.input_method orelse return;
+pub fn sendInputMethodState(self: *Self) void {
+    const input_method = self.input_method.?;
+    const wlr_text_input = self.text_input.?.wlr_text_input;
 
-    // TODO: only send each of those if they were modified
-    // after activation, all supported features must be sent
+    // TODO Send these events only if something changed.
+    // On activation all events must be sent for all active features.
 
     if (wlr_text_input.active_features.surrounding_text) {
         if (wlr_text_input.current.surrounding.text) |text| {
@@ -178,33 +171,39 @@ pub fn sendInputMethodState(self: *Self, wlr_text_input: *wlr.TextInputV3) void
     input_method.sendDone();
 }
 
-pub fn setSurfaceFocus(self: *Self, wlr_surface: ?*wlr.Surface) void {
-    var new_text_input: ?*TextInput = null;
+pub fn focus(self: *Self, new_focus: ?*wlr.Surface) void {
+    // Send leave events
+    {
+        var it = self.text_inputs.first;
+        while (it) |node| : (it = node.next) {
+            const text_input = &node.data;
 
-    var it = self.text_inputs.first;
-    while (it) |node| : (it = node.next) {
-        const text_input = &node.data;
-        if (text_input.wlr_text_input.focused_surface) |surface| {
-            if (wlr_surface != surface) {
-                text_input.relay.disableTextInput(text_input);
+            if (text_input.wlr_text_input.focused_surface) |surface| {
+                // This function should not be called unless focus changes
+                assert(surface != new_focus);
                 text_input.wlr_text_input.sendLeave();
-            } else {
-                log.debug("input relay setSurfaceFocus already focused", .{});
-                continue;
             }
         }
+    }
 
-        if (wlr_surface) |surface| {
-            if (text_input.wlr_text_input.resource.getClient() == surface.resource.getClient()) {
-                assert(new_text_input == null);
-                new_text_input = text_input;
-            }
-        }
+    // Clear currently enabled text input
+    if (self.text_input != null) {
+        self.disableTextInput();
     }
 
-    if (new_text_input) |text_input| {
+    // Send enter events if we have an input method.
+    // No text input for the new surface should be enabled yet as the client
+    // should wait until it receives an enter event.
+    if (new_focus) |surface| {
         if (self.input_method != null) {
-            text_input.wlr_text_input.sendEnter(wlr_surface.?);
+            var it = self.text_inputs.first;
+            while (it) |node| : (it = node.next) {
+                const text_input = &node.data;
+
+                if (text_input.wlr_text_input.resource.getClient() == surface.resource.getClient()) {
+                    text_input.wlr_text_input.sendEnter(surface);
+                }
+            }
         }
     }
 }
diff --git a/river/Seat.zig b/river/Seat.zig
index 880dc3c..6058dd1 100644
--- a/river/Seat.zig
+++ b/river/Seat.zig
@@ -268,7 +268,7 @@ pub fn setFocusRaw(self: *Self, new_focus: FocusTarget) void {
     }
 
     self.keyboardEnterOrLeave(target_surface);
-    self.relay.setSurfaceFocus(target_surface);
+    self.relay.focus(target_surface);
 
     if (target_surface) |surface| {
         const pointer_constraints = server.input_manager.pointer_constraints;
diff --git a/river/TextInput.zig b/river/TextInput.zig
index 3f7e794..3375500 100644
--- a/river/TextInput.zig
+++ b/river/TextInput.zig
@@ -56,46 +56,47 @@ pub fn init(self: *Self, relay: *InputRelay, wlr_text_input: *wlr.TextInputV3) v
 fn handleEnable(listener: *wl.Listener(*wlr.TextInputV3), _: *wlr.TextInputV3) void {
     const self = @fieldParentPtr(Self, "enable", listener);
 
-    if (self.relay.input_method) |im| {
-        im.sendActivate();
-    } else {
-        log.debug("enabling text input but input method is gone", .{});
+    if (self.relay.text_input != null) {
+        log.err("client requested to enable more than one text input on a single seat, ignoring request", .{});
         return;
     }
 
-    // must send surrounding_text if supported
-    // must send content_type if supported
-    self.relay.sendInputMethodState(self.wlr_text_input);
+    self.relay.text_input = self;
+
+    if (self.relay.input_method) |input_method| {
+        input_method.sendActivate();
+        self.relay.sendInputMethodState();
+    }
 }
 
 fn handleCommit(listener: *wl.Listener(*wlr.TextInputV3), _: *wlr.TextInputV3) void {
     const self = @fieldParentPtr(Self, "commit", listener);
-    if (!self.wlr_text_input.current_enabled) {
-        log.debug("inactive text input tried to commit an update", .{});
+
+    if (self.relay.text_input != self) {
+        log.err("inactive text input tried to commit an update, client bug?", .{});
         return;
     }
-    log.debug("text input committed update", .{});
-    if (self.relay.input_method == null) {
-        log.debug("committed text input but input method is gone", .{});
-        return;
+
+    if (self.relay.input_method != null) {
+        self.relay.sendInputMethodState();
     }
-    self.relay.sendInputMethodState(self.wlr_text_input);
 }
 
 fn handleDisable(listener: *wl.Listener(*wlr.TextInputV3), _: *wlr.TextInputV3) void {
     const self = @fieldParentPtr(Self, "disable", listener);
-    if (self.wlr_text_input.focused_surface == null) {
-        log.debug("disabling text input, but no longer focused", .{});
-        return;
+
+    if (self.relay.text_input == self) {
+        self.relay.disableTextInput();
     }
-    self.relay.disableTextInput(self);
 }
 
 fn handleDestroy(listener: *wl.Listener(*wlr.TextInputV3), _: *wlr.TextInputV3) void {
     const self = @fieldParentPtr(Self, "destroy", listener);
     const node = @fieldParentPtr(std.TailQueue(Self).Node, "data", self);
 
-    if (self.wlr_text_input.current_enabled) self.relay.disableTextInput(self);
+    if (self.relay.text_input == self) {
+        self.relay.disableTextInput();
+    }
 
     self.enable.link.remove();
     self.commit.link.remove();