4 files changed, 54 insertions, 48 deletions

diff --git a/src/render_helpers/border.rs b/src/render_helpers/border.rs
index 6f2b7795..b0b373e2 100644
--- a/src/render_helpers/border.rs
+++ b/src/render_helpers/border.rs
@@ -8,8 +8,8 @@ use smithay::backend::renderer::utils::{CommitCounter, DamageSet};
 use smithay::utils::{Buffer, Logical, Physical, Rectangle, Scale, Transform};
 
 use super::renderer::NiriRenderer;
-use super::shader_element::{ShaderProgram, ShaderRenderElement};
-use super::shaders::{mat3_uniform, Shaders};
+use super::shader_element::ShaderRenderElement;
+use super::shaders::{mat3_uniform, ProgramType, Shaders};
 use crate::backend::tty::{TtyFrame, TtyRenderer, TtyRendererError};
 
 /// Renders a wide variety of borders and border parts.
@@ -40,7 +40,6 @@ struct Parameters {
 impl BorderRenderElement {
     #[allow(clippy::too_many_arguments)]
     pub fn new(
-        shader: &ShaderProgram,
         scale: Scale<f64>,
         area: Rectangle<i32, Logical>,
         gradient_area: Rectangle<i32, Logical>,
@@ -51,8 +50,7 @@ impl BorderRenderElement {
         border_width: f32,
         corner_radius: CornerRadius,
     ) -> Self {
-        let mut inner = ShaderRenderElement::empty(Kind::Unspecified);
-        inner.update_shader(Some(shader));
+        let inner = ShaderRenderElement::empty(ProgramType::Border, Kind::Unspecified);
         let mut rv = Self {
             inner,
             params: Parameters {
@@ -72,7 +70,7 @@ impl BorderRenderElement {
     }
 
     pub fn empty() -> Self {
-        let inner = ShaderRenderElement::empty(Kind::Unspecified);
+        let inner = ShaderRenderElement::empty(ProgramType::Border, Kind::Unspecified);
         Self {
             inner,
             params: Parameters {
@@ -89,8 +87,8 @@ impl BorderRenderElement {
         }
     }
 
-    pub fn update_shader(&mut self, shader: Option<&ShaderProgram>) {
-        self.inner.update_shader(shader);
+    pub fn damage_all(&mut self) {
+        self.inner.damage_all();
     }
 
     #[allow(clippy::too_many_arguments)]
@@ -187,12 +185,10 @@ impl BorderRenderElement {
         );
     }
 
-    pub fn has_shader(&self) -> bool {
-        self.inner.has_shader()
-    }
-
-    pub fn shader(renderer: &mut impl NiriRenderer) -> Option<&ShaderProgram> {
-        Shaders::get(renderer).border.as_ref()
+    pub fn has_shader(renderer: &mut impl NiriRenderer) -> bool {
+        Shaders::get(renderer)
+            .program(ProgramType::Border)
+            .is_some()
     }
 }
 
diff --git a/src/render_helpers/resize.rs b/src/render_helpers/resize.rs
index c728689a..021c59bc 100644
--- a/src/render_helpers/resize.rs
+++ b/src/render_helpers/resize.rs
@@ -8,8 +8,8 @@ use smithay::backend::renderer::utils::{CommitCounter, DamageSet};
 use smithay::utils::{Buffer, Logical, Physical, Rectangle, Scale, Size, Transform};
 
 use super::renderer::{AsGlesFrame, NiriRenderer};
-use super::shader_element::{ShaderProgram, ShaderRenderElement};
-use super::shaders::{mat3_uniform, Shaders};
+use super::shader_element::ShaderRenderElement;
+use super::shaders::{mat3_uniform, ProgramType, Shaders};
 use crate::backend::tty::{TtyFrame, TtyRenderer, TtyRendererError};
 
 #[derive(Debug)]
@@ -18,7 +18,6 @@ pub struct ResizeRenderElement(ShaderRenderElement);
 impl ResizeRenderElement {
     #[allow(clippy::too_many_arguments)]
     pub fn new(
-        shader: ShaderProgram,
         area: Rectangle<i32, Logical>,
         scale: Scale<f64>,
         texture_prev: (GlesTexture, Rectangle<i32, Physical>),
@@ -92,7 +91,7 @@ impl ResizeRenderElement {
 
         // Create the shader.
         Self(ShaderRenderElement::new(
-            Some(shader),
+            ProgramType::Resize,
             area,
             size,
             None,
@@ -117,8 +116,10 @@ impl ResizeRenderElement {
         ))
     }
 
-    pub fn shader(renderer: &mut impl NiriRenderer) -> Option<ShaderProgram> {
-        Shaders::get(renderer).resize()
+    pub fn has_shader(renderer: &mut impl NiriRenderer) -> bool {
+        Shaders::get(renderer)
+            .program(ProgramType::Resize)
+            .is_some()
     }
 }
 
diff --git a/src/render_helpers/shader_element.rs b/src/render_helpers/shader_element.rs
index 75ed2780..4b4af625 100644
--- a/src/render_helpers/shader_element.rs
+++ b/src/render_helpers/shader_element.rs
@@ -14,12 +14,13 @@ use smithay::utils::{Buffer, Logical, Physical, Rectangle, Scale, Size};
 
 use super::renderer::AsGlesFrame;
 use super::resources::Resources;
+use super::shaders::{ProgramType, Shaders};
 use crate::backend::tty::{TtyFrame, TtyRenderer, TtyRendererError};
 
 /// Renders a shader with optional texture input, on the primary GPU.
 #[derive(Debug, Clone)]
 pub struct ShaderRenderElement {
-    shader: Option<ShaderProgram>,
+    program: ProgramType,
     id: Id,
     commit_counter: CommitCounter,
     area: Rectangle<i32, Logical>,
@@ -197,7 +198,7 @@ impl ShaderProgram {
 impl ShaderRenderElement {
     #[allow(clippy::too_many_arguments)]
     pub fn new(
-        shader: Option<ShaderProgram>,
+        program: ProgramType,
         area: Rectangle<i32, Logical>,
         size: Size<f64, Buffer>,
         opaque_regions: Option<Vec<Rectangle<i32, Logical>>>,
@@ -207,7 +208,7 @@ impl ShaderRenderElement {
         kind: Kind,
     ) -> Self {
         Self {
-            shader,
+            program,
             id: Id::new(),
             commit_counter: CommitCounter::default(),
             area,
@@ -220,9 +221,9 @@ impl ShaderRenderElement {
         }
     }
 
-    pub fn empty(kind: Kind) -> Self {
+    pub fn empty(program: ProgramType, kind: Kind) -> Self {
         Self {
-            shader: None,
+            program,
             id: Id::new(),
             commit_counter: CommitCounter::default(),
             area: Rectangle::default(),
@@ -235,12 +236,7 @@ impl ShaderRenderElement {
         }
     }
 
-    pub fn update_shader(&mut self, shader: Option<&ShaderProgram>) {
-        if self.shader.as_ref() == shader {
-            return;
-        }
-
-        self.shader = shader.cloned();
+    pub fn damage_all(&mut self) {
         self.commit_counter.increment();
     }
 
@@ -260,10 +256,6 @@ impl ShaderRenderElement {
 
         self.commit_counter.increment();
     }
-
-    pub fn has_shader(&self) -> bool {
-        self.shader.is_some()
-    }
 }
 
 impl Element for ShaderRenderElement {
@@ -309,7 +301,7 @@ impl RenderElement<GlesRenderer> for ShaderRenderElement {
     ) -> Result<(), GlesError> {
         let frame = frame.as_gles_frame();
 
-        let Some(shader) = &self.shader else {
+        let Some(shader) = Shaders::get_from_frame(frame).program(self.program) else {
             return Ok(());
         };
 
@@ -389,14 +381,14 @@ impl RenderElement<GlesRenderer> for ShaderRenderElement {
         let has_debug = !frame.debug_flags().is_empty();
         let has_tint = frame.debug_flags().contains(DebugFlags::TINT);
 
-        let program = if has_debug {
-            &shader.0.debug
-        } else {
-            &shader.0.normal
-        };
-
         // render
         frame.with_context(move |gl| -> Result<(), GlesError> {
+            let program = if has_debug {
+                &shader.0.debug
+            } else {
+                &shader.0.normal
+            };
+
             unsafe {
                 for (i, texture) in self.textures.values().enumerate() {
                     gl.ActiveTexture(ffi::TEXTURE0 + i as u32);
diff --git a/src/render_helpers/shaders/mod.rs b/src/render_helpers/shaders/mod.rs
index 5937fcc1..034d8db1 100644
--- a/src/render_helpers/shaders/mod.rs
+++ b/src/render_helpers/shaders/mod.rs
@@ -2,7 +2,8 @@ use std::cell::RefCell;
 
 use glam::Mat3;
 use smithay::backend::renderer::gles::{
-    GlesError, GlesRenderer, GlesTexProgram, Uniform, UniformName, UniformType, UniformValue,
+    GlesError, GlesFrame, GlesRenderer, GlesTexProgram, Uniform, UniformName, UniformType,
+    UniformValue,
 };
 
 use super::renderer::NiriRenderer;
@@ -15,6 +16,12 @@ pub struct Shaders {
     pub custom_resize: RefCell<Option<ShaderProgram>>,
 }
 
+#[derive(Debug, Clone, Copy)]
+pub enum ProgramType {
+    Border,
+    Resize,
+}
+
 impl Shaders {
     fn compile(renderer: &mut GlesRenderer) -> Self {
         let _span = tracy_client::span!("Shaders::compile");
@@ -68,6 +75,12 @@ impl Shaders {
         }
     }
 
+    pub fn get_from_frame<'a>(frame: &'a mut GlesFrame<'_>) -> &'a Self {
+        let data = frame.egl_context().user_data();
+        data.get()
+            .expect("shaders::init() must be called when creating the renderer")
+    }
+
     pub fn get(renderer: &mut impl NiriRenderer) -> &Self {
         let renderer = renderer.as_gles_renderer();
         let data = renderer.egl_context().user_data();
@@ -82,11 +95,15 @@ impl Shaders {
         self.custom_resize.replace(program)
     }
 
-    pub fn resize(&self) -> Option<ShaderProgram> {
-        self.custom_resize
-            .borrow()
-            .clone()
-            .or_else(|| self.resize.clone())
+    pub fn program(&self, program: ProgramType) -> Option<ShaderProgram> {
+        match program {
+            ProgramType::Border => self.border.clone(),
+            ProgramType::Resize => self
+                .custom_resize
+                .borrow()
+                .clone()
+                .or_else(|| self.resize.clone()),
+        }
     }
 }