From 780f1f9d113c59685d6f1df567bc14c801523aff Mon Sep 17 00:00:00 2001
From: Mikolaj Wielgus <wielgusmikolaj@gmail.com>
Date: Fri, 20 Mar 2026 12:01:03 +0100
Subject: [PATCH] Add separate methods to find all layer joints, segments,
 polygons

---
 topola/src/layout.rs | 36 +++++++++++++++++++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/topola/src/layout.rs b/topola/src/layout.rs
index 34c26f4..d237b40 100644
--- a/topola/src/layout.rs
+++ b/topola/src/layout.rs
@@ -7,7 +7,7 @@ use std::collections::BTreeMap;
 use derive_getters::{Dissolve, Getters};
 use derive_more::Constructor;
 use rstar::{
-    RTree,
+    AABB, RTree,
     primitives::{GeomWithData, Rectangle},
 };
 use serde::{Deserialize, Serialize};
@@ -252,6 +252,40 @@ impl Layout {
             })
     }
 
+    pub fn layer_joints(&self, layer: usize) -> impl Iterator<Item = JointId> + '_ {
+        let envelope = Self::whole_layer_aabb(layer);
+        self.joints_rtree
+            .as_ref()
+            .locate_in_envelope_intersecting(&envelope)
+            .map(|geom_with_data| geom_with_data.data)
+            .filter(move |&id| self.joint(id).layer == layer)
+    }
+
+    pub fn layer_segments(&self, layer: usize) -> impl Iterator<Item = SegmentId> + '_ {
+        let envelope = Self::whole_layer_aabb(layer);
+        self.segments_rtree
+            .as_ref()
+            .locate_in_envelope_intersecting(&envelope)
+            .map(|geom_with_data| geom_with_data.data)
+            .filter(move |&id| self.segment(id).layer == layer)
+    }
+
+    pub fn layer_polygons(&self, layer: usize) -> impl Iterator<Item = PolygonId> + '_ {
+        let envelope = Self::whole_layer_aabb(layer);
+        self.polygons_rtree
+            .as_ref()
+            .locate_in_envelope_intersecting(&envelope)
+            .map(|geom_with_data| geom_with_data.data)
+            .filter(move |&id| self.polygon(id).layer == layer)
+    }
+
+    fn whole_layer_aabb(layer: usize) -> AABB<[i64; 3]> {
+        AABB::from_corners(
+            [i64::MIN, i64::MIN, layer as i64],
+            [i64::MAX, i64::MAX, layer as i64],
+        )
+    }
+
     pub fn joint(&self, joint_id: JointId) -> &Joint {
         self.joints.get(&joint_id.index()).unwrap()
     }