Set DPMS and remove fade surfaces in one place

src/fade_black.rs

@@ -11,10 +11,7 @@ use wayland_protocols::wp::{
     single_pixel_buffer::v1::client::wp_single_pixel_buffer_manager_v1,
     viewporter::client::wp_viewport,
 };
-use wayland_protocols_wlr::{
+use wayland_protocols_wlr::layer_shell::v1::client::{zwlr_layer_shell_v1, zwlr_layer_surface_v1};
-    layer_shell::v1::client::{zwlr_layer_shell_v1, zwlr_layer_surface_v1},
-    output_power_management::v1::client::zwlr_output_power_v1,
-};
 
 use crate::{State, StateInner};
 
@@ -133,13 +130,13 @@ impl Dispatch<wl_callback::WlCallback, wl_surface::WlSurface> for State {
                     if let Some(fade_surface) = &mut output.fade_surface {
                         if &fade_surface.surface == surface {
                             if fade_surface.is_done() {
-                                output
-                                    .output_power
-                                    .set_mode(zwlr_output_power_v1::Mode::Off);
-                                output.fade_surface = None;
-
                                 // All outputs are done fading
-                                if state.outputs.iter().all(|o| o.fade_surface.is_none()) {
+                                if state
+                                    .outputs
+                                    .iter()
+                                    .flat_map(|o| o.fade_surface.as_ref())
+                                    .all(|s| s.is_done())
+                                {
                                     state.fade_done();
                                 }
                             } else {

src/main.rs

@@ -150,6 +150,13 @@ impl State {
 
     // Fade surfaces on all outputs have finished fading out
     fn fade_done(&mut self) {
+        for output in &mut self.outputs {
+            output
+                .output_power
+                .set_mode(zwlr_output_power_v1::Mode::Off);
+            output.fade_surface = None;
+        }
+
         if let Some(command) = self
             .system_actions
             .get(&shortcuts::action::System::LockScreen)