gamma-control: Misc. clean ups and fixes

5 files changed, 175 insertions, 184 deletions

diff --git a/Cargo.toml b/Cargo.toml
index 7a487a0c..231ffd24 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -45,6 +45,7 @@ arrayvec = "0.7.4"
 async-channel = { version = "2.2.0", optional = true }
 async-io = { version = "1.13.0", optional = true }
 bitflags = "2.4.2"
+bytemuck = "1.14.3"
 calloop = { version = "0.13.0", features = ["executor", "futures-io"] }
 clap = { workspace = true, features = ["string"] }
 directories = "5.0.1"
@@ -73,7 +74,6 @@ tracy-client.workspace = true
 url = { version = "2.5.0", optional = true }
 xcursor = "0.3.5"
 zbus = { version = "~3.15.2", optional = true }
-bytemuck = "1.14.3"
 
 [dependencies.smithay]
 workspace = true
diff --git a/src/backend/tty.rs b/src/backend/tty.rs
index 6d954a62..6ac84ca5 100644
--- a/src/backend/tty.rs
+++ b/src/backend/tty.rs
@@ -1,6 +1,7 @@
 use std::cell::RefCell;
 use std::collections::{HashMap, HashSet};
 use std::fmt::Write;
+use std::iter::zip;
 use std::panic::{catch_unwind, AssertUnwindSafe};
 use std::path::Path;
 use std::rc::Rc;
@@ -702,6 +703,11 @@ impl Tty {
             Err(err) => debug!("error setting max bpc: {err:?}"),
         }
 
+        // Reset gamma to linear in case it was set before.
+        if let Err(err) = set_gamma_for_crtc(&device.drm, crtc, None) {
+            debug!("error resetting gamma: {err:?}");
+        }
+
         let surface = device
             .drm
             .create_surface(crtc, mode, &[connector.handle()])?;
@@ -1244,63 +1250,28 @@ impl Tty {
         }
     }
 
-    pub fn get_gamma_length(&self, output: &Output) -> Option<u32> {
+    pub fn get_gamma_size(&self, output: &Output) -> anyhow::Result<u32> {
         let tty_state = output.user_data().get::<TtyOutputState>().unwrap();
-        let device = self.devices.get(&tty_state.node)?;
-        let crtc_info = device.drm.get_crtc(tty_state.crtc.clone()).ok()?;
-        return Some(crtc_info.gamma_length());
-    }
+        let crtc = tty_state.crtc;
 
-    pub fn get_gamma(&self, output: &Output) -> Option<Vec<u16>> {
-        let tty_state = output.user_data().get::<TtyOutputState>().unwrap();
-        let device = self.devices.get(&tty_state.node)?;
-        let crtc_info = device.drm.get_crtc(tty_state.crtc.clone()).ok()?;
-        let crtc = tty_state.crtc.clone();
-        let gamma_length = crtc_info.gamma_length() as usize;
-        let mut red = vec![0; gamma_length];
-        let mut green = vec![0; gamma_length];
-        let mut blue = vec![0; gamma_length];
-        if let Err(err) = device.drm.get_gamma(crtc, &mut red, &mut green, &mut blue) {
-            warn!("error getting gamma for crtc {crtc:?}: {err:?}");
-            return None;
-        }
-        red.extend(green);
-        red.extend(blue);
-        return Some(red);
+        let device = self
+            .devices
+            .get(&tty_state.node)
+            .context("missing device")?;
+        let info = device
+            .drm
+            .get_crtc(crtc)
+            .context("error getting crtc info")?;
+        Ok(info.gamma_length())
     }
 
-    pub fn set_gamma(
-        &self,
-        niri: &mut Niri,
-        output: &Output,
-        ramp: Vec<u16>,
-    ) -> anyhow::Result<()> {
+    pub fn set_gamma(&self, output: &Output, ramp: Option<&[u16]>) -> anyhow::Result<()> {
         let tty_state = output.user_data().get::<TtyOutputState>().unwrap();
         let device = self
             .devices
             .get(&tty_state.node)
             .context("missing device")?;
-
-        let crtc_info = device.drm.get_crtc(tty_state.crtc.clone())?;
-        let crtc = tty_state.crtc.clone();
-        let gamma_length = crtc_info.gamma_length() as usize;
-
-        ensure!(ramp.len() == gamma_length * 3, "wrong gamma length");
-        let mut red = ramp.clone();
-        red.truncate(gamma_length);
-        let mut green = ramp.clone();
-        green.drain(0..gamma_length);
-        green.truncate(gamma_length);
-        let mut blue = ramp;
-        blue.drain(0..gamma_length * 2);
-        blue.truncate(gamma_length);
-
-        device
-            .drm
-            .set_gamma(crtc, &mut red, &mut green, &mut blue)
-            .context("error setting gamma")?;
-        niri.queue_redraw(output.clone());
-        Ok(())
+        set_gamma_for_crtc(&device.drm, tty_state.crtc, ramp)
     }
 
     fn refresh_ipc_outputs(&self) {
@@ -1870,6 +1841,49 @@ fn set_max_bpc(device: &DrmDevice, connector: connector::Handle, bpc: u64) -> an
     Err(anyhow!("couldn't find max bpc property"))
 }
 
+pub fn set_gamma_for_crtc(
+    device: &DrmDevice,
+    crtc: crtc::Handle,
+    ramp: Option<&[u16]>,
+) -> anyhow::Result<()> {
+    let _span = tracy_client::span!("set_gamma_for_crtc");
+
+    let info = device.get_crtc(crtc).context("error getting crtc info")?;
+    let gamma_length = info.gamma_length() as usize;
+
+    let mut temp;
+    let ramp = if let Some(ramp) = ramp {
+        ensure!(ramp.len() == gamma_length * 3, "wrong gamma length");
+        ramp
+    } else {
+        let _span = tracy_client::span!("generate linear gamma");
+
+        // The legacy API provides no way to reset the gamma, so set a linear one manually.
+        temp = vec![0u16; gamma_length * 3];
+
+        let (red, rest) = temp.split_at_mut(gamma_length);
+        let (green, blue) = rest.split_at_mut(gamma_length);
+        let denom = gamma_length as u64 - 1;
+        for (i, ((r, g), b)) in zip(zip(red, green), blue).enumerate() {
+            let value = (0xFFFFu64 * i as u64 / denom) as u16;
+            *r = value;
+            *g = value;
+            *b = value;
+        }
+
+        &temp
+    };
+
+    let (red, ramp) = ramp.split_at(gamma_length);
+    let (green, blue) = ramp.split_at(gamma_length);
+
+    device
+        .set_gamma(crtc, red, green, blue)
+        .context("error setting gamma")?;
+
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs
index 54a9b451..d027194c 100644
--- a/src/handlers/mod.rs
+++ b/src/handlers/mod.rs
@@ -8,7 +8,6 @@ use std::os::fd::OwnedFd;
 use std::sync::Arc;
 use std::thread;
 
-use anyhow::anyhow;
 use smithay::backend::allocator::dmabuf::Dmabuf;
 use smithay::backend::drm::DrmNode;
 use smithay::desktop::{PopupKind, PopupManager};
@@ -16,7 +15,6 @@ use smithay::input::pointer::{CursorIcon, CursorImageStatus, PointerHandle};
 use smithay::input::{keyboard, Seat, SeatHandler, SeatState};
 use smithay::output::Output;
 use smithay::reexports::wayland_protocols::xdg::shell::server::xdg_toplevel;
-use smithay::reexports::wayland_server::backend::ObjectId;
 use smithay::reexports::wayland_server::protocol::wl_data_source::WlDataSource;
 use smithay::reexports::wayland_server::protocol::wl_output::WlOutput;
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
@@ -449,52 +447,27 @@ impl GammaControlHandler for State {
         &mut self.niri.gamma_control_manager_state
     }
 
-    fn set_gamma(
-        &mut self,
-        wl_output: &WlOutput,
-        ramp: Vec<u16>,
-        gamma_size: u32,
-    ) -> anyhow::Result<()> {
-        if ramp.len() != gamma_size as usize * 3 {
-            error!(
-                "gamma length wrong (expected {}, got {})",
-                gamma_size,
-                ramp.len()
-            );
-            return Err(anyhow!("Gamma length wrong"));
-        }
-
-        let Some(output) = Output::from_resource(wl_output) else {
-            return Err(anyhow!("No Output matching WlOutput"));
-        };
-
-        self.backend.tty().set_gamma(&mut self.niri, &output, ramp)
-    }
-
-    fn get_gamma(&mut self, wl_output: &WlOutput) -> Option<Vec<u16>> {
-        let output = Output::from_resource(wl_output)?;
-
-        let gamma_ramp = self.backend.tty().get_gamma(&output);
-        if gamma_ramp.is_none() {
-            warn!("Failed to get gamma ramp");
+    fn get_gamma_size(&mut self, output: &Output) -> Option<u32> {
+        match self.backend.tty().get_gamma_size(output) {
+            Ok(size) => Some(size),
+            Err(err) => {
+                warn!(
+                    "error getting gamma size for output {}: {err:?}",
+                    output.name()
+                );
+                None
+            }
         }
-        gamma_ramp
     }
 
-    fn destroy(&mut self, output_id: ObjectId) {
-        self.niri
-            .gamma_control_manager_state
-            .destroy_gamma_control(output_id)
-    }
-
-    fn get_gamma_size(&mut self, wl_output: &WlOutput) -> Option<u32> {
-        let output = Output::from_resource(wl_output)?;
-        let gamma_size = self.backend.tty().get_gamma_length(&output);
-        if gamma_size.is_none() {
-            warn!("Failed to get gamma size");
+    fn set_gamma(&mut self, output: &Output, ramp: Option<&[u16]>) -> Option<()> {
+        match self.backend.tty().set_gamma(output, ramp) {
+            Ok(()) => Some(()),
+            Err(err) => {
+                warn!("error setting gamma for output {}: {err:?}", output.name());
+                None
+            }
         }
-        gamma_size
     }
 }
-
 delegate_gamma_control!(State);
diff --git a/src/niri.rs b/src/niri.rs
index 1e5b5c16..37c1e40c 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -931,14 +931,11 @@ impl Niri {
         let viewporter_state = ViewporterState::new::<State>(&display_handle);
         let xdg_foreign_state = XdgForeignState::new::<State>(&display_handle);
 
-        let gamma_control_manager_state = GammaControlManagerState::new::<State, _>(
-            &display_handle,
-            match backend {
-                Backend::Tty(_) => true,
-                _ => false,
-            },
-            |client| !client.get_data::<ClientState>().unwrap().restricted,
-        );
+        let is_tty = matches!(backend, Backend::Tty(_));
+        let gamma_control_manager_state =
+            GammaControlManagerState::new::<State, _>(&display_handle, move |client| {
+                is_tty && !client.get_data::<ClientState>().unwrap().restricted
+            });
 
         let mut seat: Seat<State> = seat_state.new_wl_seat(&display_handle, backend.seat_name());
         seat.add_keyboard(
@@ -1326,6 +1323,7 @@ impl Niri {
         self.layout.remove_output(output);
         self.global_space.unmap_output(output);
         self.reposition_outputs(None);
+        self.gamma_control_manager_state.output_removed(output);
 
         let state = self.output_state.remove(output).unwrap();
         self.output_by_name.remove(&output.name()).unwrap();
diff --git a/src/protocols/gamma_control.rs b/src/protocols/gamma_control.rs
index 947b5e28..07e5a848 100644
--- a/src/protocols/gamma_control.rs
+++ b/src/protocols/gamma_control.rs
@@ -2,9 +2,9 @@ use std::collections::HashMap;
 use std::fs::File;
 use std::io::Read;
 
+use smithay::output::Output;
 use smithay::reexports::wayland_protocols_wlr;
-use smithay::reexports::wayland_server::backend::{ClientId, ObjectId};
-use smithay::reexports::wayland_server::protocol::wl_output::WlOutput;
+use smithay::reexports::wayland_server::backend::ClientId;
 use smithay::reexports::wayland_server::{
     Client, DataInit, Dispatch, DisplayHandle, GlobalDispatch, New, Resource,
 };
@@ -17,36 +17,26 @@ use zwlr_gamma_control_v1::ZwlrGammaControlV1;
 const VERSION: u32 = 1;
 
 pub struct GammaControlManagerState {
-    gamma_controls: HashMap<WlOutput, ZwlrGammaControlV1>,
+    // Active gamma controls only. Failed ones are removed.
+    gamma_controls: HashMap<Output, ZwlrGammaControlV1>,
 }
 
 pub struct GammaControlManagerGlobalData {
     filter: Box<dyn for<'c> Fn(&'c Client) -> bool + Send + Sync>,
-    can_view: bool,
 }
 
 pub trait GammaControlHandler {
     fn gamma_control_manager_state(&mut self) -> &mut GammaControlManagerState;
-    fn set_gamma(
-        &mut self,
-        output: &WlOutput,
-        ramp: Vec<u16>,
-        gamma_size: u32,
-    ) -> anyhow::Result<()>;
-    fn get_gamma(&mut self, output: &WlOutput) -> Option<Vec<u16>>;
-    fn destroy(&mut self, output_id: ObjectId);
-    fn get_gamma_size(&mut self, output: &WlOutput) -> Option<u32>;
+    fn get_gamma_size(&mut self, output: &Output) -> Option<u32>;
+    fn set_gamma(&mut self, output: &Output, ramp: Option<&[u16]>) -> Option<()>;
 }
 
 pub struct GammaControlState {
-    gamma_size: Option<u32>,
-    previous_gamma_ramp: Option<Vec<u16>>,
-    output: WlOutput,
-    failed: bool,
+    gamma_size: u32,
 }
 
 impl GammaControlManagerState {
-    pub fn new<D, F>(display: &DisplayHandle, can_view: bool, filter: F) -> Self
+    pub fn new<D, F>(display: &DisplayHandle, filter: F) -> Self
     where
         D: GlobalDispatch<ZwlrGammaControlManagerV1, GammaControlManagerGlobalData>,
         D: Dispatch<ZwlrGammaControlManagerV1, ()>,
@@ -57,7 +47,6 @@ impl GammaControlManagerState {
     {
         let global_data = GammaControlManagerGlobalData {
             filter: Box::new(filter),
-            can_view,
         };
         display.create_global::<D, ZwlrGammaControlManagerV1, _>(VERSION, global_data);
 
@@ -65,8 +54,11 @@ impl GammaControlManagerState {
             gamma_controls: HashMap::new(),
         }
     }
-    pub fn destroy_gamma_control(&mut self, output_id: ObjectId) {
-        self.gamma_controls.remove(&output_id);
+
+    pub fn output_removed(&mut self, output: &Output) {
+        if let Some(gamma_control) = self.gamma_controls.remove(output) {
+            gamma_control.failed();
+        }
     }
 }
 
@@ -91,7 +83,7 @@ where
     }
 
     fn can_view(client: Client, global_data: &GammaControlManagerGlobalData) -> bool {
-        global_data.can_view && (global_data.filter)(&client)
+        (global_data.filter)(&client)
     }
 }
 
@@ -113,45 +105,30 @@ where
     ) {
         match request {
             zwlr_gamma_control_manager_v1::Request::GetGammaControl { id, output } => {
-                let gamma_size = state.get_gamma_size(&output);
-                let previous_gamma_ramp = state.get_gamma(&output);
-
-                if state
-                    .gamma_control_manager_state()
-                    .gamma_controls
-                    .contains_key(&output)
-                    || gamma_size.is_none()
-                    || previous_gamma_ramp.is_none()
-                {
-                    data_init
-                        .init(
-                            id,
-                            GammaControlState {
-                                gamma_size: gamma_size.clone(),
-                                previous_gamma_ramp: None,
-                                output: output.clone(),
-                                failed: true,
-                            },
-                        )
-                        .failed();
-                    return;
+                if let Some(output) = Output::from_resource(&output) {
+                    // We borrow state in the middle.
+                    #[allow(clippy::map_entry)]
+                    if !state
+                        .gamma_control_manager_state()
+                        .gamma_controls
+                        .contains_key(&output)
+                    {
+                        if let Some(gamma_size) = state.get_gamma_size(&output) {
+                            let zwlr_gamma_control =
+                                data_init.init(id, GammaControlState { gamma_size });
+                            zwlr_gamma_control.gamma_size(gamma_size);
+                            state
+                                .gamma_control_manager_state()
+                                .gamma_controls
+                                .insert(output, zwlr_gamma_control);
+                            return;
+                        }
+                    }
                 }
 
-                let zwlr_gamma_control = data_init.init(
-                    id,
-                    GammaControlState {
-                        gamma_size: gamma_size.clone(),
-                        previous_gamma_ramp,
-                        output: output.clone(),
-                        failed: false,
-                    },
-                );
-
-                zwlr_gamma_control.gamma_size(gamma_size.unwrap());
-                state
-                    .gamma_control_manager_state()
-                    .gamma_controls
-                    .insert(output, zwlr_gamma_control);
+                data_init
+                    .init(id, GammaControlState { gamma_size: 0 })
+                    .failed();
             }
             zwlr_gamma_control_manager_v1::Request::Destroy => (),
             _ => unreachable!(),
@@ -176,24 +153,55 @@ where
     ) {
         match request {
             zwlr_gamma_control_v1::Request::SetGamma { fd } => {
-                if data.failed {
+                let gamma_controls = &mut state.gamma_control_manager_state().gamma_controls;
+                let Some((output, _)) = gamma_controls.iter().find(|(_, x)| *x == resource) else {
                     return;
-                }
-                debug!("setting gamma for output {:?}", data.output);
-                let buf = &mut Vec::new();
-                if File::from(fd).read_to_end(buf).is_err() {
-                    warn!("failed to read gamma data for output {:?}", data.output);
-                    resource.failed();
-                    return;
-                }
+                };
+                let output = output.clone();
+
+                trace!("setting gamma for output {}", output.name());
 
-                let gamma = bytemuck::cast_slice(buf).to_vec();
-                let gamma_size = data.gamma_size.unwrap();
+                // Start with a u16 slice so it's aligned correctly.
+                let mut buf = vec![0u16; data.gamma_size as usize * 3];
+                let buf = bytemuck::cast_slice_mut(&mut buf);
+                let mut file = File::from(fd);
+                {
+                    let _span = tracy_client::span!("read gamma from fd");
+
+                    if let Err(err) = file.read_exact(buf) {
+                        warn!("failed to read gamma data: {err:?}");
+                        resource.failed();
+                        gamma_controls.remove(&output);
+                        let _ = state.set_gamma(&output, None);
+                        return;
+                    }
+
+                    // Verify that there's no more data.
+                    match file.read(&mut [0]) {
+                        Ok(0) => (),
+                        Ok(_) => {
+                            warn!("gamma data is too large");
+                            resource.failed();
+                            gamma_controls.remove(&output);
+                            let _ = state.set_gamma(&output, None);
+                            return;
+                        }
+                        Err(err) => {
+                            warn!("error reading gamma data: {err:?}");
+                            resource.failed();
+                            gamma_controls.remove(&output);
+                            let _ = state.set_gamma(&output, None);
+                            return;
+                        }
+                    }
+                }
+                let gamma = bytemuck::cast_slice(buf);
 
-                if let Err(err) = state.set_gamma(&data.output, gamma, gamma_size) {
-                    warn!("error setting gamma: {err:?}");
+                if state.set_gamma(&output, Some(gamma)).is_none() {
                     resource.failed();
-                    return;
+                    let gamma_controls = &mut state.gamma_control_manager_state().gamma_controls;
+                    gamma_controls.remove(&output);
+                    let _ = state.set_gamma(&output, None);
                 }
             }
             zwlr_gamma_control_v1::Request::Destroy => (),
@@ -204,19 +212,17 @@ where
     fn destroyed(
         state: &mut D,
         _client: ClientId,
-        _resource: &ZwlrGammaControlV1,
-        data: &GammaControlState,
+        resource: &ZwlrGammaControlV1,
+        _data: &GammaControlState,
     ) {
-        if data.failed {
+        let gamma_controls = &mut state.gamma_control_manager_state().gamma_controls;
+        let Some((output, _)) = gamma_controls.iter().find(|(_, x)| *x == resource) else {
             return;
-        }
-        let ramp = data.previous_gamma_ramp.as_ref().unwrap();
-
-        if let Err(err) = state.set_gamma(&data.output, ramp.to_vec(), data.gamma_size.unwrap()) {
-            warn!("error resetting gamma: {err:?}");
-        }
+        };
+        let output = output.clone();
+        gamma_controls.remove(&output);
 
-        state.destroy(data.output.id());
+        let _ = state.set_gamma(&output, None);
     }
 }