Implement custom shader for window-open

9 files changed, 296 insertions, 42 deletions

diff --git a/src/backend/tty.rs b/src/backend/tty.rs
index 65ecd458..05164041 100644
--- a/src/backend/tty.rs
+++ b/src/backend/tty.rs
@@ -508,6 +508,9 @@ impl Tty {
             if let Some(src) = config.animations.window_close.custom_shader.as_deref() {
                 shaders::set_custom_close_program(gles_renderer, Some(src));
             }
+            if let Some(src) = config.animations.window_open.custom_shader.as_deref() {
+                shaders::set_custom_open_program(gles_renderer, Some(src));
+            }
             drop(config);
 
             niri.layout.update_shaders();
diff --git a/src/backend/winit.rs b/src/backend/winit.rs
index 047ebb06..432904e2 100644
--- a/src/backend/winit.rs
+++ b/src/backend/winit.rs
@@ -143,6 +143,9 @@ impl Winit {
         if let Some(src) = config.animations.window_close.custom_shader.as_deref() {
             shaders::set_custom_close_program(renderer, Some(src));
         }
+        if let Some(src) = config.animations.window_open.custom_shader.as_deref() {
+            shaders::set_custom_open_program(renderer, Some(src));
+        }
         drop(config);
 
         niri.layout.update_shaders();
diff --git a/src/layout/mod.rs b/src/layout/mod.rs
index ec02bfc6..f4b88560 100644
--- a/src/layout/mod.rs
+++ b/src/layout/mod.rs
@@ -58,6 +58,7 @@ use crate::window::ResolvedWindowRules;
 pub mod closing_window;
 pub mod focus_ring;
 pub mod monitor;
+pub mod opening_window;
 pub mod tile;
 pub mod workspace;
 
diff --git a/src/layout/opening_window.rs b/src/layout/opening_window.rs
new file mode 100644
index 00000000..1a8b3b99
--- /dev/null
+++ b/src/layout/opening_window.rs
@@ -0,0 +1,156 @@
+use std::collections::HashMap;
+use std::time::Duration;
+
+use anyhow::Context as _;
+use glam::{Mat3, Vec2};
+use smithay::backend::allocator::Fourcc;
+use smithay::backend::renderer::element::texture::TextureRenderElement;
+use smithay::backend::renderer::element::utils::{
+    Relocate, RelocateRenderElement, RescaleRenderElement,
+};
+use smithay::backend::renderer::element::{Id, Kind, RenderElement};
+use smithay::backend::renderer::gles::{GlesRenderer, Uniform};
+use smithay::backend::renderer::{Renderer as _, Texture};
+use smithay::utils::{Logical, Point, Rectangle, Scale, Size, Transform};
+
+use crate::animation::Animation;
+use crate::niri_render_elements;
+use crate::render_helpers::primary_gpu_texture::PrimaryGpuTextureRenderElement;
+use crate::render_helpers::render_to_encompassing_texture;
+use crate::render_helpers::shader_element::ShaderRenderElement;
+use crate::render_helpers::shaders::{mat3_uniform, ProgramType, Shaders};
+
+#[derive(Debug)]
+pub struct OpenAnimation {
+    anim: Animation,
+    random_seed: f32,
+}
+
+niri_render_elements! {
+    OpeningWindowRenderElement => {
+        Texture = RelocateRenderElement<RescaleRenderElement<PrimaryGpuTextureRenderElement>>,
+        Shader = ShaderRenderElement,
+    }
+}
+
+impl OpenAnimation {
+    pub fn new(anim: Animation) -> Self {
+        Self {
+            anim,
+            random_seed: fastrand::f32(),
+        }
+    }
+
+    pub fn advance_animations(&mut self, current_time: Duration) {
+        self.anim.set_current_time(current_time);
+    }
+
+    pub fn is_done(&self) -> bool {
+        self.anim.is_done()
+    }
+
+    // We can't depend on view_rect here, because the result of window opening can be snapshot and
+    // then rendered elsewhere.
+    pub fn render(
+        &self,
+        renderer: &mut GlesRenderer,
+        elements: &[impl RenderElement<GlesRenderer>],
+        geo_size: Size<i32, Logical>,
+        location: Point<i32, Logical>,
+        scale: Scale<f64>,
+    ) -> anyhow::Result<OpeningWindowRenderElement> {
+        let progress = self.anim.value();
+        let clamped_progress = self.anim.clamped_value().clamp(0., 1.);
+
+        let (texture, _sync_point, geo) = render_to_encompassing_texture(
+            renderer,
+            scale,
+            Transform::Normal,
+            Fourcc::Abgr8888,
+            elements,
+        )
+        .context("error rendering to texture")?;
+
+        let offset = geo.loc.to_f64().to_logical(scale);
+        let texture_size = geo.size.to_f64().to_logical(scale);
+
+        if Shaders::get(renderer).program(ProgramType::Open).is_some() {
+            let mut area = Rectangle::from_loc_and_size(location.to_f64() + offset, texture_size);
+
+            // Expand the area a bit to allow for more varied effects.
+            let mut target_size = area.size.upscale(1.5);
+            target_size.w = f64::max(area.size.w + 1000., target_size.w);
+            target_size.h = f64::max(area.size.h + 1000., target_size.h);
+            let diff = target_size.to_point() - area.size.to_point();
+            area.loc -= diff.downscale(2.);
+            area.size += diff.to_size();
+
+            let area = area.to_i32_up();
+            let area_loc = Vec2::new(area.loc.x as f32, area.loc.y as f32);
+            let area_size = Vec2::new(area.size.w as f32, area.size.h as f32);
+
+            let geo_loc = Vec2::new(location.x as f32, location.y as f32);
+            let geo_size = Vec2::new(geo_size.w as f32, geo_size.h as f32);
+
+            let input_to_geo = Mat3::from_scale(area_size / geo_size)
+                * Mat3::from_translation((area_loc - geo_loc) / area_size);
+
+            let tex_scale = Vec2::new(scale.x as f32, scale.y as f32);
+            let tex_loc = Vec2::new(offset.x as f32, offset.y as f32);
+            let tex_size = Vec2::new(texture.width() as f32, texture.height() as f32) / tex_scale;
+
+            let geo_to_tex =
+                Mat3::from_translation(-tex_loc / tex_size) * Mat3::from_scale(geo_size / tex_size);
+
+            return Ok(ShaderRenderElement::new(
+                ProgramType::Open,
+                area.size,
+                None,
+                1.,
+                vec![
+                    mat3_uniform("niri_input_to_geo", input_to_geo),
+                    Uniform::new("niri_geo_size", geo_size.to_array()),
+                    mat3_uniform("niri_geo_to_tex", geo_to_tex),
+                    Uniform::new("niri_progress", progress as f32),
+                    Uniform::new("niri_clamped_progress", clamped_progress as f32),
+                    Uniform::new("niri_random_seed", self.random_seed),
+                ],
+                HashMap::from([(String::from("niri_tex"), texture.clone())]),
+                Kind::Unspecified,
+            )
+            .with_location(area.loc)
+            .into());
+        }
+
+        let elem = TextureRenderElement::from_static_texture(
+            Id::new(),
+            renderer.id(),
+            Point::from((0., 0.)),
+            texture.clone(),
+            scale.x as i32,
+            Transform::Normal,
+            Some(clamped_progress as f32),
+            None,
+            None,
+            None,
+            Kind::Unspecified,
+        );
+
+        let elem = PrimaryGpuTextureRenderElement(elem);
+
+        let center = geo_size.to_point().to_f64().downscale(2.);
+        let elem = RescaleRenderElement::from_element(
+            elem,
+            (center - offset).to_physical_precise_round(scale),
+            (progress / 2. + 0.5).max(0.),
+        );
+
+        let elem = RelocateRenderElement::from_element(
+            elem,
+            (location.to_f64() + offset).to_physical_precise_round(scale),
+            Relocate::Relative,
+        );
+
+        Ok(elem.into())
+    }
+}
diff --git a/src/layout/tile.rs b/src/layout/tile.rs
index d53eb2b1..719f12bc 100644
--- a/src/layout/tile.rs
+++ b/src/layout/tile.rs
@@ -5,12 +5,12 @@ use std::time::Duration;
 use niri_config::CornerRadius;
 use smithay::backend::allocator::Fourcc;
 use smithay::backend::renderer::element::solid::{SolidColorBuffer, SolidColorRenderElement};
-use smithay::backend::renderer::element::utils::RescaleRenderElement;
 use smithay::backend::renderer::element::{Element, Kind};
 use smithay::backend::renderer::gles::GlesRenderer;
 use smithay::utils::{Logical, Point, Rectangle, Scale, Size, Transform};
 
 use super::focus_ring::{FocusRing, FocusRingRenderElement};
+use super::opening_window::{OpenAnimation, OpeningWindowRenderElement};
 use super::{
     LayoutElement, LayoutElementRenderElement, LayoutElementRenderSnapshot, Options,
     RESIZE_ANIMATION_THRESHOLD,
@@ -20,7 +20,6 @@ use crate::niri_render_elements;
 use crate::render_helpers::border::BorderRenderElement;
 use crate::render_helpers::clipped_surface::{ClippedSurfaceRenderElement, RoundedCornerDamage};
 use crate::render_helpers::damage::ExtraDamage;
-use crate::render_helpers::offscreen::OffscreenRenderElement;
 use crate::render_helpers::renderer::NiriRenderer;
 use crate::render_helpers::resize::ResizeRenderElement;
 use crate::render_helpers::snapshot::RenderSnapshot;
@@ -54,7 +53,7 @@ pub struct Tile<W: LayoutElement> {
     fullscreen_size: Size<i32, Logical>,
 
     /// The animation upon opening a window.
-    open_animation: Option<Animation>,
+    open_animation: Option<OpenAnimation>,
 
     /// The animation of the window resizing.
     resize_animation: Option<ResizeAnimation>,
@@ -80,7 +79,7 @@ niri_render_elements! {
         LayoutElement = LayoutElementRenderElement<R>,
         FocusRing = FocusRingRenderElement,
         SolidColor = SolidColorRenderElement,
-        Offscreen = RescaleRenderElement<OffscreenRenderElement>,
+        Opening = OpeningWindowRenderElement,
         Resize = ResizeRenderElement,
         Border = BorderRenderElement,
         ClippedSurface = ClippedSurfaceRenderElement<R>,
@@ -201,9 +200,9 @@ impl<W: LayoutElement> Tile<W> {
     }
 
     pub fn advance_animations(&mut self, current_time: Duration) {
-        if let Some(anim) = &mut self.open_animation {
-            anim.set_current_time(current_time);
-            if anim.is_done() {
+        if let Some(open) = &mut self.open_animation {
+            open.advance_animations(current_time);
+            if open.is_done() {
                 self.open_animation = None;
             }
         }
@@ -299,16 +298,12 @@ impl<W: LayoutElement> Tile<W> {
     }
 
     pub fn start_open_animation(&mut self) {
-        self.open_animation = Some(Animation::new(
+        self.open_animation = Some(OpenAnimation::new(Animation::new(
             0.,
             1.,
             0.,
-            self.options.animations.window_open.0,
-        ));
-    }
-
-    pub fn open_animation(&self) -> &Option<Animation> {
-        &self.open_animation
+            self.options.animations.window_open.anim,
+        )));
     }
 
     pub fn resize_animation(&self) -> Option<&Animation> {
@@ -797,39 +792,34 @@ impl<W: LayoutElement> Tile<W> {
     ) -> impl Iterator<Item = TileRenderElement<R>> {
         let _span = tracy_client::span!("Tile::render");
 
-        if let Some(anim) = &self.open_animation {
+        let mut open_anim_elem = None;
+        let mut window_elems = None;
+
+        if let Some(open) = &self.open_animation {
             let renderer = renderer.as_gles_renderer();
-            let elements = self.render_inner(renderer, location, scale, focus_ring, target);
+            let elements =
+                self.render_inner(renderer, Point::from((0, 0)), scale, focus_ring, target);
             let elements = elements.collect::<Vec<TileRenderElement<_>>>();
+            match open.render(renderer, &elements, self.tile_size(), location, scale) {
+                Ok(elem) => {
+                    self.window()
+                        .set_offscreen_element_id(Some(elem.id().clone()));
+                    open_anim_elem = Some(elem.into());
+                }
+                Err(err) => {
+                    warn!("error rendering window opening animation: {err:?}");
+                }
+            }
+        }
 
-            let elem = OffscreenRenderElement::new(
-                renderer,
-                scale.x as i32,
-                &elements,
-                anim.clamped_value().clamp(0., 1.) as f32,
-            );
-            self.window()
-                .set_offscreen_element_id(Some(elem.id().clone()));
-
-            let mut center = location;
-            center.x += self.tile_size().w / 2;
-            center.y += self.tile_size().h / 2;
-
-            Some(TileRenderElement::Offscreen(
-                RescaleRenderElement::from_element(
-                    elem,
-                    center.to_physical_precise_round(scale),
-                    (anim.value() / 2. + 0.5).max(0.),
-                ),
-            ))
-            .into_iter()
-            .chain(None.into_iter().flatten())
-        } else {
+        if open_anim_elem.is_none() {
             self.window().set_offscreen_element_id(None);
-
-            let elements = self.render_inner(renderer, location, scale, focus_ring, target);
-            None.into_iter().chain(Some(elements).into_iter().flatten())
+            window_elems = Some(self.render_inner(renderer, location, scale, focus_ring, target));
         }
+
+        open_anim_elem
+            .into_iter()
+            .chain(window_elems.into_iter().flatten())
     }
 
     pub fn store_unmap_snapshot_if_empty(
diff --git a/src/niri.rs b/src/niri.rs
index 7fc94365..4e887570 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -951,6 +951,16 @@ impl State {
             shaders_changed = true;
         }
 
+        if config.animations.window_open.custom_shader
+            != old_config.animations.window_open.custom_shader
+        {
+            let src = config.animations.window_open.custom_shader.as_deref();
+            self.backend.with_primary_renderer(|renderer| {
+                shaders::set_custom_open_program(renderer, src);
+            });
+            shaders_changed = true;
+        }
+
         if config.debug != old_config.debug {
             debug_config_changed = true;
         }
diff --git a/src/render_helpers/shaders/mod.rs b/src/render_helpers/shaders/mod.rs
index 21287a55..12e9d4ed 100644
--- a/src/render_helpers/shaders/mod.rs
+++ b/src/render_helpers/shaders/mod.rs
@@ -15,6 +15,7 @@ pub struct Shaders {
     pub resize: Option<ShaderProgram>,
     pub custom_resize: RefCell<Option<ShaderProgram>>,
     pub custom_close: RefCell<Option<ShaderProgram>>,
+    pub custom_open: RefCell<Option<ShaderProgram>>,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -22,6 +23,7 @@ pub enum ProgramType {
     Border,
     Resize,
     Close,
+    Open,
 }
 
 impl Shaders {
@@ -75,6 +77,7 @@ impl Shaders {
             resize,
             custom_resize: RefCell::new(None),
             custom_close: RefCell::new(None),
+            custom_open: RefCell::new(None),
         }
     }
 
@@ -105,6 +108,13 @@ impl Shaders {
         self.custom_close.replace(program)
     }
 
+    pub fn replace_custom_open_program(
+        &self,
+        program: Option<ShaderProgram>,
+    ) -> Option<ShaderProgram> {
+        self.custom_open.replace(program)
+    }
+
     pub fn program(&self, program: ProgramType) -> Option<ShaderProgram> {
         match program {
             ProgramType::Border => self.border.clone(),
@@ -114,6 +124,7 @@ impl Shaders {
                 .clone()
                 .or_else(|| self.resize.clone()),
             ProgramType::Close => self.custom_close.borrow().clone(),
+            ProgramType::Open => self.custom_open.borrow().clone(),
         }
     }
 }
@@ -216,6 +227,49 @@ pub fn set_custom_close_program(renderer: &mut GlesRenderer, src: Option<&str>)
     }
 }
 
+fn compile_open_program(
+    renderer: &mut GlesRenderer,
+    src: &str,
+) -> Result<ShaderProgram, GlesError> {
+    let mut program = include_str!("open_prelude.frag").to_string();
+    program.push_str(src);
+    program.push_str(include_str!("open_epilogue.frag"));
+
+    ShaderProgram::compile(
+        renderer,
+        &program,
+        &[
+            UniformName::new("niri_input_to_geo", UniformType::Matrix3x3),
+            UniformName::new("niri_geo_size", UniformType::_2f),
+            UniformName::new("niri_geo_to_tex", UniformType::Matrix3x3),
+            UniformName::new("niri_progress", UniformType::_1f),
+            UniformName::new("niri_clamped_progress", UniformType::_1f),
+            UniformName::new("niri_random_seed", UniformType::_1f),
+        ],
+        &["niri_tex"],
+    )
+}
+
+pub fn set_custom_open_program(renderer: &mut GlesRenderer, src: Option<&str>) {
+    let program = if let Some(src) = src {
+        match compile_open_program(renderer, src) {
+            Ok(program) => Some(program),
+            Err(err) => {
+                warn!("error compiling custom open shader: {err:?}");
+                return;
+            }
+        }
+    } else {
+        None
+    };
+
+    if let Some(prev) = Shaders::get(renderer).replace_custom_open_program(program) {
+        if let Err(err) = prev.destroy(renderer) {
+            warn!("error destroying previous custom open shader: {err:?}");
+        }
+    }
+}
+
 pub fn mat3_uniform(name: &str, mat: Mat3) -> Uniform {
     Uniform::new(
         name,
diff --git a/src/render_helpers/shaders/open_epilogue.frag b/src/render_helpers/shaders/open_epilogue.frag
new file mode 100644
index 00000000..1e5f5375
--- /dev/null
+++ b/src/render_helpers/shaders/open_epilogue.frag
@@ -0,0 +1,16 @@
+
+void main() {
+    vec3 coords_geo = niri_input_to_geo * vec3(niri_v_coords, 1.0);
+    vec3 size_geo = vec3(niri_geo_size, 1.0);
+
+    vec4 color = open_color(coords_geo, size_geo);
+
+    color = color * niri_alpha;
+
+#if defined(DEBUG_FLAGS)
+    if (niri_tint == 1.0)
+        color = vec4(0.0, 0.2, 0.0, 0.2) + color * 0.8;
+#endif
+
+    gl_FragColor = color;
+}
diff --git a/src/render_helpers/shaders/open_prelude.frag b/src/render_helpers/shaders/open_prelude.frag
new file mode 100644
index 00000000..db42190b
--- /dev/null
+++ b/src/render_helpers/shaders/open_prelude.frag
@@ -0,0 +1,21 @@
+precision mediump float;
+
+#if defined(DEBUG_FLAGS)
+uniform float niri_tint;
+#endif
+
+varying vec2 niri_v_coords;
+uniform vec2 niri_size;
+
+uniform mat3 niri_input_to_geo;
+uniform vec2 niri_geo_size;
+
+uniform sampler2D niri_tex;
+uniform mat3 niri_geo_to_tex;
+
+uniform float niri_progress;
+uniform float niri_clamped_progress;
+uniform float niri_random_seed;
+
+uniform float niri_alpha;
+