session-lock: handle output unplugging better

 river/DragIcon.zig                 |  6 ++--
 river/LayerSurface.zig             |  2 +-
 river/LockManager.zig              |  2 +-
 river/LockSurface.zig              | 71 +++++++++++++++++++++++---------------
 river/Output.zig                   | 10 ++++--
 river/Root.zig                     |  2 +-
 river/Seat.zig                     |  2 +-
 river/Subsurface.zig               |  8 ++---
 river/XdgToplevel.zig              |  4 +--
 river/XwaylandOverrideRedirect.zig |  2 +-
 river/XwaylandView.zig             |  2 +-
 river/command/config.zig           |  8 ++---
 river/render.zig                   |  2 +-
 13 files changed, 70 insertions(+), 51 deletions(-)

diff --git a/river/DragIcon.zig b/river/DragIcon.zig
index 3b5ab3d..535facb 100644
--- a/river/DragIcon.zig
+++ b/river/DragIcon.zig
@@ -72,7 +72,7 @@ fn handleMap(listener: *wl.Listener(*wlr.Drag.Icon), wlr_drag_icon: *wlr.Drag.Ic
 
     wlr_drag_icon.surface.events.commit.add(&drag_icon.commit);
     var it = server.root.outputs.first;
-    while (it) |node| : (it = node.next) node.data.damage.addWhole();
+    while (it) |node| : (it = node.next) node.data.damage.?.addWhole();
 }
 
 fn handleUnmap(listener: *wl.Listener(*wlr.Drag.Icon), _: *wlr.Drag.Icon) void {
@@ -80,12 +80,12 @@ fn handleUnmap(listener: *wl.Listener(*wlr.Drag.Icon), _: *wlr.Drag.Icon) void {
 
     drag_icon.commit.link.remove();
     var it = server.root.outputs.first;
-    while (it) |node| : (it = node.next) node.data.damage.addWhole();
+    while (it) |node| : (it = node.next) node.data.damage.?.addWhole();
 }
 
 fn handleCommit(_: *wl.Listener(*wlr.Surface), _: *wlr.Surface) void {
     var it = server.root.outputs.first;
-    while (it) |node| : (it = node.next) node.data.damage.addWhole();
+    while (it) |node| : (it = node.next) node.data.damage.?.addWhole();
 }
 
 fn handleNewSubsurface(listener: *wl.Listener(*wlr.Subsurface), wlr_subsurface: *wlr.Subsurface) void {
diff --git a/river/LayerSurface.zig b/river/LayerSurface.zig
index 7c78ed3..2cc02db 100644
--- a/river/LayerSurface.zig
+++ b/river/LayerSurface.zig
@@ -164,7 +164,7 @@ fn handleCommit(listener: *wl.Listener(*wlr.Surface), _: *wlr.Surface) void {
         server.root.startTransaction();
     }
 
-    self.output.damage.addWhole();
+    self.output.damage.?.addWhole();
 }
 
 fn handleNewPopup(listener: *wl.Listener(*wlr.XdgPopup), wlr_xdg_popup: *wlr.XdgPopup) void {
diff --git a/river/LockManager.zig b/river/LockManager.zig
index 4918638..a761983 100644
--- a/river/LockManager.zig
+++ b/river/LockManager.zig
@@ -120,5 +120,5 @@ fn handleSurface(
     assert(manager.locked);
     assert(manager.lock != null);
 
-    LockSurface.create(wlr_lock_surface);
+    LockSurface.create(wlr_lock_surface, manager.lock.?);
 }
diff --git a/river/LockSurface.zig b/river/LockSurface.zig
index a9e4b81..84b0839 100644
--- a/river/LockSurface.zig
+++ b/river/LockSurface.zig
@@ -24,17 +24,19 @@ const server = &@import("main.zig").server;
 const util = @import("util.zig");
 
 const Output = @import("Output.zig");
+const Seat = @import("Seat.zig");
 const Subsurface = @import("Subsurface.zig");
 
 wlr_lock_surface: *wlr.SessionLockSurfaceV1,
+lock: *wlr.SessionLockV1,
 
 output_mode: wl.Listener(*wlr.Output) = wl.Listener(*wlr.Output).init(handleOutputMode),
 map: wl.Listener(void) = wl.Listener(void).init(handleMap),
-destroy: wl.Listener(void) = wl.Listener(void).init(handleDestroy),
+surface_destroy: wl.Listener(void) = wl.Listener(void).init(handleDestroy),
 commit: wl.Listener(*wlr.Surface) = wl.Listener(*wlr.Surface).init(handleCommit),
 new_subsurface: wl.Listener(*wlr.Subsurface) = wl.Listener(*wlr.Subsurface).init(handleSubsurface),
 
-pub fn create(wlr_lock_surface: *wlr.SessionLockSurfaceV1) void {
+pub fn create(wlr_lock_surface: *wlr.SessionLockSurfaceV1, lock: *wlr.SessionLockV1) void {
     const lock_surface = util.gpa.create(LockSurface) catch {
         wlr_lock_surface.resource.getClient().postNoMemory();
         return;
@@ -42,10 +44,13 @@ pub fn create(wlr_lock_surface: *wlr.SessionLockSurfaceV1) void {
 
     lock_surface.* = .{
         .wlr_lock_surface = wlr_lock_surface,
+        .lock = lock,
     };
+    wlr_lock_surface.data = @ptrToInt(lock_surface);
+
     wlr_lock_surface.output.events.mode.add(&lock_surface.output_mode);
     wlr_lock_surface.events.map.add(&lock_surface.map);
-    wlr_lock_surface.events.destroy.add(&lock_surface.destroy);
+    wlr_lock_surface.events.destroy.add(&lock_surface.surface_destroy);
     wlr_lock_surface.surface.events.commit.add(&lock_surface.commit);
     wlr_lock_surface.surface.events.new_subsurface.add(&lock_surface.new_subsurface);
 
@@ -54,6 +59,38 @@ pub fn create(wlr_lock_surface: *wlr.SessionLockSurfaceV1) void {
     Subsurface.handleExisting(wlr_lock_surface.surface, .{ .lock_surface = lock_surface });
 }
 
+pub fn destroy(lock_surface: *LockSurface) void {
+    lock_surface.output().lock_surface = null;
+    if (lock_surface.output().damage) |damage| damage.addWhole();
+
+    {
+        var surface_it = lock_surface.lock.surfaces.iterator(.forward);
+        const new_focus: Seat.FocusTarget = while (surface_it.next()) |surface| {
+            if (surface != lock_surface.wlr_lock_surface)
+                break .{ .lock_surface = @intToPtr(*LockSurface, surface.data) };
+        } else .none;
+
+        var seat_it = server.input_manager.seats.first;
+        while (seat_it) |node| : (seat_it = node.next) {
+            const seat = &node.data;
+            if (seat.focused == .lock_surface and seat.focused.lock_surface == lock_surface) {
+                seat.setFocusRaw(new_focus);
+            }
+            seat.cursor.updateState();
+        }
+    }
+
+    lock_surface.output_mode.link.remove();
+    lock_surface.map.link.remove();
+    lock_surface.surface_destroy.link.remove();
+    lock_surface.commit.link.remove();
+    lock_surface.new_subsurface.link.remove();
+
+    Subsurface.destroySubsurfaces(lock_surface.wlr_lock_surface.surface);
+
+    util.gpa.destroy(lock_surface);
+}
+
 pub fn output(lock_surface: *LockSurface) *Output {
     return @intToPtr(*Output, lock_surface.wlr_lock_surface.output.data);
 }
@@ -84,37 +121,15 @@ fn handleMap(listener: *wl.Listener(void)) void {
 }
 
 fn handleDestroy(listener: *wl.Listener(void)) void {
-    const lock_surface = @fieldParentPtr(LockSurface, "destroy", listener);
+    const lock_surface = @fieldParentPtr(LockSurface, "surface_destroy", listener);
 
-    lock_surface.output().lock_surface = null;
-    lock_surface.output().damage.addWhole();
-
-    {
-        var it = server.input_manager.seats.first;
-        while (it) |node| : (it = node.next) {
-            const seat = &node.data;
-            if (seat.focused == .lock_surface and seat.focused.lock_surface == lock_surface) {
-                seat.setFocusRaw(.none);
-            }
-            seat.cursor.updateState();
-        }
-    }
-
-    lock_surface.output_mode.link.remove();
-    lock_surface.map.link.remove();
-    lock_surface.destroy.link.remove();
-    lock_surface.commit.link.remove();
-    lock_surface.new_subsurface.link.remove();
-
-    Subsurface.destroySubsurfaces(lock_surface.wlr_lock_surface.surface);
-
-    util.gpa.destroy(lock_surface);
+    lock_surface.destroy();
 }
 
 fn handleCommit(listener: *wl.Listener(*wlr.Surface), _: *wlr.Surface) void {
     const lock_surface = @fieldParentPtr(LockSurface, "commit", listener);
 
-    lock_surface.output().damage.addWhole();
+    lock_surface.output().damage.?.addWhole();
 }
 
 fn handleSubsurface(listener: *wl.Listener(*wlr.Subsurface), subsurface: *wlr.Subsurface) void {
diff --git a/river/Output.zig b/river/Output.zig
index 15adc7e..59edff7 100644
--- a/river/Output.zig
+++ b/river/Output.zig
@@ -55,7 +55,7 @@ const State = struct {
 };
 
 wlr_output: *wlr.Output,
-damage: *wlr.OutputDamage,
+damage: ?*wlr.OutputDamage,
 
 /// All layer surfaces on the output, indexed by the layer enum.
 layers: [4]std.TailQueue(LayerSurface) = [1]std.TailQueue(LayerSurface){.{}} ** 4,
@@ -136,8 +136,8 @@ pub fn init(self: *Self, wlr_output: *wlr.Output) !void {
     wlr_output.events.enable.add(&self.enable);
     wlr_output.events.mode.add(&self.mode);
 
-    self.damage.events.frame.add(&self.frame);
-    self.damage.events.destroy.add(&self.damage_destroy);
+    self.damage.?.events.frame.add(&self.frame);
+    self.damage.?.events.destroy.add(&self.damage_destroy);
 
     // Ensure that a cursor image at the output's scale factor is loaded
     // for each seat.
@@ -431,6 +431,8 @@ fn handleDamageDestroy(listener: *wl.Listener(*wlr.OutputDamage), _: *wlr.Output
     self.frame.link.remove();
     // Ensure that it is safe to call remove() again in handleDestroy()
     self.frame.link = .{ .prev = &self.frame.link, .next = &self.frame.link };
+
+    self.damage = null;
 }
 
 fn handleDestroy(listener: *wl.Listener(*wlr.Output), _: *wlr.Output) void {
@@ -452,6 +454,8 @@ fn handleDestroy(listener: *wl.Listener(*wlr.Output), _: *wlr.Output) void {
         }
     }
 
+    if (self.lock_surface) |surface| surface.destroy();
+
     // Remove all listeners
     self.destroy.link.remove();
     self.enable.link.remove();
diff --git a/river/Root.zig b/river/Root.zig
index 21849a7..b578a25 100644
--- a/river/Root.zig
+++ b/river/Root.zig
@@ -392,7 +392,7 @@ fn commitTransaction(self: *Self) void {
         if (view_tags_changed) output.sendViewTags();
         if (urgent_tags_dirty) output.sendUrgentTags();
 
-        output.damage.addWhole();
+        output.damage.?.addWhole();
     }
     server.input_manager.updateCursorState();
     server.idle_inhibitor_manager.idleInhibitCheckActive();
diff --git a/river/Seat.zig b/river/Seat.zig
index 8d5c237..1a4b7e6 100644
--- a/river/Seat.zig
+++ b/river/Seat.zig
@@ -47,7 +47,7 @@ const XwaylandOverrideRedirect = @import("XwaylandOverrideRedirect.zig");
 const log = std.log.scoped(.seat);
 const PointerConstraint = @import("PointerConstraint.zig");
 
-const FocusTarget = union(enum) {
+pub const FocusTarget = union(enum) {
     view: *View,
     xwayland_override_redirect: *XwaylandOverrideRedirect,
     layer: *LayerSurface,
diff --git a/river/Subsurface.zig b/river/Subsurface.zig
index 3456e48..58b9380 100644
--- a/river/Subsurface.zig
+++ b/river/Subsurface.zig
@@ -37,12 +37,12 @@ pub const Parent = union(enum) {
 
     pub fn damageWholeOutput(parent: Parent) void {
         switch (parent) {
-            .xdg_toplevel => |xdg_toplevel| xdg_toplevel.view.output.damage.addWhole(),
-            .layer_surface => |layer_surface| layer_surface.output.damage.addWhole(),
-            .lock_surface => |lock_surface| lock_surface.output().damage.addWhole(),
+            .xdg_toplevel => |xdg_toplevel| xdg_toplevel.view.output.damage.?.addWhole(),
+            .layer_surface => |layer_surface| layer_surface.output.damage.?.addWhole(),
+            .lock_surface => |lock_surface| lock_surface.output().damage.?.addWhole(),
             .drag_icon => |_| {
                 var it = server.root.outputs.first;
-                while (it) |node| : (it = node.next) node.data.damage.addWhole();
+                while (it) |node| : (it = node.next) node.data.damage.?.addWhole();
             },
         }
     }
diff --git a/river/XdgToplevel.zig b/river/XdgToplevel.zig
index 5579323..7c39823 100644
--- a/river/XdgToplevel.zig
+++ b/river/XdgToplevel.zig
@@ -298,7 +298,7 @@ fn handleCommit(listener: *wl.Listener(*wlr.Surface), _: *wlr.Surface) void {
                 // before some change occured that caused shouldTrackConfigure() to return false.
                 view.dropSavedBuffers();
 
-                view.output.damage.addWhole();
+                view.output.damage.?.addWhole();
                 server.input_manager.updateCursorState();
             }
         } else {
@@ -309,7 +309,7 @@ fn handleCommit(listener: *wl.Listener(*wlr.Surface), _: *wlr.Surface) void {
             view.sendFrameDone();
         }
     } else {
-        view.output.damage.addWhole();
+        view.output.damage.?.addWhole();
         const size_changed = !std.meta.eql(view.surface_box, new_box);
         view.surface_box = new_box;
         // If the client has decided to resize itself and the view is floating,
diff --git a/river/XwaylandOverrideRedirect.zig b/river/XwaylandOverrideRedirect.zig
index 9710824..6b61100 100644
--- a/river/XwaylandOverrideRedirect.zig
+++ b/river/XwaylandOverrideRedirect.zig
@@ -130,7 +130,7 @@ fn handleUnmap(listener: *wl.Listener(*wlr.XwaylandSurface), _: *wlr.XwaylandSur
 
 fn handleCommit(_: *wl.Listener(*wlr.Surface), _: *wlr.Surface) void {
     var it = server.root.outputs.first;
-    while (it) |node| : (it = node.next) node.data.damage.addWhole();
+    while (it) |node| : (it = node.next) node.data.damage.?.addWhole();
 }
 
 fn handleSetOverrideRedirect(
diff --git a/river/XwaylandView.zig b/river/XwaylandView.zig
index 84c8d7c..04ed396 100644
--- a/river/XwaylandView.zig
+++ b/river/XwaylandView.zig
@@ -300,7 +300,7 @@ fn handleSetOverrideRedirect(
 fn handleCommit(listener: *wl.Listener(*wlr.Surface), surface: *wlr.Surface) void {
     const self = @fieldParentPtr(Self, "commit", listener);
 
-    self.view.output.damage.addWhole();
+    self.view.output.damage.?.addWhole();
 
     self.view.surface_box = .{
         .x = 0,
diff --git a/river/command/config.zig b/river/command/config.zig
index d2d2b43..83a2d96 100644
--- a/river/command/config.zig
+++ b/river/command/config.zig
@@ -48,7 +48,7 @@ pub fn backgroundColor(
     server.config.background_color = try parseRgba(args[1]);
 
     var it = server.root.outputs.first;
-    while (it) |node| : (it = node.next) node.data.damage.addWhole();
+    while (it) |node| : (it = node.next) node.data.damage.?.addWhole();
 }
 
 pub fn borderColorFocused(
@@ -62,7 +62,7 @@ pub fn borderColorFocused(
     server.config.border_color_focused = try parseRgba(args[1]);
 
     var it = server.root.outputs.first;
-    while (it) |node| : (it = node.next) node.data.damage.addWhole();
+    while (it) |node| : (it = node.next) node.data.damage.?.addWhole();
 }
 
 pub fn borderColorUnfocused(
@@ -76,7 +76,7 @@ pub fn borderColorUnfocused(
     server.config.border_color_unfocused = try parseRgba(args[1]);
 
     var it = server.root.outputs.first;
-    while (it) |node| : (it = node.next) node.data.damage.addWhole();
+    while (it) |node| : (it = node.next) node.data.damage.?.addWhole();
 }
 
 pub fn borderColorUrgent(
@@ -90,7 +90,7 @@ pub fn borderColorUrgent(
     server.config.border_color_urgent = try parseRgba(args[1]);
 
     var it = server.root.outputs.first;
-    while (it) |node| : (it = node.next) node.data.damage.addWhole();
+    while (it) |node| : (it = node.next) node.data.damage.?.addWhole();
 }
 
 pub fn setCursorWarp(
diff --git a/river/render.zig b/river/render.zig
index 4581675..10104d5 100644
--- a/river/render.zig
+++ b/river/render.zig
@@ -51,7 +51,7 @@ pub fn renderOutput(output: *Output) void {
     var damage_region: pixman.Region32 = undefined;
     damage_region.init();
     defer damage_region.deinit();
-    output.damage.attachRender(&needs_frame, &damage_region) catch {
+    output.damage.?.attachRender(&needs_frame, &damage_region) catch {
         log.err("failed to attach renderer", .{});
         return;
     };