From 13dd000dba391cb5b2aa7dfe772f39d1a47aa8eb Mon Sep 17 00:00:00 2001
From: Mikolaj Wielgus <wielgusmikolaj@gmail.com>
Date: Fri, 12 Apr 2024 23:50:47 +0000
Subject: [PATCH] geometry,layout,drawing: add accessor for node's groupings

---
 src/drawing/drawing.rs     |  4 ++++
 src/geometry/geometry.rs   | 14 ++++++++++++++
 src/geometry/grouping.rs   |  1 +
 src/geometry/with_rtree.rs |  4 ++++
 src/layout/layout.rs       |  7 +++++++
 5 files changed, 30 insertions(+)

diff --git a/src/drawing/drawing.rs b/src/drawing/drawing.rs
index adc513e..a80507e 100644
--- a/src/drawing/drawing.rs
+++ b/src/drawing/drawing.rs
@@ -887,4 +887,8 @@ impl<GW: Copy, R: RulesTrait> GroupingManagerTrait<GW, GenericIndex<GW>> for Dra
         self.geometry_with_rtree
             .assign_to_grouping(primitive, grouping);
     }
+
+    fn groupings<W>(&self, node: GenericIndex<W>) -> impl Iterator<Item = GenericIndex<GW>> {
+        self.geometry_with_rtree.groupings(node)
+    }
 }
diff --git a/src/geometry/geometry.rs b/src/geometry/geometry.rs
index 878a191..321282d 100644
--- a/src/geometry/geometry.rs
+++ b/src/geometry/geometry.rs
@@ -500,4 +500,18 @@ impl<
             GeometryLabel::Grouping,
         );
     }
+
+    fn groupings<W>(&self, node: GenericIndex<W>) -> impl Iterator<Item = GenericIndex<GW>> {
+        self.graph
+            .neighbors(node.node_index())
+            .filter(move |ni| {
+                matches!(
+                    self.graph
+                        .edge_weight(self.graph.find_edge(node.node_index(), *ni).unwrap())
+                        .unwrap(),
+                    GeometryLabel::Grouping
+                )
+            })
+            .map(|ni| GenericIndex::new(ni))
+    }
 }
diff --git a/src/geometry/grouping.rs b/src/geometry/grouping.rs
index e779995..52680b9 100644
--- a/src/geometry/grouping.rs
+++ b/src/geometry/grouping.rs
@@ -4,4 +4,5 @@ pub trait GroupingManagerTrait<GW: Copy, GI: GetNodeIndex + Copy> {
     fn add_grouping(&mut self, weight: GW) -> GenericIndex<GW>;
     fn remove_grouping(&mut self, grouping: GenericIndex<GW>);
     fn assign_to_grouping<W>(&mut self, node: GenericIndex<W>, grouping: GenericIndex<GW>);
+    fn groupings<W>(&self, node: GenericIndex<W>) -> impl Iterator<Item = GenericIndex<GW>>;
 }
diff --git a/src/geometry/with_rtree.rs b/src/geometry/with_rtree.rs
index 39a810d..6779bcd 100644
--- a/src/geometry/with_rtree.rs
+++ b/src/geometry/with_rtree.rs
@@ -410,4 +410,8 @@ impl<
     fn assign_to_grouping<W>(&mut self, primitive: GenericIndex<W>, grouping: GenericIndex<GW>) {
         self.geometry.assign_to_grouping(primitive, grouping);
     }
+
+    fn groupings<W>(&self, node: GenericIndex<W>) -> impl Iterator<Item = GenericIndex<GW>> {
+        self.geometry.groupings(node)
+    }
 }
diff --git a/src/layout/layout.rs b/src/layout/layout.rs
index 1c9b185..834fd86 100644
--- a/src/layout/layout.rs
+++ b/src/layout/layout.rs
@@ -233,4 +233,11 @@ impl<R: RulesTrait> GroupingManagerTrait<ZoneWeight, GenericIndex<ZoneWeight>> f
     ) {
         self.drawing.assign_to_grouping(primitive, grouping);
     }
+
+    fn groupings<W>(
+        &self,
+        node: GenericIndex<W>,
+    ) -> impl Iterator<Item = GenericIndex<ZoneWeight>> {
+        self.drawing.groupings(node)
+    }
 }