Allow read access to shapes added to painter this frame (#3866)

Co-authored-by: Emil Ernerfeldt <emil.ernerfeldt@gmail.com>

crates/egui/src/context.rs

@@ -701,6 +701,12 @@ impl Context {
         self.write(move |ctx| writer(&mut ctx.viewport().graphics))
     }
 
+    /// Read-only access to [`GraphicLayers`], where painted [`crate::Shape`]s are written to.
+    #[inline]
+    pub(crate) fn graphics<R>(&self, reader: impl FnOnce(&GraphicLayers) -> R) -> R {
+        self.write(move |ctx| reader(&ctx.viewport().graphics))
+    }
+
     /// Read-only access to [`PlatformOutput`].
     ///
     /// This is what egui outputs each frame.

crates/egui/src/layers.rs

@@ -111,7 +111,7 @@ impl LayerId {
 
 /// A unique identifier of a specific [`Shape`] in a [`PaintList`].
 #[derive(Clone, Copy, PartialEq, Eq)]
-pub struct ShapeIdx(usize);
+pub struct ShapeIdx(pub usize);
 
 /// A list of [`Shape`]s paired with a clip rectangle.
 #[derive(Clone, Default)]
@@ -158,6 +158,11 @@ impl PaintList {
             shape.translate(delta);
         }
     }
+
+    /// Read-only access to all held shapes.
+    pub fn all_entries(&self) -> impl ExactSizeIterator<Item = &ClippedShape> {
+        self.0.iter()
+    }
 }
 
 #[derive(Clone, Default)]
@@ -170,6 +175,10 @@ impl GraphicLayers {
             .or_default()
     }
 
+    pub fn get(&self, layer_id: LayerId) -> Option<&PaintList> {
+        self.0[layer_id.order as usize].get(&layer_id.id)
+    }
+
     pub fn drain(&mut self, area_order: &[LayerId]) -> Vec<ClippedShape> {
         crate::profile_function!();
 

crates/egui/src/painter.rs

@@ -7,7 +7,7 @@ use crate::{
 };
 use epaint::{
     text::{Fonts, Galley},
-    CircleShape, RectShape, Rounding, Shape, Stroke,
+    CircleShape, ClippedShape, RectShape, Rounding, Shape, Stroke,
 };
 
 /// Helper to paint shapes and text to a specific region on a specific layer.
@@ -193,6 +193,17 @@ impl Painter {
         self.transform_shape(&mut shape);
         self.paint_list(|l| l.set(idx, self.clip_rect, shape));
     }
+
+    /// Access all shapes added this frame.
+    pub fn for_each_shape(&self, mut reader: impl FnMut(&ClippedShape)) {
+        self.ctx.graphics(|g| {
+            if let Some(list) = g.get(self.layer_id) {
+                for c in list.all_entries() {
+                    reader(c);
+                }
+            }
+        });
+    }
 }
 
 /// ## Debug painting