2 files changed, 112 insertions, 24 deletions

diff --git a/src/render_helpers/mod.rs b/src/render_helpers/mod.rs
index ad3e91b8..de40c274 100644
--- a/src/render_helpers/mod.rs
+++ b/src/render_helpers/mod.rs
@@ -1,7 +1,6 @@
 use std::ptr;
 
 use anyhow::{ensure, Context};
-use niri_config::BlockOutFrom;
 use smithay::backend::allocator::Fourcc;
 use smithay::backend::renderer::element::solid::{SolidColorBuffer, SolidColorRenderElement};
 use smithay::backend::renderer::element::texture::{TextureBuffer, TextureRenderElement};
@@ -27,6 +26,7 @@ pub mod render_elements;
 pub mod renderer;
 pub mod resources;
 pub mod shaders;
+pub mod snapshot;
 pub mod surface;
 
 /// What we're rendering for.
@@ -49,19 +49,6 @@ pub struct BakedBuffer<B> {
     pub dst: Option<Size<i32, Logical>>,
 }
 
-/// Snapshot of a render.
-#[derive(Debug)]
-pub struct RenderSnapshot<C, B> {
-    /// Contents for a normal render.
-    pub contents: Vec<C>,
-
-    /// Blocked-out contents.
-    pub blocked_out_contents: Vec<B>,
-
-    /// Where the contents were blocked out from at the time of the snapshot.
-    pub block_out_from: Option<BlockOutFrom>,
-}
-
 pub trait ToRenderElement {
     type RenderElement;
 
@@ -118,16 +105,6 @@ impl ToRenderElement for BakedBuffer<SolidColorBuffer> {
     }
 }
 
-impl<C, B> Default for RenderSnapshot<C, B> {
-    fn default() -> Self {
-        Self {
-            contents: Default::default(),
-            blocked_out_contents: Default::default(),
-            block_out_from: Default::default(),
-        }
-    }
-}
-
 pub fn render_to_encompassing_texture(
     renderer: &mut GlesRenderer,
     scale: Scale<f64>,
diff --git a/src/render_helpers/snapshot.rs b/src/render_helpers/snapshot.rs
new file mode 100644
index 00000000..b394dcd9
--- /dev/null
+++ b/src/render_helpers/snapshot.rs
@@ -0,0 +1,111 @@
+use std::cell::OnceCell;
+
+use niri_config::BlockOutFrom;
+use smithay::backend::allocator::Fourcc;
+use smithay::backend::renderer::element::{Kind, RenderElement};
+use smithay::backend::renderer::gles::{GlesRenderer, GlesTexture};
+use smithay::utils::{Logical, Physical, Point, Rectangle, Scale, Size, Transform};
+
+use super::{render_to_encompassing_texture, RenderTarget, ToRenderElement};
+
+/// Snapshot of a render.
+#[derive(Debug)]
+pub struct RenderSnapshot<C, B> {
+    /// Contents for a normal render.
+    ///
+    /// Relative to the geometry.
+    pub contents: Vec<C>,
+
+    /// Blocked-out contents.
+    ///
+    /// Relative to the geometry.
+    pub blocked_out_contents: Vec<B>,
+
+    /// Where the contents were blocked out from at the time of the snapshot.
+    pub block_out_from: Option<BlockOutFrom>,
+
+    /// Visual size of the element at the point of the snapshot.
+    pub size: Size<i32, Logical>,
+
+    /// Contents rendered into a texture (lazily).
+    pub texture: OnceCell<Option<(GlesTexture, Rectangle<i32, Physical>)>>,
+
+    /// Blocked-out contents rendered into a texture (lazily).
+    pub blocked_out_texture: OnceCell<Option<(GlesTexture, Rectangle<i32, Physical>)>>,
+}
+
+impl<C, B, EC, EB> RenderSnapshot<C, B>
+where
+    C: ToRenderElement<RenderElement = EC>,
+    B: ToRenderElement<RenderElement = EB>,
+    EC: RenderElement<GlesRenderer>,
+    EB: RenderElement<GlesRenderer>,
+{
+    pub fn texture(
+        &self,
+        renderer: &mut GlesRenderer,
+        scale: Scale<f64>,
+        target: RenderTarget,
+    ) -> Option<&(GlesTexture, Rectangle<i32, Physical>)> {
+        let block_out = match self.block_out_from {
+            None => false,
+            Some(BlockOutFrom::Screencast) => target == RenderTarget::Screencast,
+            Some(BlockOutFrom::ScreenCapture) => target != RenderTarget::Output,
+        };
+
+        if block_out {
+            self.blocked_out_texture.get_or_init(|| {
+                let _span = tracy_client::span!("RenderSnapshot::Texture");
+
+                let elements: Vec<_> = self
+                    .blocked_out_contents
+                    .iter()
+                    .map(|baked| {
+                        baked.to_render_element(Point::from((0, 0)), scale, 1., Kind::Unspecified)
+                    })
+                    .collect();
+
+                match render_to_encompassing_texture(
+                    renderer,
+                    scale,
+                    Transform::Normal,
+                    Fourcc::Abgr8888,
+                    &elements,
+                ) {
+                    Ok((texture, _sync_point, geo)) => Some((texture, geo)),
+                    Err(err) => {
+                        warn!("error rendering blocked-out contents to texture: {err:?}");
+                        None
+                    }
+                }
+            })
+        } else {
+            self.texture.get_or_init(|| {
+                let _span = tracy_client::span!("RenderSnapshot::Texture");
+
+                let elements: Vec<_> = self
+                    .contents
+                    .iter()
+                    .map(|baked| {
+                        baked.to_render_element(Point::from((0, 0)), scale, 1., Kind::Unspecified)
+                    })
+                    .collect();
+
+                match render_to_encompassing_texture(
+                    renderer,
+                    scale,
+                    Transform::Normal,
+                    Fourcc::Abgr8888,
+                    &elements,
+                ) {
+                    Ok((texture, _sync_point, geo)) => Some((texture, geo)),
+                    Err(err) => {
+                        warn!("error rendering contents to texture: {err:?}");
+                        None
+                    }
+                }
+            })
+        }
+        .as_ref()
+    }
+}