From b452b2f6347e13c00c817b769cb5bf145af19c22 Mon Sep 17 00:00:00 2001
From: Ottatop <ottatop1227@gmail.com>
Date: Sat, 2 Sep 2023 17:49:57 -0500
Subject: [PATCH] Render layer surfaces properly

---
 src/input.rs  |  5 ++++
 src/render.rs | 74 ++++++++++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/src/input.rs b/src/input.rs
index dd6b8ef..d7b2482 100644
--- a/src/input.rs
+++ b/src/input.rs
@@ -102,9 +102,12 @@ impl State {
         });
 
         // I think I'm going a bit too far with the functional stuff
+
+        // The topmost fullscreen window
         top_fullscreen_window
             .map(|window| (FocusTarget::from(window.clone()), output_geo.loc))
             .or_else(|| {
+                // The topmost layer surface in Overlay or Top
                 layers
                     .layer_under(wlr_layer::Layer::Overlay, point)
                     .or_else(|| layers.layer_under(wlr_layer::Layer::Top, point))
@@ -114,6 +117,7 @@ impl State {
                     })
             })
             .or_else(|| {
+                // The topmost window
                 self.space
                     .elements()
                     .rev()
@@ -133,6 +137,7 @@ impl State {
                     })
             })
             .or_else(|| {
+                // The topmost layer surface in Bottom or Background
                 layers
                     .layer_under(wlr_layer::Layer::Bottom, point)
                     .or_else(|| layers.layer_under(wlr_layer::Layer::Background, point))
diff --git a/src/render.rs b/src/render.rs
index 3be0410..46cabe3 100644
--- a/src/render.rs
+++ b/src/render.rs
@@ -11,6 +11,7 @@ use smithay::{
         ImportAll, ImportMem, Renderer, Texture,
     },
     desktop::{
+        layer_map_for_output,
         space::{SpaceRenderElements, SurfaceTree},
         Space,
     },
@@ -22,7 +23,7 @@ use smithay::{
     },
     render_elements,
     utils::{IsAlive, Logical, Physical, Point, Scale},
-    wayland::{compositor, input_method::InputMethodHandle},
+    wayland::{compositor, input_method::InputMethodHandle, shell::wlr_layer},
 };
 
 use crate::{state::WithState, tag::Tag, window::WindowElement};
@@ -76,6 +77,58 @@ where
     }
 }
 
+struct LayerRenderElements<R> {
+    background: Vec<WaylandSurfaceRenderElement<R>>,
+    bottom: Vec<WaylandSurfaceRenderElement<R>>,
+    top: Vec<WaylandSurfaceRenderElement<R>>,
+    overlay: Vec<WaylandSurfaceRenderElement<R>>,
+}
+
+fn layer_render_elements<R>(output: &Output, renderer: &mut R) -> LayerRenderElements<R>
+where
+    R: Renderer + ImportAll,
+    <R as Renderer>::TextureId: 'static,
+{
+    let layer_map = layer_map_for_output(output);
+    let mut overlay = vec![];
+    let mut top = vec![];
+    let mut bottom = vec![];
+    let mut background = vec![];
+
+    let layer_elements = layer_map
+        .layers()
+        .filter_map(|surface| {
+            layer_map
+                .layer_geometry(surface)
+                .map(|geo| (surface, geo.loc))
+        })
+        .map(|(surface, loc)| {
+            let render_elements = surface.render_elements::<WaylandSurfaceRenderElement<R>>(
+                renderer,
+                loc.to_physical(1),
+                Scale::from(1.0),
+                1.0,
+            );
+            (surface.layer(), render_elements)
+        });
+
+    for (layer, elements) in layer_elements {
+        match layer {
+            wlr_layer::Layer::Background => background.extend(elements),
+            wlr_layer::Layer::Bottom => bottom.extend(elements),
+            wlr_layer::Layer::Top => top.extend(elements),
+            wlr_layer::Layer::Overlay => overlay.extend(elements),
+        }
+    }
+
+    LayerRenderElements {
+        background,
+        bottom,
+        top,
+        overlay,
+    }
+}
+
 #[allow(clippy::too_many_arguments)]
 pub fn generate_render_elements<R, T>(
     space: &Space<WindowElement>,
@@ -216,29 +269,44 @@ where
             //     space::space_render_elements(renderer, [space], output, 1.0)
             //         .expect("Failed to get render elements");
 
+            let LayerRenderElements {
+                background,
+                bottom,
+                top,
+                overlay,
+            } = layer_render_elements(output, renderer);
+
             let tags = space
                 .outputs()
                 .flat_map(|op| {
                     op.with_state(|state| state.focused_tags().cloned().collect::<Vec<_>>())
                 })
                 .collect::<Vec<_>>();
-            let space_render_elements: Vec<WaylandSurfaceRenderElement<R>> =
+            let window_render_elements: Vec<WaylandSurfaceRenderElement<R>> =
                 Tag::tag_render_elements(&tags, windows, space, renderer);
 
             let mut output_render_elements =
                 Vec::<OutputRenderElements<R, WaylandSurfaceRenderElement<R>>>::new();
 
+            // Elements render from top to bottom
+
             output_render_elements.extend(
                 custom_render_elements
                     .into_iter()
                     .map(OutputRenderElements::from),
             );
+
             output_render_elements.extend(
-                space_render_elements
+                overlay
                     .into_iter()
+                    .chain(top)
+                    .chain(window_render_elements)
+                    .chain(bottom)
+                    .chain(background)
                     .map(CustomRenderElements::from)
                     .map(OutputRenderElements::from),
             );
+
             output_render_elements
         }
     };