lattice_graph/lattice_abstract/
edges.rs

1use std::iter::FusedIterator;
2
3use petgraph::visit::{EdgeRef, IntoEdgeReferences, IntoEdges, IntoEdgesDirected};
4
5use super::*;
6
7/// Edge reference data. See [`IntoEdgeReferences`] or [`IntoEdges`].
8#[derive(Debug, PartialEq, Eq)]
9pub struct EdgeReference<'a, C, E, D, A> {
10    pub(crate) source_id: C,
11    pub(crate) target_id: C,
12    pub(crate) edge_weight: &'a E,
13    pub(crate) direction: D,
14    pub(crate) axis: PhantomData<fn() -> A>,
15}
16
17impl<'a, C, E, D, A> EdgeReference<'a, C, E, D, A> {
18    /// Get a reference to the edge reference's direction.
19    pub fn direction(&self) -> &D {
20        &self.direction
21    }
22}
23
24impl<'a, C: Clone, E, D: Clone, A> Clone for EdgeReference<'a, C, E, D, A> {
25    fn clone(&self) -> Self {
26        Self {
27            source_id: self.source_id.clone(),
28            target_id: self.target_id.clone(),
29            edge_weight: self.edge_weight,
30            direction: self.direction.clone(),
31            axis: PhantomData,
32        }
33    }
34}
35
36impl<'a, C: Copy, E, D: Copy, A> Copy for EdgeReference<'a, C, E, D, A> {}
37
38impl<'a, C, E, D, A> EdgeRef for EdgeReference<'a, C, E, D, A>
39where
40    C: Copy,
41    D: AxisDirection + Copy,
42    A: Axis<Direction = D>,
43{
44    type NodeId = C;
45    type EdgeId = (C, A);
46    type Weight = E;
47
48    fn source(&self) -> Self::NodeId {
49        self.source_id
50    }
51
52    fn target(&self) -> Self::NodeId {
53        self.target_id
54    }
55
56    fn weight(&self) -> &Self::Weight {
57        self.edge_weight
58    }
59
60    fn id(&self) -> Self::EdgeId {
61        (self.source_id, A::from_direction(self.direction))
62    }
63}
64
65/// Edges connected to a node. See [`IntoEdges`].
66// Type parameter `C` is to derive `Debug`. (I don't want to impl manually).
67#[derive(Debug)]
68pub struct Edges<'a, N, E, S: Shape, C = <S as Shape>::Coordinate, Dt = AxisDirMarker> {
69    graph: &'a LatticeGraph<N, E, S>,
70    node: C,
71    offset: Offset,
72    state: usize,
73    directed: Dt,
74}
75
76/// Marker Used for [`Edges`] inside the [`IntoEdgeReferences`]
77#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
78pub struct AxisMarker;
79/// Marker Used for [`Edges`] as [`IntoEdges`]
80#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
81pub struct AxisDirMarker;
82/// Marker for [`Edges`].
83pub trait DtMarker {
84    const DIRECTED: bool;
85    // trick to be used in [`IntoEdgesDirected`]
86    #[inline]
87    unsafe fn get_raw_id<S: Shape>(
88        &self,
89        s: &S,
90        d: &<<S as Shape>::Axis as Axis>::Direction,
91        source: Offset,
92        target: <S as Shape>::Coordinate,
93        st: usize,
94    ) -> (Offset, usize) {
95        if <S as Shape>::Axis::is_forward_direction(d) {
96            (source, st)
97        } else {
98            (
99                s.to_offset_unchecked(target),
100                st - <S as Shape>::Axis::COUNT,
101            )
102        }
103    }
104    #[inline(always)]
105    fn need_reverse(&self) -> bool {
106        false
107    }
108}
109
110impl DtMarker for AxisMarker {
111    const DIRECTED: bool = true;
112}
113impl DtMarker for AxisDirMarker {
114    const DIRECTED: bool = false;
115}
116
117/// [`petgraph::Direction`] as marker for [`Edges`] used in [`IntoEdgesDirected`].
118impl DtMarker for petgraph::Direction {
119    const DIRECTED: bool = false;
120    // const MAYREVERSE: bool = true;
121    unsafe fn get_raw_id<S: Shape>(
122        &self,
123        s: &S,
124        d: &<<S as Shape>::Axis as Axis>::Direction,
125        source: Offset,
126        target: <S as Shape>::Coordinate,
127        st: usize,
128    ) -> (Offset, usize) {
129        match (self, <S as Shape>::Axis::DIRECTED) {
130            (petgraph::EdgeDirection::Incoming, true) => {
131                let da = <S as Shape>::Axis::from_direction(d.clone())
132                    .backward()
133                    .dir_to_index();
134                let o = s.to_offset_unchecked(target);
135                (o, da)
136            }
137            _ => AxisDirMarker.get_raw_id(s, d, source, target, st),
138        }
139    }
140    fn need_reverse(&self) -> bool {
141        self == &petgraph::Direction::Incoming
142    }
143}
144
145impl<'a, N, E, S, C, D, A, Dt> Edges<'a, N, E, S, C, Dt>
146where
147    C: Copy,
148    S: Shape<Coordinate = C, Axis = A>,
149    A: Axis<Direction = D>,
150    D: AxisDirection,
151{
152    fn new(g: &'a LatticeGraph<N, E, S>, a: C) -> Edges<'a, N, E, S, C, Dt>
153    where
154        Dt: Default,
155    {
156        Self::new_d(g, a, Dt::default())
157    }
158
159    fn new_d(g: &'a LatticeGraph<N, E, S>, a: C, d: Dt) -> Edges<'a, N, E, S, C, Dt> {
160        let offset = g.s.to_offset(a);
161        Edges {
162            graph: g,
163            node: a,
164            state: if offset.is_ok() {
165                0
166            } else {
167                S::Axis::UNDIRECTED_COUNT
168            },
169            offset: offset.unwrap_or_else(|_| unsafe { unreachable_debug_checked() }),
170            directed: d,
171        }
172    }
173
174    unsafe fn new_unchecked(g: &'a LatticeGraph<N, E, S>, a: C) -> Edges<'a, N, E, S, C, Dt>
175    where
176        Dt: Default,
177    {
178        let offset = g.s.to_offset(a);
179        Edges {
180            graph: g,
181            node: a,
182            state: 0,
183            offset: offset.unwrap_or_else(|_| unreachable_debug_checked()),
184            directed: Dt::default(),
185        }
186    }
187}
188
189impl<'a, N, E, S, C, D, A, Dt> Iterator for Edges<'a, N, E, S, C, Dt>
190where
191    C: Copy,
192    S: Shape<Coordinate = C, Axis = A>,
193    A: Axis<Direction = D>,
194    D: AxisDirection,
195    Dt: DtMarker,
196{
197    type Item = EdgeReference<'a, C, E, D, A>;
198
199    fn next(&mut self) -> Option<Self::Item> {
200        while self.state
201            < if Dt::DIRECTED {
202                A::COUNT
203            } else {
204                A::UNDIRECTED_COUNT
205            }
206        {
207            unsafe {
208                let d = D::dir_from_index_unchecked(self.state);
209                let n = self.graph.s.move_coord(self.node, d.clone());
210                let st = self.state;
211                self.state += 1;
212                if let Ok(target) = n {
213                    let (nx, ne) =
214                        self.directed
215                            .get_raw_id(&self.graph.s, &d, self.offset, target, st);
216                    debug_assert_eq!(A::from_direction(d.clone()).to_index(), ne);
217                    //let ne = S::Axis::from_direction(d.clone()).to_index();
218                    let e = self.graph.edge_weight_unchecked_raw((nx, ne));
219                    let (source_id, target_id) = if self.directed.need_reverse() {
220                        (target, self.node)
221                    } else {
222                        (self.node, target)
223                    };
224                    return Some(EdgeReference {
225                        source_id,
226                        target_id,
227                        edge_weight: e,
228                        direction: d,
229                        axis: PhantomData,
230                    });
231                }
232            }
233        }
234        None
235    }
236
237    fn size_hint(&self) -> (usize, Option<usize>) {
238        let x = if Dt::DIRECTED {
239            A::COUNT
240        } else {
241            A::UNDIRECTED_COUNT
242        } - self.state;
243        (0, Some(x))
244    }
245}
246
247impl<'a, N, E, S, C, Dt> FusedIterator for Edges<'a, N, E, S, C, Dt>
248where
249    Self: Iterator,
250    S: Shape,
251{
252}
253
254impl<'a, N, E, S, C, D, A> IntoEdges for &'a LatticeGraph<N, E, S>
255where
256    C: Copy,
257    S: Shape<Coordinate = C, Axis = A>,
258    A: Axis<Direction = D>,
259    D: AxisDirection + Copy,
260{
261    type Edges = Edges<'a, N, E, S>;
262
263    fn edges(self, a: Self::NodeId) -> Self::Edges {
264        Edges::new(self, a)
265    }
266}
267
268/// Edges connected to a node with [`Direction`](`petgraph::Direction`). See [`IntoEdgesDirected`].
269pub type EdgesDirected<'a, N, E, S> =
270    Edges<'a, N, E, S, <S as Shape>::Coordinate, petgraph::Direction>;
271impl<'a, N, E, S, C, D, A> IntoEdgesDirected for &'a LatticeGraph<N, E, S>
272where
273    C: Copy,
274    S: Shape<Coordinate = C, Axis = A>,
275    A: Axis<Direction = D>,
276    D: AxisDirection + Copy,
277{
278    type EdgesDirected = EdgesDirected<'a, N, E, S>;
279
280    fn edges_directed(self, a: Self::NodeId, dir: petgraph::Direction) -> Self::EdgesDirected {
281        Edges::new_d(self, a, dir)
282    }
283}
284
285/// Iterator for all edges of [`LatticeGraph`]. See [`IntoEdgeReferences`](`IntoEdgeReferences::edge_references`).
286// Type parameter `C` is to derive `Debug`. (I don't want to impl manually).
287#[derive(Debug)]
288pub struct EdgeReferences<'a, N, E, S: Shape, C = <S as Shape>::Coordinate> {
289    g: &'a LatticeGraph<N, E, S>,
290    e: Option<Edges<'a, N, E, S, C, AxisMarker>>,
291    index: usize,
292}
293
294impl<'a, N, E, S, C, D, A> Iterator for EdgeReferences<'a, N, E, S, C>
295where
296    C: Copy,
297    S: Shape<Coordinate = C, Axis = A>,
298    D: AxisDirection + Copy,
299    A: Axis<Direction = D>,
300{
301    type Item = EdgeReference<'a, C, E, D, A>;
302
303    fn next(&mut self) -> Option<Self::Item> {
304        loop {
305            if let Some(ref mut e) = self.e {
306                let next = e.next();
307                if next.is_some() {
308                    return next;
309                }
310            }
311            if self.index < self.g.s.node_count() {
312                let x = self.g.s.index_to_coordinate(self.index);
313                self.index += 1;
314                //self.e = Some(self.g.edges(x));
315                self.e = Some(unsafe { Edges::new_unchecked(self.g, x) });
316            } else {
317                return None;
318            }
319        }
320    }
321
322    fn size_hint(&self) -> (usize, Option<usize>) {
323        let node_len = self.g.node_count() - self.index;
324        let maxlen = node_len * S::Axis::UNDIRECTED_COUNT
325            + self
326                .e
327                .as_ref()
328                .map(|x| x.size_hint().1.unwrap_or(0))
329                .unwrap_or(0);
330        (0, Some(maxlen))
331    }
332}
333
334impl<'a, N, E, S, C, D, A> FusedIterator for EdgeReferences<'a, N, E, S, C>
335where
336    C: Copy,
337    S: Shape<Coordinate = C, Axis = A>,
338    D: AxisDirection + Copy,
339    A: Axis<Direction = D>,
340{
341}
342
343impl<'a, N, E, S, C, D, A> IntoEdgeReferences for &'a LatticeGraph<N, E, S>
344where
345    C: Copy,
346    S: Shape<Coordinate = C, Axis = A>,
347    A: Axis<Direction = D>,
348    D: AxisDirection + Copy,
349{
350    type EdgeRef = EdgeReference<'a, C, E, D, A>;
351    type EdgeReferences = EdgeReferences<'a, N, E, S>;
352
353    fn edge_references(self) -> Self::EdgeReferences {
354        EdgeReferences {
355            g: self,
356            e: None,
357            index: 0,
358        }
359    }
360}
lattice_graph/lattice_abstract/edges.rs

lattice_graph/lattice_abstract/
edges.rs
